import { describe, expect, it } from "vitest";
import {
  MAX_ADOPT_SNAPSHOT_BYTES,
  MAX_ADOPT_SNAPSHOT_ENTRIES,
  MAX_ADOPT_SNAPSHOT_ENVELOPE_BYTES,
  MAX_ADOPT_SNAPSHOT_ID_CHARS,
  MAX_ADOPT_SNAPSHOT_TOMBSTONES,
  MAX_SNAPSHOT_CHUNKS,
  SNAPSHOT_FRAME_WIRE_CEILING,
  applyAdoptSnapshotFrame,
  createAdoptSnapshotBudget,
  isAdoptSnapshot,
  receiveFrame,
  receiveVaultSnapshot,
  requestAdoption,
  sendVaultSnapshot,
  type AdoptionReceiverDeps,
} from "./vault-adopt";
import type { PairingTransport } from "./pairing-handshake";
import { padSyncFrame } from "./sync-padding";

/** Snapshot frames are bucket-padded on the wire (PB-26-604). */
function sendSnapshotFrame(channel: PairingTransport, frame: unknown): Promise<void> {
  return channel.send(padSyncFrame(new TextEncoder().encode(JSON.stringify(frame))));
}

function memoryTransportPair(): [PairingTransport, PairingTransport] {
  const a: Uint8Array[] = [];
  const b: Uint8Array[] = [];
  const aReaders: Array<(message: Uint8Array) => void> = [];
  const bReaders: Array<(message: Uint8Array) => void> = [];

  const transport = (
    inbox: Uint8Array[],
    outbox: Uint8Array[],
    inboxReaders: Array<(message: Uint8Array) => void>,
    outboxReaders: Array<(message: Uint8Array) => void>,
  ): PairingTransport => ({
    async send(message) {
      const reader = outboxReaders.shift();
      if (reader) reader(message);
      else outbox.push(message);
    },
    receive() {
      const message = inbox.shift();
      return message
        ? Promise.resolve(message)
        : new Promise<Uint8Array>((resolve) => inboxReaders.push(resolve));
    },
  });

  return [transport(a, b, aReaders, bReaders), transport(b, a, bReaders, aReaders)];
}

describe("inline vault snapshot", () => {
  it("round-trips an empty vault (0 entries)", async () => {
    const [source, receiver] = memoryTransportPair();
    const [result] = await Promise.all([
      receiveVaultSnapshot(receiver),
      sendVaultSnapshot(source, {}, {}),
    ]);
    expect(result.entries).toEqual({});
    expect(result.tombstones).toEqual({});
  });

  it("round-trips 50 entries + 5 tombstones in a single chunk", async () => {
    const [source, receiver] = memoryTransportPair();
    const entries: Record<string, string> = {};
    for (let i = 0; i < 50; i++) entries[`id-${i}`] = `cipher-${i}`;
    const tombstones: Record<string, number> = {};
    for (let i = 0; i < 5; i++) tombstones[`dead-${i}`] = 1000 + i;
    const [result] = await Promise.all([
      receiveVaultSnapshot(receiver),
      sendVaultSnapshot(source, entries, tombstones),
    ]);
    expect(result.entries).toEqual(entries);
    expect(result.tombstones).toEqual(tombstones);
  });

  it("splits a large vault into 3+ chunks and reassembles it", async () => {
    const [source, receiver] = memoryTransportPair();
    let frames = 0;
    const countingSource: PairingTransport = {
      send: (msg) => {
        frames++;
        return source.send(msg);
      },
      receive: () => source.receive(),
    };
    const big = "x".repeat(200 * 1024); // 200 KiB per entry forces chunking
    const entries: Record<string, string> = {};
    for (let i = 0; i < 5; i++) entries[`id-${i}`] = big;
    const tombstones: Record<string, number> = { "dead-0": 42 };
    const [result] = await Promise.all([
      receiveVaultSnapshot(receiver),
      sendVaultSnapshot(countingSource, entries, tombstones),
    ]);
    expect(frames).toBeGreaterThanOrEqual(3);
    expect(result.entries).toEqual(entries);
    expect(result.tombstones).toEqual(tombstones);
  });

  it("rejects oversized duress before sending any snapshot frame", async () => {
    const [source] = memoryTransportPair();
    let frames = 0;
    const countingSource: PairingTransport = {
      send: (msg) => {
        frames++;
        return source.send(msg);
      },
      receive: () => source.receive(),
    };

    await expect(sendVaultSnapshot(
      countingSource,
      { e1: "cipher1" },
      {},
      "x".repeat(768 * 1024),
    )).rejects.toThrow("ADOPT_SNAPSHOT_DURESS_FRAME_TOO_LARGE");
    expect(frames).toBe(0);
  });

  it("rejects out-of-order chunks", async () => {
    const [source, receiver] = memoryTransportPair();
    const send = sendSnapshotFrame(source, {
      type: "ADOPT_SNAPSHOT",
      chunk_index: 1,
      chunk_count: 2,
      entries: { a: "x" },
      tombstones: {},
    });
    await expect(receiveVaultSnapshot(receiver)).rejects.toThrow(/OUT_OF_ORDER/);
    await send;
  });

  it("rejects a chunk_count exceeding MAX_SNAPSHOT_CHUNKS", async () => {
    const [source, receiver] = memoryTransportPair();
    const send = sendSnapshotFrame(source, {
      type: "ADOPT_SNAPSHOT",
      chunk_index: 0,
      chunk_count: MAX_SNAPSHOT_CHUNKS + 1,
      entries: {},
      tombstones: {},
    });
    await expect(receiveVaultSnapshot(receiver)).rejects.toThrow(/TOO_MANY_CHUNKS/);
    await send;
  });

  it("validates ADOPT_SNAPSHOT frames", () => {
    expect(
      isAdoptSnapshot({
        type: "ADOPT_SNAPSHOT",
        chunk_index: 0,
        chunk_count: 1,
        entries: { a: "x" },
        tombstones: {},
      }),
    ).toBe(true);
    // Missing tombstones / chunk_index.
    expect(
      isAdoptSnapshot({ type: "ADOPT_SNAPSHOT", chunk_index: 0, chunk_count: 1, entries: {} }),
    ).toBe(false);
    expect(
      isAdoptSnapshot({ type: "ADOPT_SNAPSHOT", chunk_count: 1, entries: {}, tombstones: {} }),
    ).toBe(false);
    // Wrong type tag.
    expect(
      isAdoptSnapshot({ type: "ADOPT_GRANT", chunk_index: 0, chunk_count: 1, entries: {}, tombstones: {} }),
    ).toBe(false);
    // Invalid indices and value types.
    expect(
      isAdoptSnapshot({ type: "ADOPT_SNAPSHOT", chunk_index: -1, chunk_count: 1, entries: {}, tombstones: {} }),
    ).toBe(false);
    expect(
      isAdoptSnapshot({ type: "ADOPT_SNAPSHOT", chunk_index: 0, chunk_count: 0, entries: {}, tombstones: {} }),
    ).toBe(false);
    expect(
      isAdoptSnapshot({ type: "ADOPT_SNAPSHOT", chunk_index: 0, chunk_count: 1, entries: { a: 5 }, tombstones: {} }),
    ).toBe(false);
    expect(isAdoptSnapshot(null)).toBe(false);
  });

  it("advertises snapshot_capable in the ADOPT_REQUEST frame", async () => {
    const [receiverChannel, sourceChannel] = memoryTransportPair();
    const deps: AdoptionReceiverDeps = {
      deviceName: "recv",
      generateDeviceId: () => "recv-id",
      async confirmSas() {
        return true;
      },
      async createVaultFromRootKey() {
        return { vaultId: "v" };
      },
      async deriveLocalIdentity() {
        return {
          deviceId: "recv-id",
          deviceName: "recv",
          deviceType: "extension",
          longTermPublicKeyHex: "ab".repeat(32),
        };
      },
      async upsertPairedDevice() {},
      onError() {},
    };
    void requestAdoption(receiverChannel, deps);
    const first = await receiveFrame(sourceChannel);
    expect(first).toMatchObject({ type: "ADOPT_REQUEST", snapshot_capable: true });
  });

  it("round-trips a duress slot on the final chunk (single-chunk vault)", async () => {
    const [source, receiver] = memoryTransportPair();
    const duressJson = JSON.stringify({
      version: 1,
      salt: "aa".repeat(16),
      nonce: "bb".repeat(12),
      ciphertext: "cc".repeat(32),
      tag: "dd".repeat(16),
    });
    const [result] = await Promise.all([
      receiveVaultSnapshot(receiver),
      sendVaultSnapshot(source, { e1: "cipher1" }, { t1: 100 }, duressJson),
    ]);
    expect(result.entries).toEqual({ e1: "cipher1" });
    expect(result.tombstones).toEqual({ t1: 100 });
    expect(result.duressSlot).toBe(duressJson);
  });

  it("round-trips a duress slot across multiple chunks (final chunk only)", async () => {
    const [source, receiver] = memoryTransportPair();
    const big = "x".repeat(200 * 1024);
    const entries: Record<string, string> = {};
    for (let i = 0; i < 5; i++) entries[`id-${i}`] = big;
    const duressJson = JSON.stringify({
      version: 1,
      salt: "11".repeat(16),
      nonce: "22".repeat(12),
      ciphertext: "33".repeat(32),
      tag: "44".repeat(16),
    });
    const [result] = await Promise.all([
      receiveVaultSnapshot(receiver),
      sendVaultSnapshot(source, entries, { dead: 1 }, duressJson),
    ]);
    expect(result.entries).toEqual(entries);
    expect(result.duressSlot).toBe(duressJson);
  });

  it("omits duress_slot when no duress is configured (null)", async () => {
    const [source, receiver] = memoryTransportPair();
    const [result] = await Promise.all([
      receiveVaultSnapshot(receiver),
      sendVaultSnapshot(source, { e1: "v1" }, {}),
    ]);
    expect(result.duressSlot).toBeNull();
  });

  it("isAdoptSnapshot accepts frames with valid duress_slot string", () => {
    expect(
      isAdoptSnapshot({
        type: "ADOPT_SNAPSHOT",
        chunk_index: 0,
        chunk_count: 1,
        entries: {},
        tombstones: {},
        duress_slot: '{"version":1}',
      }),
    ).toBe(true);
  });

  it("isAdoptSnapshot accepts frames without duress_slot", () => {
    expect(
      isAdoptSnapshot({
        type: "ADOPT_SNAPSHOT",
        chunk_index: 0,
        chunk_count: 1,
        entries: {},
        tombstones: {},
      }),
    ).toBe(true);
  });

  it("isAdoptSnapshot rejects frames with non-string duress_slot", () => {
    expect(
      isAdoptSnapshot({
        type: "ADOPT_SNAPSHOT",
        chunk_index: 0,
        chunk_count: 1,
        entries: {},
        tombstones: {},
        duress_slot: 42,
      }),
    ).toBe(false);
    expect(
      isAdoptSnapshot({
        type: "ADOPT_SNAPSHOT",
        chunk_index: 0,
        chunk_count: 1,
        entries: {},
        tombstones: {},
        duress_slot: { version: 1 },
      }),
    ).toBe(false);
  });
});

describe("adoption snapshot aggregate budget (PB-26-206)", () => {
  // Cross-platform boundary vectors V1..V11. The same wire shapes and verdicts
  // are asserted in Android's VaultAdoptSnapshotTest — both clients must agree
  // on accept/reject for every vector. Vectors use the shared helper directly
  // where a full 10 MiB wire transfer would only prove the same comparison.

  it("V1: frame bytes at the aggregate limit pass, limit+1 rejects before merge", () => {
    const atLimit = createAdoptSnapshotBudget();
    atLimit.bytes = MAX_ADOPT_SNAPSHOT_BYTES - 100;
    applyAdoptSnapshotFrame(atLimit, { a: "x" }, {}, null, 100);
    expect(atLimit.bytes).toBe(MAX_ADOPT_SNAPSHOT_BYTES);

    const over = createAdoptSnapshotBudget();
    over.bytes = MAX_ADOPT_SNAPSHOT_BYTES - 100;
    expect(() => applyAdoptSnapshotFrame(over, { a: "x" }, {}, null, 101))
      .toThrow("ADOPT_SNAPSHOT_OVER_BUDGET");
    // Rejected before merge: no partial adopted state escapes.
    expect(over.entries).toEqual({});
  });

  it("V2: 25,000 entries pass, the 25,001st rejects", () => {
    const entries: Record<string, string> = {};
    for (let i = 0; i < MAX_ADOPT_SNAPSHOT_ENTRIES; i++) entries[`e${i}`] = "v";
    const acc = createAdoptSnapshotBudget();
    applyAdoptSnapshotFrame(acc, entries, {}, null, 100);
    expect(acc.entryCount).toBe(MAX_ADOPT_SNAPSHOT_ENTRIES);
    expect(() => applyAdoptSnapshotFrame(acc, { oneMore: "v" }, {}, null, 1))
      .toThrow("ADOPT_SNAPSHOT_TOO_MANY_ENTRIES");
    expect(acc.entries.oneMore).toBeUndefined();
  });

  it("V3: 25,000 tombstones pass, the 25,001st rejects", () => {
    const tombstones: Record<string, number> = {};
    for (let i = 0; i < MAX_ADOPT_SNAPSHOT_TOMBSTONES; i++) tombstones[`dead-${i}`] = i + 1;
    const acc = createAdoptSnapshotBudget();
    applyAdoptSnapshotFrame(acc, {}, tombstones, null, 100);
    expect(acc.tombstoneCount).toBe(MAX_ADOPT_SNAPSHOT_TOMBSTONES);
    expect(() => applyAdoptSnapshotFrame(acc, {}, { oneMore: 1 }, null, 1))
      .toThrow("ADOPT_SNAPSHOT_TOO_MANY_TOMBSTONES");
    expect(acc.tombstones.oneMore).toBeUndefined();
  });

  it("V4: entry/tombstone IDs of 512 chars pass, 513 reject", () => {
    const ok = "k".repeat(MAX_ADOPT_SNAPSHOT_ID_CHARS);
    const acc = createAdoptSnapshotBudget();
    applyAdoptSnapshotFrame(acc, { [ok]: "v" }, { [ok]: 1 }, null, 10);
    expect(acc.entries[ok]).toBe("v");

    const tooLong = "k".repeat(MAX_ADOPT_SNAPSHOT_ID_CHARS + 1);
    expect(() => applyAdoptSnapshotFrame(
      createAdoptSnapshotBudget(), { [tooLong]: "v" }, {}, null, 10,
    )).toThrow("ADOPT_SNAPSHOT_ID_TOO_LONG");
    expect(() => applyAdoptSnapshotFrame(
      createAdoptSnapshotBudget(), {}, { [tooLong]: 1 }, null, 10,
    )).toThrow("ADOPT_SNAPSHOT_ID_TOO_LONG");
  });

  it("V5: envelope at the base64-overhead ceiling passes, +1 rejects", () => {
    const atLimit = "c".repeat(Math.floor(MAX_ADOPT_SNAPSHOT_ENVELOPE_BYTES));
    const acc = createAdoptSnapshotBudget();
    applyAdoptSnapshotFrame(acc, { e1: atLimit }, {}, null, 10);
    expect(acc.entries.e1).toBe(atLimit);

    const over = "c".repeat(Math.floor(MAX_ADOPT_SNAPSHOT_ENVELOPE_BYTES) + 1);
    expect(() => applyAdoptSnapshotFrame(
      createAdoptSnapshotBudget(), { e1: over }, {}, null, 10,
    )).toThrow("ADOPT_SNAPSHOT_ENVELOPE_TOO_LARGE");
  });

  it("V6: duress slot at the wire-frame ceiling passes, +1 rejects", () => {
    const acc = createAdoptSnapshotBudget();
    applyAdoptSnapshotFrame(acc, {}, {}, "d".repeat(SNAPSHOT_FRAME_WIRE_CEILING), 10);
    expect(acc.duressSlot).toHaveLength(SNAPSHOT_FRAME_WIRE_CEILING);

    expect(() => applyAdoptSnapshotFrame(
      createAdoptSnapshotBudget(), {}, {}, "d".repeat(SNAPSHOT_FRAME_WIRE_CEILING + 1), 10,
    )).toThrow("ADOPT_SNAPSHOT_DURESS_FRAME_TOO_LARGE");
  });

  it("V7: non-finite tombstone values reject", () => {
    expect(() => applyAdoptSnapshotFrame(
      createAdoptSnapshotBudget(), {}, { dead: Number.NaN }, null, 10,
    )).toThrow("ADOPT_SNAPSHOT_INVALID_TOMBSTONE");
    expect(() => applyAdoptSnapshotFrame(
      createAdoptSnapshotBudget(), {}, { dead: Number.POSITIVE_INFINITY }, null, 10,
    )).toThrow("ADOPT_SNAPSHOT_INVALID_TOMBSTONE");
  });

  it("V8: high-count/low-byte snapshot round-trips (25,000 tiny entries)", async () => {
    const [source, receiver] = memoryTransportPair();
    const entries: Record<string, string> = {};
    for (let i = 0; i < MAX_ADOPT_SNAPSHOT_ENTRIES; i++) entries[`e${i}`] = "v";
    const [result] = await Promise.all([
      receiveVaultSnapshot(receiver),
      sendVaultSnapshot(source, entries, {}),
    ]);
    expect(Object.keys(result.entries)).toHaveLength(MAX_ADOPT_SNAPSHOT_ENTRIES);
  });

  it("V9: low-count/high-byte stream over the aggregate rejects with no partial result", async () => {
    // 16 envelopes of ~674 KiB each: every frame fits the 768 KiB wire
    // ceiling, but the ~10.8 MiB total exceeds the 10 MiB aggregate.
    const [source, receiver] = memoryTransportPair();
    const envelope = "c".repeat(690_000);
    const chunks = 16;
    const send = (async () => {
      for (let i = 0; i < chunks; i++) {
        await sendSnapshotFrame(source, {
          type: "ADOPT_SNAPSHOT",
          chunk_index: i,
          chunk_count: chunks,
          entries: { [`big-${i}`]: envelope },
          tombstones: {},
        });
      }
    })();
    await expect(receiveVaultSnapshot(receiver)).rejects.toThrow("ADOPT_SNAPSHOT_OVER_BUDGET");
    await send;
  });

  it("V10: over-budget sender preflights with zero frames sent (sync fallback, never truncation)", async () => {
    const [source] = memoryTransportPair();
    let frames = 0;
    const countingSource: PairingTransport = {
      send: (msg) => {
        frames++;
        return source.send(msg);
      },
      receive: () => source.receive(),
    };
    const envelope = "c".repeat(690_000);
    const entries: Record<string, string> = {};
    for (let i = 0; i < 16; i++) entries[`big-${i}`] = envelope;
    await expect(sendVaultSnapshot(countingSource, entries, {}))
      .rejects.toThrow("ADOPT_SNAPSHOT_OVER_BUDGET");
    expect(frames).toBe(0);
  });

  it("V11: duress bytes count against the aggregate budget", () => {
    const acc = createAdoptSnapshotBudget();
    acc.bytes = MAX_ADOPT_SNAPSHOT_BYTES - 100;
    // A 200-byte duress slot does not fit the remaining 100-byte budget.
    expect(() => applyAdoptSnapshotFrame(acc, {}, {}, "d".repeat(200), 200))
      .toThrow("ADOPT_SNAPSHOT_OVER_BUDGET");
  });
});

/**
 * PB-26-604 — adoption snapshot frames were sent as raw JSON, so a relay
 * observing Noise ciphertext saw plaintext_size + 16 per message and could
 * single out the distinctive duress frame. They now ride the sync transport's
 * bucket ladder.
 */
describe("PB-26-604 adoption snapshot padding", () => {
  it("pads every frame to a bucket so similar logical frames look alike", async () => {
    const [source] = memoryTransportPair();
    let wireSizes: number[] = [];
    const probe: PairingTransport = {
      send: (msg) => {
        wireSizes.push(msg.length);
        return source.send(msg);
      },
      receive: () => source.receive(),
    };

    await sendVaultSnapshot(probe, { "id-0": "x".repeat(1000) }, {}, null);
    const withoutDuress = wireSizes;

    wireSizes = [];
    await sendVaultSnapshot(probe, { "id-0": "x".repeat(1000) }, {}, "d".repeat(1000));
    const withDuress = wireSizes;

    // Every frame is bucket-padded: no wire size is the raw JSON length.
    for (const size of [...withoutDuress, ...withDuress]) {
      expect(size % 256).toBe(0);
    }
    // The duress frame lands in the same bucket as an entries frame of similar
    // logical size, so the relay cannot tell which transfer carried duress.
    expect(withDuress.at(-1)).toBe(withoutDuress.at(-1));
  });
});