import { describe, expect, it } from "vitest";
import { generateKeyPair } from "@stablelib/x25519";
import {
  derivePairingSAS,
  memoryTransportPair,
  runPairingInitiator,
  runPairingResponder,
  type LocalIdentity,
} from "./pairing-handshake";
import { deriveSyncIdentityKeypair, pubkeyToHex } from "./sync-identity";

const MASTER = Uint8Array.from(Array.from({ length: 32 }, (_, i) => (i * 5 + 3) & 0xff));

function identityFor(deviceId: string, deviceName: string, type: LocalIdentity["deviceType"]): LocalIdentity {
  const kp = deriveSyncIdentityKeypair(MASTER, deviceId);
  return {
    deviceId,
    deviceName,
    deviceType: type,
    longTermPublicKeyHex: pubkeyToHex(kp.publicKey),
    signingPublicKeyHex: "b".repeat(64),
  };
}

describe("pairing handshake (Noise NK + identity exchange)", () => {
  it("exchanges long-term sync identities between initiator and responder", async () => {
    // Responder owns the ephemeral that the initiator targets.
    const respEphemeral = generateKeyPair();

    const initIdent = identityFor("android-001", "June's Pixel 9", "android");
    const respIdent = identityFor("ext-001", "Firefox on MacBook", "extension");

    const [t1, t2] = memoryTransportPair();
    const respPromise = runPairingResponder(
      respEphemeral.secretKey,
      respEphemeral.publicKey,
      respIdent,
      t2,
    );
    const initPromise = runPairingInitiator(respEphemeral.publicKey, initIdent, t1);

    const [respResult, initResult] = await Promise.all([respPromise, initPromise]);
    const initLearned = initResult.device;
    const respLearned = respResult.device;
    expect(initLearned.deviceId).toBe("ext-001");
    expect(initLearned.deviceName).toBe("Firefox on MacBook");
    expect(initLearned.deviceType).toBe("extension");
    expect(initLearned.longTermPublicKeyHex).toBe(respIdent.longTermPublicKeyHex);
    expect(initLearned.signingPublicKeyHex).toBe(respIdent.signingPublicKeyHex);
    expect(initLearned.trustSource).toBe("direct");

    expect(respLearned.deviceId).toBe("android-001");
    expect(respLearned.deviceName).toBe("June's Pixel 9");
    expect(respLearned.deviceType).toBe("android");
    expect(respLearned.longTermPublicKeyHex).toBe(initIdent.longTermPublicKeyHex);
    expect(respLearned.signingPublicKeyHex).toBe(initIdent.signingPublicKeyHex);
    expect(respLearned.trustSource).toBe("direct");

    // Both ends derive the SAME 8-digit SAS from the shared handshake hash —
    // this is what the user compares across devices.
    expect(initResult.sas).toMatch(/^\d{8}$/);
    expect(initResult.sas).toBe(respResult.sas);
  });

  it("carries a Pro license from the device that has one to the peer", async () => {
    const respEphemeral = generateKeyPair();
    // The initiator (e.g. an existing Pro phone) advertises its signed license.
    const initIdent = { ...identityFor("android-pro", "Pro Pixel", "android"), license: "PPK-AAA-BBB" };
    const respIdent = identityFor("ext-free", "Free Firefox", "extension"); // no license

    const [t1, t2] = memoryTransportPair();
    const respPromise = runPairingResponder(respEphemeral.secretKey, respEphemeral.publicKey, respIdent, t2);
    const initPromise = runPairingInitiator(respEphemeral.publicKey, initIdent, t1);
    const [respResult, initResult] = await Promise.all([respPromise, initPromise]);

    // The free responder learns the peer's license to activate; the initiator
    // sees none (we sent it nothing). License is never written onto the device.
    expect(respResult.peerLicense).toBe("PPK-AAA-BBB");
    expect(initResult.peerLicense).toBeUndefined();
    expect("license" in respResult.device).toBe(false);
  });

  it("rejects a malformed identity frame", async () => {
    const respEphemeral = generateKeyPair();
    const respIdent = identityFor("ext-002", "Firefox", "extension");
    const [t1, t2] = memoryTransportPair();

    // Initiator sends garbage instead of a PAIR_IDENTITY frame
    const respPromise = runPairingResponder(
      respEphemeral.secretKey,
      respEphemeral.publicKey,
      respIdent,
      t2,
    );
    const evilInitiator = (async () => {
      const { noiseNKInitiator } = await import("./noise-js");
      const session = await noiseNKInitiator(respEphemeral.publicKey, t1.send, t1.receive);
      await t1.send(session.encrypt(new TextEncoder().encode("{}"))); // not a PAIR_IDENTITY
    })();

    await expect(respPromise).rejects.toThrow(/PAIR_IDENTITY/);
    await evilInitiator.catch(() => undefined);
  });

  it("rejects an identity with a malformed long_term_public_key_hex", async () => {
    const respEphemeral = generateKeyPair();
    const respIdent = identityFor("ext-003", "Firefox", "extension");
    const [t1, t2] = memoryTransportPair();

    const respPromise = runPairingResponder(
      respEphemeral.secretKey,
      respEphemeral.publicKey,
      respIdent,
      t2,
    );
    const evilInitiator = (async () => {
      const { noiseNKInitiator } = await import("./noise-js");
      const session = await noiseNKInitiator(respEphemeral.publicKey, t1.send, t1.receive);
      const frame = JSON.stringify({
        type: "PAIR_IDENTITY",
        device_id: "evil",
        device_name: "evil",
        device_type: "android",
        long_term_public_key_hex: "deadbeef", // too short
      });
      await t1.send(session.encrypt(new TextEncoder().encode(frame)));
    })();

    await expect(respPromise).rejects.toThrow(/32-byte hex/);
    await evilInitiator.catch(() => undefined);
  });

  it("rejects an identity with a malformed signing_public_key_hex", async () => {
    const respEphemeral = generateKeyPair();
    const respIdent = identityFor("ext-004", "Firefox", "extension");
    const [t1, t2] = memoryTransportPair();

    const respPromise = runPairingResponder(
      respEphemeral.secretKey,
      respEphemeral.publicKey,
      respIdent,
      t2,
    );
    const evilInitiator = (async () => {
      const { noiseNKInitiator } = await import("./noise-js");
      const session = await noiseNKInitiator(respEphemeral.publicKey, t1.send, t1.receive);
      const frame = JSON.stringify({
        type: "PAIR_IDENTITY",
        device_id: "evil",
        device_name: "evil",
        device_type: "android",
        long_term_public_key_hex: "a".repeat(64),
        signing_public_key_hex: "deadbeef",
      });
      await t1.send(session.encrypt(new TextEncoder().encode(frame)));
    })();

    await expect(respPromise).rejects.toThrow(/signing_public_key_hex/);
    await evilInitiator.catch(() => undefined);
  });
});

describe("derivePairingSAS (out-of-band MITM-detection code)", () => {
  it("returns a deterministic 8-digit code", () => {
    const h = Uint8Array.from(Array.from({ length: 32 }, (_, i) => (i * 7 + 1) & 0xff));
    const sas = derivePairingSAS(h);
    expect(sas).toMatch(/^\d{8}$/);
    expect(derivePairingSAS(h)).toBe(sas);
  });

  it("both peers derive the same SAS from the shared Noise handshake hash", () => {
    // After the handshake both ends hold the identical handshakeHash, so both
    // independently compute the same comparison code.
    const shared = crypto.getRandomValues(new Uint8Array(32));
    expect(derivePairingSAS(shared.slice())).toBe(derivePairingSAS(shared.slice()));
  });

  it("differs for different handshake hashes", () => {
    expect(derivePairingSAS(Uint8Array.of(0, 0, 0, 0, 0, 0, 0, 0))).not.toBe(
      derivePairingSAS(Uint8Array.of(0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff)),
    );
  });

  it("pins the cross-platform big-endian 64-bit vector", () => {
    expect(derivePairingSAS(Uint8Array.of(0, 1, 2, 3, 4, 5, 6, 7))).toBe("52306183");
    expect(derivePairingSAS(new Uint8Array(8).fill(0xff))).toBe("09551615");
  });

  it("rejects undersized input", () => {
    expect(() => derivePairingSAS(new Uint8Array(7))).toThrow(/too short/);
  });
});