// Runner pool enrollment with proof-of-possession (docs/ANTIMATTER_RUNNER.md §8).
//
// **Why this exists.** A pool is addressed by `poolId = sha256(provisioning device's static
// public key)` — the same identity string the relay uses for that device's room (its `HostID`).
// The static key is *public* (it rides in pairing QRs and every Welcome), so pool creation must
// prove possession of the matching private key or anyone holding the public key could squat the
// pool, cap it to zero, or point the owner's Settings at an attacker-held credential.
//
// Same X25519 DH challenge as the relay's enrollment (cloud/nucleic-edge/src/relayEnroll.ts) with
// a *different* proof context, so a proof captured for one service can never be replayed against
// the other:
//   1. POST /v1/runner/enroll/challenge { staticKey }   → { challengeId, ephemeralPublicKey }
//   2. client computes shared = X25519(staticPriv, ephemeralPub),
//      proof = HMAC-SHA256(shared, "nucleic-runner-enroll-v1:" + staticKeyB64)
//   3. POST /v1/runner/enroll { challengeId, proof }    → { poolId, poolSecret }
//
// Re-enrolling rotates the poolSecret (the old one dies), mirroring relay re-enrollment. The
// credential is then shared to the owner's other mesh devices over the existing E2EE mesh
// channel, the way relay memberships are distributed today.

const enc = new TextEncoder();

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

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

const PROOF_CONTEXT = "nucleic-runner-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("");
}

/** poolId is the provisioning device's identity: the hex SHA-256 of its X25519 static public
 * key — the same string as its `DeviceIdentity.hostID` / relay roomId, so "whose pool is this"
 * and "who is this on the mesh" are one identifier. */
export function derivePoolId(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. The runner-specific context ties the proof to *this* service. */
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.
 * Same one-boundary cast as the relay Worker: the runtime implements the standard
 * `{ name, public }` derive-params shape that `@cloudflare/workers-types` models differently. */
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. The ephemeral private key waits 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 poolId, mint the pool credential. Consumes the
 * challenge. Re-enrollment for an existing pool rotates the secret in place. */
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 poolId = await derivePoolId(staticKey);
  const secretBytes = new Uint8Array(32);
  crypto.getRandomValues(secretBytes);
  const poolSecret = [...secretBytes].map((b) => b.toString(16).padStart(2, "0")).join("");

  await kv.put(`poolauth:${poolId}`, await sha256Hex(enc.encode(poolSecret)), { expirationTtl: POOL_AUTH_TTL });

  return json({ poolId, poolSecret });
}

/** Resolve a `Bearer <poolId>.<poolSecret>` to its poolId, or null. Every /v1/pool/* route runs
 * through this (or the instance-token path), so the pool a caller reaches is always the
 * PoP-proven one, never a client-supplied value. */
export async function authedPool(kv: KVNamespace, bearer: string | null): Promise<string | null> {
  if (!bearer) return null;
  const dot = bearer.indexOf(".");
  if (dot <= 0) return null;
  const poolId = bearer.slice(0, dot);
  const secret = bearer.slice(dot + 1);
  if (await kv.get(`revoked:${poolId}`)) return null;
  const wantHash = await kv.get(`poolauth:${poolId}`);
  if (!wantHash) 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 poolId;
}

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;
}
