import { test } from "node:test";
import assert from "node:assert/strict";
import {
  authedPool,
  b64Decode,
  b64urlEncode,
  derivePoolId,
  enrollChallenge,
  enrollComplete,
  enrollProof,
  sha256Hex,
} from "../src/enroll.ts";

// Minimal KV stub (same shape the edge tests use).
function kvStub() {
  const map = new Map<string, string>();
  return {
    map,
    async get(key: string) {
      return map.get(key) ?? null;
    },
    async put(key: string, value: string, _opts?: unknown) {
      map.set(key, value);
    },
    async delete(key: string) {
      map.delete(key);
    },
  } as unknown as KVNamespace & { map: Map<string, string> };
}

async function makeDeviceKey() {
  const pair = (await crypto.subtle.generateKey({ name: "X25519" }, true, ["deriveBits"])) as CryptoKeyPair;
  const pubRaw = new Uint8Array((await crypto.subtle.exportKey("raw", pair.publicKey)) as ArrayBuffer);
  return { pair, pubRaw, pubB64: b64urlEncode(pubRaw) };
}

async function dh(priv: CryptoKey, peerPubRaw: Uint8Array): Promise<Uint8Array> {
  const pub = await crypto.subtle.importKey("raw", peerPubRaw, { name: "X25519" }, false, []);
  const params = { name: "X25519", public: pub } as unknown as Parameters<SubtleCrypto["deriveBits"]>[0];
  return new Uint8Array(await crypto.subtle.deriveBits(params, priv, 256));
}

test("full PoP enrollment mints a pool credential bound to the static key", async () => {
  const kv = kvStub();
  const device = await makeDeviceKey();

  const challengeRes = await enrollChallenge(kv, { staticKey: device.pubB64 });
  assert.equal(challengeRes.status, 200);
  const challenge = (await challengeRes.json()) as { challengeId: string; ephemeralPublicKey: string };

  const shared = await dh(device.pair.privateKey, b64Decode(challenge.ephemeralPublicKey));
  const proof = b64urlEncode(await enrollProof(shared, device.pubB64));

  const completeRes = await enrollComplete(kv, { challengeId: challenge.challengeId, proof });
  assert.equal(completeRes.status, 200);
  const cred = (await completeRes.json()) as { poolId: string; poolSecret: string };

  assert.equal(cred.poolId, await derivePoolId(device.pubRaw));
  assert.equal(cred.poolSecret.length, 64);

  // The credential authenticates; a wrong secret doesn't.
  assert.equal(await authedPool(kv, `${cred.poolId}.${cred.poolSecret}`), cred.poolId);
  assert.equal(await authedPool(kv, `${cred.poolId}.${"0".repeat(64)}`), null);
});

test("a wrong proof is rejected and the challenge is consumed", async () => {
  const kv = kvStub();
  const device = await makeDeviceKey();
  const challenge = (await (await enrollChallenge(kv, { staticKey: device.pubB64 })).json()) as {
    challengeId: string;
  };

  const bad = await enrollComplete(kv, { challengeId: challenge.challengeId, proof: b64urlEncode(new Uint8Array(32)) });
  assert.equal(bad.status, 401);

  // Consumed: even the right proof can't be replayed against the same challenge.
  const retry = await enrollComplete(kv, { challengeId: challenge.challengeId, proof: b64urlEncode(new Uint8Array(32)) });
  assert.equal(retry.status, 401);
});

test("challenge validates the static key shape", async () => {
  const kv = kvStub();
  assert.equal((await enrollChallenge(kv, null)).status, 422);
  assert.equal((await enrollChallenge(kv, { staticKey: 42 })).status, 422);
  assert.equal((await enrollChallenge(kv, { staticKey: b64urlEncode(new Uint8Array(16)) })).status, 422);
});

test("proof from another device's key fails (possession is the gate)", async () => {
  const kv = kvStub();
  const victim = await makeDeviceKey();
  const attacker = await makeDeviceKey();

  // Attacker holds the victim's PUBLIC key and requests a challenge for it…
  const challenge = (await (await enrollChallenge(kv, { staticKey: victim.pubB64 })).json()) as {
    challengeId: string;
    ephemeralPublicKey: string;
  };
  // …but can only complete the DH with their own private key.
  const shared = await dh(attacker.pair.privateKey, b64Decode(challenge.ephemeralPublicKey));
  const proof = b64urlEncode(await enrollProof(shared, victim.pubB64));
  const res = await enrollComplete(kv, { challengeId: challenge.challengeId, proof });
  assert.equal(res.status, 401);
});

test("re-enrollment rotates the pool secret", async () => {
  const kv = kvStub();
  const device = await makeDeviceKey();

  async function enroll(): Promise<{ poolId: string; poolSecret: string }> {
    const challenge = (await (await enrollChallenge(kv, { staticKey: device.pubB64 })).json()) as {
      challengeId: string;
      ephemeralPublicKey: string;
    };
    const shared = await dh(device.pair.privateKey, b64Decode(challenge.ephemeralPublicKey));
    const proof = b64urlEncode(await enrollProof(shared, device.pubB64));
    return (await (await enrollComplete(kv, { challengeId: challenge.challengeId, proof })).json()) as {
      poolId: string;
      poolSecret: string;
    };
  }

  const first = await enroll();
  const second = await enroll();
  assert.equal(first.poolId, second.poolId);
  assert.notEqual(first.poolSecret, second.poolSecret);
  assert.equal(await authedPool(kv, `${first.poolId}.${first.poolSecret}`), null); // rotated away
  assert.equal(await authedPool(kv, `${second.poolId}.${second.poolSecret}`), second.poolId);
});

test("revocation kills the credential", async () => {
  const kv = kvStub();
  const device = await makeDeviceKey();
  const challenge = (await (await enrollChallenge(kv, { staticKey: device.pubB64 })).json()) as {
    challengeId: string;
    ephemeralPublicKey: string;
  };
  const shared = await dh(device.pair.privateKey, b64Decode(challenge.ephemeralPublicKey));
  const proof = b64urlEncode(await enrollProof(shared, device.pubB64));
  const cred = (await (await enrollComplete(kv, { challengeId: challenge.challengeId, proof })).json()) as {
    poolId: string;
    poolSecret: string;
  };
  await kv.put(`revoked:${cred.poolId}`, "1");
  assert.equal(await authedPool(kv, `${cred.poolId}.${cred.poolSecret}`), null);
});

test("poolId derivation matches sha256 hex of the raw key (Swift DeviceIdentity.hostID)", async () => {
  const key = new Uint8Array(32).fill(7);
  assert.equal(await derivePoolId(key), await sha256Hex(key));
});
