// Relay host enrollment with proof-of-possession (mesh P2 hardening, docs/CLOUD_INFRA.md §2.1).
//
// **Why this exists.** The relay room is addressed by `roomId = sha256(host static public key)`.
// That static key is *public* — it rides in every pairing QR and is held by every paired
// device. The old `/v1/host/enroll` minted a credential to anyone who asked, so an attacker
// holding the (public) host key could enroll for the victim's roomId, mint a `role=host`
// connection token, and evict the real host from its room. They could never break Noise, but
// they could blackhole the victim's relay path and capture traffic metadata.
//
// The fix binds enrollment to the roomId's own private key via an X25519 Diffie-Hellman
// challenge: the enroller must prove possession of the *static private key* that derives the
// roomId. Because roomId is `sha256(staticPub)` and only the holder of the matching private key
// can complete the DH, an attacker with just the public key cannot enroll. This binds directly
// to the roomId source — no separate signing-key ↔ static-key binding problem (an attacker could
// otherwise pair the victim's static key with their own signing key).
//
// Flow (stateless-ish: one short-lived KV entry per in-flight challenge):
//   1. POST /v1/relay/enroll/challenge { staticKey }         → { challengeId, ephemeralPublicKey }
//   2. client computes shared = X25519(staticPriv, ephemeralPub),
//      proof = HMAC-SHA256(shared, "nucleic-relay-enroll-v1:" + staticKeyB64)
//   3. POST /v1/relay/enroll { challengeId, proof }          → { hostId, hostSecret, roomId }
//
// The Worker recomputes shared = X25519(ephemeralPriv, staticPub) and verifies the proof. On
// success it derives roomId, mints the host credential, and records `hostroom:<hostId> → roomId`
// so token minting (index.ts) derives the room server-side and never trusts a client-supplied one.

const enc = new TextEncoder();

/** How long a pending enroll challenge is valid (KV TTL). */
export const ENROLL_CHALLENGE_TTL = 120; // seconds

/** Host credential lifetime — refreshed on every authenticated mint (see index.ts). */
export const HOST_ROOM_TTL = 60 * 60 * 24 * 365; // 1 year

const PROOF_CONTEXT = "nucleic-relay-enroll-v1:";

export function b64urlEncode(bytes: Uint8Array): string {
  let s = "";
  for (const b of bytes) s += String.fromCharCode(b);
  return btoa(s).replace(/\+/g, "-").replace(/\//g, "_").replace(/=+$/, "");
}

export function b64Decode(s: string): Uint8Array {
  const norm = s.replace(/-/g, "+").replace(/_/g, "/");
  const pad = norm.length % 4 ? 4 - (norm.length % 4) : 0;
  const bin = atob(norm + "=".repeat(pad));
  const out = new Uint8Array(bin.length);
  for (let i = 0; i < bin.length; i++) out[i] = bin.charCodeAt(i);
  return out;
}

export async function sha256Hex(bytes: Uint8Array): Promise<string> {
  const digest = await crypto.subtle.digest("SHA-256", bytes);
  return [...new Uint8Array(digest)].map((b) => b.toString(16).padStart(2, "0")).join("");
}

/** roomId is the host's identity: the hex SHA-256 of its X25519 static public key. Matches the
 * Swift `DeviceIdentity.hostID`, so the room the host advertises and the room it enrolls for are
 * the same string on both sides. */
export function deriveRoomId(staticKey: Uint8Array): Promise<string> {
  return sha256Hex(staticKey);
}

/** The DH proof both sides compute: HMAC over a fixed context + the static key, keyed by the
 * X25519 shared secret. Ties the proof to *this* enrollment (the static key) so a shared secret
 * captured for one purpose can't be replayed for another. */
export async function enrollProof(shared: Uint8Array, staticKeyB64: string): Promise<Uint8Array> {
  const key = await crypto.subtle.importKey("raw", shared, { name: "HMAC", hash: "SHA-256" }, false, ["sign"]);
  const msg = enc.encode(PROOF_CONTEXT + staticKeyB64);
  return new Uint8Array(await crypto.subtle.sign("HMAC", key, msg));
}

/** X25519 shared secret between our ephemeral private key and the peer's static public key.
 * The `{ name, public }` derive-params shape is the standard WebCrypto one the Workers/Node
 * runtimes implement; `@cloudflare/workers-types` models it under a different field name, so we
 * cast at this one boundary rather than contort the call. */
async function x25519(privateKey: CryptoKey, peerPublicKey: Uint8Array): Promise<Uint8Array> {
  const pub = await crypto.subtle.importKey("raw", peerPublicKey, { name: "X25519" }, false, []);
  const params = { name: "X25519", public: pub } as unknown as Parameters<SubtleCrypto["deriveBits"]>[0];
  return new Uint8Array(await crypto.subtle.deriveBits(params, privateKey, 256));
}

function json(obj: unknown, status = 200): Response {
  return new Response(JSON.stringify(obj), { status, headers: { "content-type": "application/json" } });
}

interface ChallengeRecord {
  staticKeyB64: string;
  ephPrivJwk: JsonWebKey;
}

/** Step 1: mint a DH challenge for the claimed static key. We keep the ephemeral private key in
 * KV (short TTL) keyed by an opaque challengeId; the client never sees it. */
export async function enrollChallenge(kv: KVNamespace, body: { staticKey?: unknown } | null): Promise<Response> {
  if (typeof body?.staticKey !== "string") return new Response("staticKey required", { status: 422 });
  let staticKey: Uint8Array;
  try {
    staticKey = b64Decode(body.staticKey);
  } catch {
    return new Response("staticKey not base64", { status: 422 });
  }
  if (staticKey.length !== 32) return new Response("staticKey must be 32 bytes", { status: 422 });

  const eph = (await crypto.subtle.generateKey({ name: "X25519" }, true, ["deriveBits"])) as CryptoKeyPair;
  const ephPubRaw = new Uint8Array((await crypto.subtle.exportKey("raw", eph.publicKey)) as ArrayBuffer);
  const ephPrivJwk = (await crypto.subtle.exportKey("jwk", eph.privateKey)) as JsonWebKey;

  const challengeId = crypto.randomUUID();
  const record: ChallengeRecord = { staticKeyB64: body.staticKey, ephPrivJwk };
  await kv.put(`enroll-chal:${challengeId}`, JSON.stringify(record), { expirationTtl: ENROLL_CHALLENGE_TTL });

  return json({ challengeId, ephemeralPublicKey: b64urlEncode(ephPubRaw) });
}

/** Step 3: verify the client's proof, derive the roomId, mint the host credential, and bind the
 * host to its room so later token minting can't spoof a different room. Consumes the challenge. */
export async function enrollComplete(
  kv: KVNamespace,
  body: { challengeId?: unknown; proof?: unknown } | null,
): Promise<Response> {
  if (typeof body?.challengeId !== "string" || typeof body?.proof !== "string") {
    return new Response("challengeId + proof required", { status: 422 });
  }
  const raw = await kv.get(`enroll-chal:${body.challengeId}`);
  if (!raw) return new Response("challenge expired or unknown", { status: 401 });
  const record = JSON.parse(raw) as ChallengeRecord;

  // Single-use: delete before verifying so a wrong proof can't be brute-forced against one nonce.
  await kv.delete(`enroll-chal:${body.challengeId}`);

  const ephPriv = await crypto.subtle.importKey("jwk", record.ephPrivJwk, { name: "X25519" }, false, ["deriveBits"]);
  const staticKey = b64Decode(record.staticKeyB64);
  const shared = await x25519(ephPriv, staticKey);
  const expected = await enrollProof(shared, record.staticKeyB64);

  let given: Uint8Array;
  try {
    given = b64Decode(body.proof);
  } catch {
    return new Response("proof not base64", { status: 401 });
  }
  if (!constantTimeEqual(expected, given)) return new Response("proof mismatch", { status: 401 });

  const roomId = await deriveRoomId(staticKey);
  const hostId = crypto.randomUUID();
  const secretBytes = new Uint8Array(32);
  crypto.getRandomValues(secretBytes);
  const hostSecret = [...secretBytes].map((b) => b.toString(16).padStart(2, "0")).join("");

  await kv.put(`hostauth:${hostId}`, await sha256Hex(enc.encode(hostSecret)), { expirationTtl: HOST_ROOM_TTL });
  await kv.put(`hostroom:${hostId}`, roomId, { expirationTtl: HOST_ROOM_TTL });

  return json({ hostId, hostSecret, roomId });
}

/** Resolve a `Bearer <hostId>.<hostSecret>` to its bound roomId, or null. Used by token minting
 * so the room is always the PoP-proven one, never a client-supplied value. */
export async function authedHostRoom(kv: KVNamespace, bearer: string | null): Promise<{ hostId: string; roomId: string } | null> {
  if (!bearer) return null;
  const dot = bearer.indexOf(".");
  if (dot <= 0) return null;
  const hostId = bearer.slice(0, dot);
  const secret = bearer.slice(dot + 1);
  const wantHash = await kv.get(`hostauth:${hostId}`);
  const roomId = await kv.get(`hostroom:${hostId}`);
  if (!wantHash || !roomId) return null;
  const gotHash = await sha256Hex(enc.encode(secret));
  // Both operands are fixed-length SHA-256 hex, so a char-wise compare is constant-time.
  if (!constantTimeEqualStr(wantHash, gotHash)) return null;
  return { hostId, roomId };
}

export function constantTimeEqual(a: Uint8Array, b: Uint8Array): boolean {
  if (a.length !== b.length) return false;
  let diff = 0;
  for (let i = 0; i < a.length; i++) diff |= a[i] ^ b[i];
  return diff === 0;
}

function constantTimeEqualStr(a: string, b: string): boolean {
  if (a.length !== b.length) return false;
  let diff = 0;
  for (let i = 0; i < a.length; i++) diff |= a.charCodeAt(i) ^ b.charCodeAt(i);
  return diff === 0;
}
