import { afterEach, describe, expect, it, vi } from "vitest";
import { generateKeyPair } from "@stablelib/x25519";
import { noiseKKResponder } from "./noise-js";
import { padSyncFrame, unpadSyncFrame } from "./sync-padding";
import {
  openRelaySyncSession,
  openSyncSession,
  type SyncPeerIdentity,
} from "./sync-transport";

const encode = (value: unknown) => new TextEncoder().encode(JSON.stringify(value));
const decode = (bytes: Uint8Array) => JSON.parse(new TextDecoder().decode(bytes));
const hex = (bytes: Uint8Array) =>
  Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join("");
const RENDEZVOUS_ID = "Q".repeat(43);

function installPeer(peerIdentityKeyHex?: string) {
  const initiator = generateKeyPair();
  const responder = generateKeyPair();
  const clientFrames: Uint8Array[] = [];
  const clientReaders: Array<(frame: Uint8Array) => void> = [];
  let socket: {
    binaryType: string;
    onopen: ((event: Event) => void) | null;
    onmessage: ((event: MessageEvent) => void) | null;
    onerror: (() => void) | null;
    onclose: ((event: CloseEvent) => void) | null;
    send(data: Uint8Array | ArrayBuffer): void;
    close(): void;
  };
  let resolveUrl!: (url: string) => void;
  const urlSeen = new Promise<string>((resolve) => { resolveUrl = resolve; });
  const receiveClient = async () => {
    const queued = clientFrames.shift();
    if (queued) return queued;
    return new Promise<Uint8Array>((resolve) => clientReaders.push(resolve));
  };
  const sendServer = async (frame: Uint8Array) => {
    queueMicrotask(() => socket.onmessage?.({ data: frame.buffer } as MessageEvent));
  };

  const selectorSeen = new Promise<Record<string, unknown>>((resolve, reject) => {
    const WebSocketMock = vi.fn().mockImplementation(function(url: string) {
      resolveUrl(url);
      socket = {
        binaryType: "arraybuffer",
        onopen: null,
        onmessage: null,
        onerror: null,
        onclose: null,
        send(data) {
          const frame = data instanceof Uint8Array ? data.slice() : new Uint8Array(data.slice(0));
          const reader = clientReaders.shift();
          if (reader) reader(frame);
          else clientFrames.push(frame);
        },
        close() {
          queueMicrotask(() => socket.onclose?.({ code: 1000 } as CloseEvent));
        },
      };
      queueMicrotask(() => socket.onopen?.(new Event("open")));
      return socket;
    });
    vi.stubGlobal("WebSocket", WebSocketMock);

    void (async () => {
      try {
        const selector = decode(await receiveClient()) as Record<string, unknown>;
        resolve(selector);
        const session = await noiseKKResponder(
          responder.secretKey,
          responder.publicKey,
          initiator.publicKey,
          sendServer,
          receiveClient,
        );
        // E7: the browser pads every session frame, so this stand-in for the
        // Android listener must unpad exactly as SyncPadding does.
        const theirIdentity = decode(unpadSyncFrame(session.decrypt(await receiveClient())));
        if (theirIdentity.long_term_public_key_hex !== hex(initiator.publicKey)) {
          throw new Error("initiator identity was not bound to KK key");
        }
        await sendServer(session.encrypt(padSyncFrame(encode({
          type: "PAIR_IDENTITY",
          device_id: "phone",
          device_name: "Pixel",
          device_type: "android",
          long_term_public_key_hex: peerIdentityKeyHex ?? hex(responder.publicKey),
        }))));
      } catch (error) {
        reject(error);
      }
    })();
  });

  const identity: SyncPeerIdentity = {
    deviceId: "browser",
    deviceName: "Firefox",
    deviceType: "extension",
    longTermPublicKeyHex: hex(initiator.publicKey),
    longTermPrivateKey: initiator.secretKey,
  };
  return { identity, responder, selectorSeen, urlSeen };
}

afterEach(() => vi.unstubAllGlobals());

describe("KK WebSocket reconnect transport", () => {
  it("selects paired state then completes KK with both keys pinned", async () => {
    const peer = installPeer();
    const session = await openSyncSession(
      { ip: "127.0.0.1", port: 49142 },
      peer.identity,
      hex(peer.responder.publicKey),
      RENDEZVOUS_ID,
    );

    await expect(peer.selectorSeen).resolves.toEqual({
      type: "PEACH_RECONNECT",
      rendezvous_id: RENDEZVOUS_ID,
    });
    expect(hex(session.remoteStaticPublicKey)).toBe(hex(peer.responder.publicKey));
    await session.close();
  });

  it("rejects a post-KK identity that claims a different responder key", async () => {
    const peer = installPeer("ff".repeat(32));
    await expect(
      openSyncSession(
        { ip: "127.0.0.1", port: 49142 },
        peer.identity,
        hex(peer.responder.publicKey),
        RENDEZVOUS_ID,
      ),
    ).rejects.toThrow(/identity does not match authenticated KK key/);
  });

  it("uses the same KK wire over a credential-free wss relay target", async () => {
    const peer = installPeer();
    const session = await openRelaySyncSession(
      "wss://relay.example.test/base/",
      peer.identity,
      hex(peer.responder.publicKey),
      RENDEZVOUS_ID,
    );

    await expect(peer.urlSeen).resolves.toBe(
      `wss://relay.example.test/base/r/${RENDEZVOUS_ID}`,
    );
    await expect(peer.selectorSeen).resolves.toEqual({
      type: "PEACH_RECONNECT",
      rendezvous_id: RENDEZVOUS_ID,
    });
    await session.close();
  });
});