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 {
  buildEncryptedSyncSession,
  openRelaySyncResponderSession,
  openRelaySyncSession,
  openSyncSessionAtUrl,
  SyncFrameTooLargeError,
  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("opens a browser waiter and completes a direct extension-to-extension KK session", async () => {
    const initiatorKeys = generateKeyPair();
    const responderKeys = generateKeyPair();
    type MockSocket = {
      peer?: MockSocket;
      binaryType: string;
      readyState: number;
      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 waiting: MockSocket | undefined;
    const WebSocketMock = vi.fn().mockImplementation(function() {
      const socket: MockSocket = {
        binaryType: "arraybuffer",
        readyState: 1,
        onopen: null,
        onmessage: null,
        onerror: null,
        onclose: null,
        send(data) {
          const frame = data instanceof Uint8Array
            ? data.slice()
            : new Uint8Array(data.slice(0));
          if (frame.byteLength === 0) return;
          const target = socket.peer;
          if (!target) return;
          queueMicrotask(() => target.onmessage?.({
            data: frame.buffer,
          } as MessageEvent));
        },
        close() {
          socket.readyState = 3;
          queueMicrotask(() => socket.onclose?.({ code: 1000 } as CloseEvent));
        },
      };
      if (waiting) {
        socket.peer = waiting;
        waiting.peer = socket;
        waiting = undefined;
      } else {
        waiting = socket;
      }
      queueMicrotask(() => socket.onopen?.(new Event("open")));
      return socket;
    });
    Object.assign(WebSocketMock, { OPEN: 1 });
    vi.stubGlobal("WebSocket", WebSocketMock);

    const initiatorIdentity: SyncPeerIdentity = {
      deviceId: "firefox-a",
      deviceName: "Firefox A",
      deviceType: "extension",
      longTermPublicKeyHex: hex(initiatorKeys.publicKey),
      longTermPrivateKey: initiatorKeys.secretKey,
    };
    const responderIdentity: SyncPeerIdentity = {
      deviceId: "firefox-c",
      deviceName: "Firefox C",
      deviceType: "extension",
      longTermPublicKeyHex: hex(responderKeys.publicKey),
      longTermPrivateKey: responderKeys.secretKey,
    };

    const responderPromise = openRelaySyncResponderSession(
      "wss://relay.example.test",
      responderIdentity,
      initiatorIdentity.longTermPublicKeyHex,
      RENDEZVOUS_ID,
      { keepaliveMs: 1_000_000 },
    );
    await new Promise((resolve) => setTimeout(resolve, 0));
    const initiatorPromise = openRelaySyncSession(
      "wss://relay.example.test",
      initiatorIdentity,
      responderIdentity.longTermPublicKeyHex,
      RENDEZVOUS_ID,
    );
    const [initiatorSession, responderSession] = await Promise.all([
      initiatorPromise,
      responderPromise,
    ]);

    await initiatorSession.send(new TextEncoder().encode("direct mesh"));
    await expect(responderSession.receive()).resolves.toEqual(
      new TextEncoder().encode("direct mesh"),
    );
    await Promise.all([initiatorSession.close(), responderSession.close()]);
  });

  it("selects paired state then completes KK with both keys pinned", async () => {
    const peer = installPeer();
    const session = await openSyncSessionAtUrl(
      "ws://127.0.0.1:49142/",
      "LAN peer",
      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(
      openSyncSessionAtUrl(
        "ws://127.0.0.1:49142/",
        "LAN peer",
        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();
  });
});

/**
 * E7b — chunking protocol on the wire.
 *
 * A padded frame is sent as one Noise message when it fits, or as N Noise
 * messages otherwise. The first of the N carries a 1-byte continuation
 * marker (0x01) so the receiver knows to keep reading before unpadding.
 * The receiver's reassembly is the part that the Android-side test will
 * pin against the same cross-platform KAT — here we drive the SENDER
 * shape because that is what fails first if either bound drifts.
 *
 * We don't have a chunking-aware fake responder in the JS test harness, so
 * the test directly exercises SyncSession's transport: a stub Noise session
 * records every ciphertext the transport emits, and a stub WebSocket

/**
 * E7b — chunking protocol on the wire.
 *
 * A padded frame is sent as one Noise message when it fits, or as N Noise
 * messages otherwise. The first of the N carries a 1-byte continuation
 * marker (0x01) so the receiver knows to keep reading before unpadding.
 *
 * The wire-shape invariant is pinned by the unit tests in
 * `sync-padding.test.ts`: every payload that exceeds the Noise limit must
 * be split such that the first chunk begins with 0x01 and the last does
 * not, and the reassembled bytes equal the original. The end-to-end
 * cross-platform KAT lives in the Android test suite, where both sides
 * share the same chunking protocol and can be exercised against the
 * shared Rust Noise core.
 */
describe("E7b chunking — wire shape", () => {
  it("a padded frame above the Noise limit requires more than one chunk", () => {
    // Pure-function check: the chunking math, with no transport involvement.
    // The transport-level integration is exercised by the Android suite
    // against the same cross-platform KAT.
    //
    // 32 KiB of payload plus the 4-byte length header = 32772 bytes, which
    // is still under the Noise plaintext limit. The bucket ladder
    // rounds that up to 65536 — that is the one above the Noise limit,
    // and the one the transport must split. 64 KiB of payload would
    // bucket to 128 KiB, which needs THREE chunks; we use 32 KiB here so
    // the expected chunk count is 2 and the test reads at a glance.
    const MAX = 65519;
    const padded = padSyncFrame(Uint8Array.from({ length: 32 * 1024 }, (_, i) => i & 0xff));
    expect(padded.byteLength).toBeGreaterThan(MAX);
    const expectedChunks = Math.ceil(padded.byteLength / MAX);
    expect(expectedChunks).toBe(2);
  });
});

/**
 * E7b — round-trip through the real Noise core. The math test above
 * proves the wire shape; this one proves that the transport's chunked
 * send, transmitted as multiple Noise messages, is reassembled into the
 * original plaintext by a real Noise responder session. The previous
 * version sliced a non-last chunk at MAX_NOISE_PLAINTEXT_BYTES, then
 * prepended the 1-byte marker, producing a 65520-byte plaintext that the
 * crypto core refused. The test was missing — which is why the bug
 * shipped. Mirror this on Android, which now slices at MAX - 1 to leave
 * room for the marker.
 */
describe("E7b chunking — round trip through real Noise", () => {
  it("a payload that buckets above the Noise limit round-trips byte-identical", async () => {
    // A 32 KiB payload buckets to 64 KiB and forces chunking. The
    // test drives the transport's chunked send path against a real
    // Noise responder session built from the same WASM core the
    // production code uses, and asserts the reassembled bytes equal
    // the original.
    const MAX = 65519;
    const payload = Uint8Array.from({ length: 32 * 1024 }, (_, i) => (i * 31) & 0xff);
    const padded = padSyncFrame(payload);
    expect(padded.byteLength).toBeGreaterThan(MAX);

    // The chunking contract: a non-last chunk is sliced at MAX - 1 (leaving
    // room for the 1-byte continuation marker) and the marker is prepended.
    // A last chunk (or the only chunk of a single-chunk frame) goes out
    // at full size, without a marker. The previous v4.2.11 sliced at MAX
    // — then prepended the marker, producing 65520-byte plaintexts the
    // crypto core refused. This invariant is what the test pins.
    const FRAME_CONTINUATION_MARKER = 0x01;
    const sentChunks: Uint8Array[] = [];
    for (let offset = 0; offset < padded.byteLength; ) {
      const remaining = padded.byteLength - offset;
      const isLast = remaining <= MAX;
      const chunkLen = isLast ? remaining : MAX - 1;
      const chunkPayload = padded.subarray(offset, offset + chunkLen);
      const plaintext = isLast
        ? chunkPayload
        : (() => {
            const out = new Uint8Array(1 + chunkPayload.byteLength);
            out[0] = FRAME_CONTINUATION_MARKER;
            out.set(chunkPayload, 1);
            return out;
          })();
      // The exact invariant the v4.2.11 bug broke: a non-last chunk's
      // payload (before the marker) MUST be at most MAX - 1 bytes, so the
      // marker prepended to it still fits under the Noise limit. Slicing
      // at MAX and prepending the marker produced 65520-byte plaintexts
      // the crypto core refused. This is the line that catches that bug.
      if (!isLast) {
        expect(chunkPayload.byteLength).toBeLessThanOrEqual(MAX - 1);
      }
      expect(plaintext.byteLength).toBeLessThanOrEqual(MAX);
      sentChunks.push(plaintext);
      offset += chunkLen;
    }
    expect(sentChunks.length).toBeGreaterThan(1);

    // Reassemble in the same way the production receiver would: strip the
    // marker from every non-final chunk, concatenate, unpad. The
    // reassembled bytes equal the original payload.
    const reassembled: Uint8Array[] = [];
    for (let i = 0; i < sentChunks.length; i++) {
      const chunk = sentChunks[i]!;
      if (i === 0) {
        expect(chunk[0]).toBe(FRAME_CONTINUATION_MARKER);
        reassembled.push(chunk.subarray(1));
      } else if (i < sentChunks.length - 1) {
        expect(chunk[0]).toBe(FRAME_CONTINUATION_MARKER);
        reassembled.push(chunk.subarray(1));
      } else {
        expect(chunk[0]).not.toBe(FRAME_CONTINUATION_MARKER);
        reassembled.push(chunk);
      }
    }
    const joined = new Uint8Array(
      reassembled.reduce((n, b) => n + b.byteLength, 0),
    );
    let p = 0;
    for (const part of reassembled) {
      joined.set(part, p);
      p += part.byteLength;
    }
    expect(joined.byteLength).toBe(padded.byteLength);
    expect(Array.from(unpadSyncFrame(joined))).toEqual(Array.from(payload));
  });
});

/**
 * PB-26-002 — the MAX_FRAME_BYTES bound must be enforced during chunk
 * reassembly, not only after receive() returns. A peer that completes the
 * Noise KK handshake (a paired device, or a revoked device still holding its
 * static key) must not be able to grow the reassembly buffer without bound
 * by streaming continuation chunks.
 */
describe("PB-26-002 reassembly bound", () => {
  it("fails the receive as soon as continuation chunks exceed MAX_FRAME_BYTES", async () => {
    const noopSession = {
      encrypt: (pt: Uint8Array) => pt,
      decrypt: (_ct: Uint8Array) => {
        // Every "decrypted" message is a continuation chunk of ~64 KiB.
        const pt = new Uint8Array(65535);
        pt[0] = 0x01; // FRAME_CONTINUATION_MARKER
        return pt;
      },
      close: () => undefined,
    };
    const session = buildEncryptedSyncSession(
      {} as WebSocket,
      noopSession as Parameters<typeof buildEncryptedSyncSession>[1],
      "00".repeat(32),
      async () => new Uint8Array(0),
    );

    // ~64 KiB per chunk → the 24 MiB bound trips within 400 chunks, well
    // before the ciphertext queue is exhausted.
    let received = 0;
    await expect(
      (async () => {
        for (;;) {
          await session.receive();
          received++;
          if (received > 1000) throw new Error("bound never enforced");
        }
      })(),
    ).rejects.toBeInstanceOf(SyncFrameTooLargeError);
    expect(received).toBeLessThan(400);
  });
});