import { describe, it, expect, beforeAll, afterAll } from "vitest";
import WebSocket from "ws";
import {
  RelayServer,
  generateKeyPair,
  deviceIdFromKeys,
  deriveSharedKey,
  sign,
  MESSAGE_TYPES,
} from "../src/index.js";
import { buildEnvelope, openEnvelope } from "../src/sync/envelope.js";
import type { CookieSyncPayload, Envelope } from "../src/protocol/spec.js";

// Helper: connect and authenticate a device
function connectDevice(
  port: number,
  keys: ReturnType<typeof generateKeyPair>,
): Promise<{ ws: WebSocket; deviceId: string }> {
  return new Promise((resolve, reject) => {
    const deviceId = deviceIdFromKeys(keys);
    const ws = new WebSocket(`ws://127.0.0.1:${port}/ws`);

    ws.on("error", reject);

    ws.on("message", (data: Buffer) => {
      const msg = JSON.parse(data.toString());
      if (msg.type === "auth_challenge") {
        const challenge = Buffer.from(msg.challenge, "hex");
        const sig = sign(challenge, keys.signSec);
        ws.send(
          JSON.stringify({
            type: "auth_response",
            deviceId,
            sig: sig.toString("hex"),
          }),
        );
      } else if (msg.type === "auth_ok") {
        resolve({ ws, deviceId });
      } else if (msg.type === "error") {
        reject(new Error(msg.error));
      }
    });
  });
}

// Helper: wait for next message of a given type
function waitForMessage(
  ws: WebSocket,
  type: string,
  timeoutMs = 5000,
): Promise<Record<string, unknown>> {
  return new Promise((resolve, reject) => {
    const timer = setTimeout(() => reject(new Error(`Timeout waiting for ${type}`)), timeoutMs);
    const handler = (data: Buffer) => {
      const msg = JSON.parse(data.toString());
      if (msg.type === type) {
        clearTimeout(timer);
        ws.off("message", handler);
        resolve(msg);
      }
    };
    ws.on("message", handler);
  });
}

describe("Integration: relay server end-to-end", () => {
  let server: RelayServer;
  let port: number;

  beforeAll(async () => {
    server = new RelayServer({ port: 0 }); // random port
    await server.start();
    port = server.port;
  });

  afterAll(async () => {
    await server.stop();
  });

  it("health check works", async () => {
    const res = await fetch(`http://127.0.0.1:${port}/health`);
    const body = await res.json();
    expect(body.status).toBe("ok");
  });

  it("pairing flow: create and accept", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();

    // Alice initiates pairing
    const createRes = await fetch(`http://127.0.0.1:${port}/pair`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        deviceId: deviceIdFromKeys(alice),
        x25519PubKey: alice.encPub.toString("hex"),
      }),
    });
    expect(createRes.status).toBe(201);
    const { pairingCode } = (await createRes.json()) as { pairingCode: string };
    expect(pairingCode).toHaveLength(6);

    // Bob accepts with the code
    const acceptRes = await fetch(`http://127.0.0.1:${port}/pair/accept`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        deviceId: deviceIdFromKeys(bob),
        x25519PubKey: bob.encPub.toString("hex"),
        pairingCode,
      }),
    });
    expect(acceptRes.status).toBe(200);
    const result = (await acceptRes.json()) as {
      initiator: { deviceId: string; x25519PubKey: string };
      acceptor: { deviceId: string; x25519PubKey: string };
    };
    expect(result.initiator.deviceId).toBe(deviceIdFromKeys(alice));
    expect(result.acceptor.deviceId).toBe(deviceIdFromKeys(bob));
  });

  it("rejects invalid pairing code", async () => {
    const res = await fetch(`http://127.0.0.1:${port}/pair/accept`, {
      method: "POST",
      headers: { "Content-Type": "application/json" },
      body: JSON.stringify({
        deviceId: "fake",
        x25519PubKey: "fake",
        pairingCode: "000000",
      }),
    });
    expect(res.status).toBe(404);
  });

  it("authenticates devices via WebSocket", async () => {
    const alice = generateKeyPair();
    const { ws, deviceId } = await connectDevice(port, alice);
    expect(deviceId).toBe(deviceIdFromKeys(alice));
    ws.close();
  });

  it("rejects bad auth signatures", async () => {
    const alice = generateKeyPair();
    const eve = generateKeyPair(); // wrong keys

    await expect(
      new Promise<void>((resolve, reject) => {
        const ws = new WebSocket(`ws://127.0.0.1:${port}/ws`);
        ws.on("message", (data: Buffer) => {
          const msg = JSON.parse(data.toString());
          if (msg.type === "auth_challenge") {
            // Sign with wrong key
            const challenge = Buffer.from(msg.challenge, "hex");
            const sig = sign(challenge, eve.signSec);
            ws.send(
              JSON.stringify({
                type: "auth_response",
                deviceId: deviceIdFromKeys(alice), // claim to be alice
                sig: sig.toString("hex"),
              }),
            );
          }
        });
        ws.on("close", (code: number) => {
          if (code === 4003) reject(new Error("Auth failed as expected"));
          else resolve();
        });
      }),
    ).rejects.toThrow("Auth failed");
  });

  it("relays encrypted cookie sync between two devices", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();

    const aliceConn = await connectDevice(port, alice);
    const bobConn = await connectDevice(port, bob);

    const aliceDeviceId = deviceIdFromKeys(alice);
    const bobDeviceId = deviceIdFromKeys(bob);

    // Alice sends a cookie sync to Bob
    const payload: CookieSyncPayload = {
      action: "set",
      cookies: [
        {
          domain: "example.com",
          name: "session",
          value: "tok_abc123",
          path: "/",
          secure: true,
          httpOnly: true,
          sameSite: "lax",
          expiresAt: null,
        },
      ],
      lamportTs: 1,
    };

    const envelope = buildEnvelope(
      MESSAGE_TYPES.COOKIE_SYNC,
      payload,
      alice,
      bob.encPub,
      bobDeviceId,
    );

    // Bob listens for the message
    const bobMessagePromise = waitForMessage(bobConn.ws, MESSAGE_TYPES.COOKIE_SYNC);

    // Alice sends
    aliceConn.ws.send(JSON.stringify(envelope));

    // Alice gets ACK
    const ack = await waitForMessage(aliceConn.ws, MESSAGE_TYPES.ACK);
    expect(ack.delivered).toBe(true);

    // Bob receives the encrypted envelope
    const received = (await bobMessagePromise) as unknown as Envelope;
    expect(received.from).toBe(aliceDeviceId);
    expect(received.to).toBe(bobDeviceId);

    // Bob decrypts
    const decrypted = openEnvelope(received, bob, alice.encPub);
    expect(decrypted.action).toBe("set");
    expect(decrypted.cookies).toHaveLength(1);
    expect(decrypted.cookies[0].domain).toBe("example.com");
    expect(decrypted.cookies[0].value).toBe("tok_abc123");

    aliceConn.ws.close();
    bobConn.ws.close();
  });

  it("queues messages for offline devices and delivers on reconnect", async () => {
    const alice = generateKeyPair();
    const bob = generateKeyPair();
    const bobDeviceId = deviceIdFromKeys(bob);

    // Alice connects, Bob is offline
    const aliceConn = await connectDevice(port, alice);

    const payload: CookieSyncPayload = {
      action: "set",
      cookies: [
        {
          domain: "queued.com",
          name: "token",
          value: "queued_val",
          path: "/",
          secure: true,
          httpOnly: false,
          sameSite: "none",
          expiresAt: null,
        },
      ],
      lamportTs: 1,
    };

    const envelope = buildEnvelope(
      MESSAGE_TYPES.COOKIE_SYNC,
      payload,
      alice,
      bob.encPub,
      bobDeviceId,
    );

    // Send while Bob is offline
    aliceConn.ws.send(JSON.stringify(envelope));

    const ack = await waitForMessage(aliceConn.ws, MESSAGE_TYPES.ACK);
    expect(ack.delivered).toBe(false); // queued, not delivered

    // Bob comes online — should receive the queued message
    const bobMessagePromise = new Promise<Envelope>((resolve) => {
      const bobWs = new WebSocket(`ws://127.0.0.1:${port}/ws`);
      bobWs.on("message", (data: Buffer) => {
        const msg = JSON.parse(data.toString());
        if (msg.type === "auth_challenge") {
          const challenge = Buffer.from(msg.challenge, "hex");
          const sig = sign(challenge, bob.signSec);
          bobWs.send(
            JSON.stringify({
              type: "auth_response",
              deviceId: bobDeviceId,
              sig: sig.toString("hex"),
            }),
          );
        } else if (msg.type === MESSAGE_TYPES.COOKIE_SYNC) {
          resolve(msg as unknown as Envelope);
          bobWs.close();
        }
      });
    });

    const received = await bobMessagePromise;
    const decrypted = openEnvelope(received, bob, alice.encPub);
    expect(decrypted.cookies[0].value).toBe("queued_val");

    aliceConn.ws.close();
  });
});