/**
 * Pure protocol + crypto for talking to a raw `browser-cli serve` endpoint.
 *
 * This module imports nothing from n8n and only touches Node's `crypto` plus a
 * lazily-loaded ML-KEM implementation, so it can be unit-tested with a plain
 * esbuild/node toolchain. The socket mechanics live in `serveClient.ts`.
 *
 * The wire protocol mirrors the Python client in `browser_cli/remote` and
 * `browser_cli/auth`:
 *
 *   1. Server sends a framed `challenge` JSON: {nonce, min_client_version, pq_kex?}.
 *   2. Client replies with one framed message. With a private key this is an
 *      Ed25519 signature over `nonce + sha256(canonical_json(msg))`, optionally
 *      bound to an ML-KEM-768 shared secret. When the server offers `pq_kex`
 *      (it always does once authorized_keys is set) the request body is also
 *      ChaCha20-Poly1305 encrypted under that secret.
 *   3. Server replies with one framed payload, encrypted the same way.
 *
 * Every value the client signs must serialize byte-for-byte like Python's
 * `json.dumps(sort_keys=True, separators=(",", ":"))` with `ensure_ascii=True`
 * or the signature is rejected — see `canonicalJson` and `protocol.test.ts`.
 */
import {
  createPrivateKey,
  createPublicKey,
  createHash,
  createHmac,
  createCipheriv,
  createDecipheriv,
  randomBytes,
  sign as nodeSign,
  type KeyObject,
} from 'node:crypto';

export const PQ_KEX_ALG = 'ML-KEM-768';
export const PQ_TRANSPORT_ALG = 'ML-KEM-768+ChaCha20Poly1305';

/** Auth-protocol fields that are never part of the signed canonical payload. */
const AUTH_FIELDS = new Set(['pubkey', 'sig', 'pq_kex', 'encrypted']);

// --- Framing ---------------------------------------------------------------

/** Prefix `payload` with browser-cli's 4-byte little-endian length header. */
export function frame(payload: Buffer): Buffer {
  const header = Buffer.allocUnsafe(4);
  header.writeUInt32LE(payload.length, 0);
  return Buffer.concat([header, payload]);
}

// --- Canonical JSON (matches Python json.dumps sort_keys + ensure_ascii) ----

/** Deterministic JSON string identical to the Python signing canonicalization. */
export function canonicalJson(value: unknown): string {
  return encode(value);
}

function encode(value: unknown): string {
  if (value === null) return 'null';
  const type = typeof value;
  if (type === 'string') return encodeString(value as string);
  if (type === 'boolean') return value ? 'true' : 'false';
  if (type === 'number') {
    const n = value as number;
    if (!Number.isFinite(n)) throw new Error('cannot encode non-finite number');
    // Integers match Python exactly; non-integers are rare in args and fall
    // back to JS formatting (documented limitation, see protocol.test.ts).
    return Number.isInteger(n) ? String(n) : JSON.stringify(n);
  }
  if (Array.isArray(value)) return '[' + value.map(encode).join(',') + ']';
  if (type === 'object') {
    const obj = value as Record<string, unknown>;
    const keys = Object.keys(obj)
      .filter((key) => obj[key] !== undefined)
      .sort();
    return '{' + keys.map((key) => encodeString(key) + ':' + encode(obj[key])).join(',') + '}';
  }
  throw new Error(`cannot encode value of type ${type} in canonical JSON`);
}

/** JSON-encode a string and escape every non-ASCII char as \uXXXX, like Python. */
function encodeString(str: string): string {
  return JSON.stringify(str).replace(/[€-￿]/g, (char) =>
    '\\u' + char.charCodeAt(0).toString(16).padStart(4, '0'),
  );
}

// --- Ed25519 ---------------------------------------------------------------

/**
 * Reconstruct a clean PEM from a value mangled by a credential field:
 * surrounding whitespace, newlines escaped to literal `\n`, or — common with
 * single-line/password inputs — the internal line breaks stripped entirely so
 * the base64 body and the BEGIN/END markers run together.
 */
export function normalizePem(input: string): string {
  let pem = (input || '').trim().replace(/\\n/g, '\n');
  const match = pem.match(/-----BEGIN ([A-Z0-9 ]+?)-----([\s\S]*?)-----END \1-----/);
  if (match) {
    const label = match[1].trim();
    const body = (match[2].match(/[A-Za-z0-9+/=]+/g) || []).join('');
    const wrapped = body.match(/.{1,64}/g) || [];
    pem = `-----BEGIN ${label}-----\n${wrapped.join('\n')}\n-----END ${label}-----\n`;
  }
  return pem;
}

/**
 * Load a PKCS8 PEM Ed25519 private key, tolerating the ways credential fields
 * mangle multi-line secrets (see {@link normalizePem}).
 */
export function loadPrivateKey(privatePem: string): KeyObject {
  const pem = normalizePem(privatePem);
  try {
    return createPrivateKey(pem);
  } catch (err) {
    throw new Error(
      'Invalid Ed25519 private key: expected a PKCS8 PEM block ' +
        '("-----BEGIN PRIVATE KEY-----"), e.g. the file from `browser-cli auth keygen`. ' +
        `(${(err as Error).message})`,
    );
  }
}

/** Raw 32-byte Ed25519 public key (hex) derived from a PKCS8 PEM private key. */
export function ed25519PublicKeyHex(privatePem: string): string {
  const publicKey = createPublicKey(loadPrivateKey(privatePem));
  const jwk = publicKey.export({ format: 'jwk' }) as { x?: string };
  if (!jwk.x) throw new Error('private key is not an Ed25519 key');
  return Buffer.from(jwk.x, 'base64url').toString('hex');
}

/** Bytes signed for auth: nonce + sha256(canonical) [+ sha256(label + secret)]. */
export function authMessage(nonceHex: string, msg: Record<string, unknown>, pqSecret: Buffer | null): Buffer {
  const nonce = Buffer.from(nonceHex, 'hex');
  const canonical = createHash('sha256').update(canonicalJson(stripAuthFields(msg)), 'utf8').digest();
  let data = Buffer.concat([nonce, canonical]);
  if (pqSecret) {
    const bound = createHash('sha256')
      .update(Buffer.concat([Buffer.from('browser-cli ml-kem-768 v1', 'ascii'), pqSecret]))
      .digest();
    data = Buffer.concat([data, bound]);
  }
  return data;
}

/** Ed25519 signature (hex) over the canonical auth payload. */
export function signAuth(
  privatePem: string,
  nonceHex: string,
  msg: Record<string, unknown>,
  pqSecret: Buffer | null,
): string {
  return nodeSign(null, authMessage(nonceHex, msg, pqSecret), loadPrivateKey(privatePem)).toString('hex');
}

function stripAuthFields(msg: Record<string, unknown>): Record<string, unknown> {
  const out: Record<string, unknown> = {};
  for (const [key, value] of Object.entries(msg)) {
    if (!AUTH_FIELDS.has(key)) out[key] = value;
  }
  return out;
}

// --- ML-KEM-768 transport encryption ---------------------------------------

// Node has no native ML-KEM, so load @noble/post-quantum at runtime. The
// indirection through `Function` keeps a real dynamic `import()` even after
// TypeScript downlevels this module to CommonJS (a plain `import()` would be
// rewritten to `require()` and fail on the ESM-only package).
const importEsm = new Function('specifier', 'return import(specifier)') as (specifier: string) => Promise<any>;
let mlkemPromise: Promise<any> | null = null;

async function mlKem768(): Promise<any> {
  if (!mlkemPromise) {
    mlkemPromise = importEsm('@noble/post-quantum/ml-kem.js').then((mod) => mod.ml_kem768);
  }
  return mlkemPromise;
}

/** Encapsulate to the server's ML-KEM public key. Returns (ciphertext hex, secret). */
export async function pqEncapsulate(serverPublicKeyHex: string): Promise<{ ciphertextHex: string; secret: Buffer }> {
  const ml = await mlKem768();
  const { cipherText, sharedSecret } = ml.encapsulate(Buffer.from(serverPublicKeyHex, 'hex'));
  return { ciphertextHex: Buffer.from(cipherText).toString('hex'), secret: Buffer.from(sharedSecret) };
}

/** HKDF-SHA256 with a 32-zero-byte salt (Python `salt=None`) and the given info. */
export function pqTransportKey(secret: Buffer, direction: string): Buffer {
  const salt = Buffer.alloc(32, 0);
  const prk = createHmac('sha256', salt).update(secret).digest();
  const info = Buffer.concat([Buffer.from(`browser-cli pq transport v1 ${direction}`, 'ascii'), Buffer.from([1])]);
  return createHmac('sha256', prk).update(info).digest().subarray(0, 32);
}

export interface PqEnvelope {
  alg?: string;
  nonce: string;
  ciphertext: string;
}

/** ChaCha20-Poly1305 encrypt an app-layer frame (ciphertext field is ct||tag). */
export function pqEncrypt(secret: Buffer, direction: string, plaintext: Buffer): PqEnvelope {
  const key = pqTransportKey(secret, direction);
  const nonce = randomBytes(12);
  const cipher = createCipheriv('chacha20-poly1305', key, nonce, { authTagLength: 16 });
  const body = Buffer.concat([cipher.update(plaintext), cipher.final()]);
  const tag = cipher.getAuthTag();
  return {
    alg: PQ_TRANSPORT_ALG,
    nonce: nonce.toString('hex'),
    ciphertext: Buffer.concat([body, tag]).toString('hex'),
  };
}

/** Inverse of {@link pqEncrypt}. */
export function pqDecrypt(secret: Buffer, direction: string, envelope: PqEnvelope): Buffer {
  if (!envelope || envelope.alg !== PQ_TRANSPORT_ALG) {
    throw new Error('unsupported encrypted transport envelope');
  }
  const key = pqTransportKey(secret, direction);
  const nonce = Buffer.from(envelope.nonce, 'hex');
  const blob = Buffer.from(envelope.ciphertext, 'hex');
  const tag = blob.subarray(blob.length - 16);
  const body = blob.subarray(0, blob.length - 16);
  const decipher = createDecipheriv('chacha20-poly1305', key, nonce, { authTagLength: 16 });
  decipher.setAuthTag(tag);
  return Buffer.concat([decipher.update(body), decipher.final()]);
}

// --- Handshake payload + response decoding ---------------------------------

export interface Challenge {
  type?: string;
  nonce?: string;
  min_client_version?: string;
  pq_kex?: { alg?: string; public_key?: string };
}

export interface AuthPayload {
  payload: Record<string, unknown>;
  pqSecret: Buffer | null;
}

function pqPublicKey(challenge: Challenge): string | null {
  const kex = challenge.pq_kex;
  if (kex && kex.alg === PQ_KEX_ALG && kex.public_key) return String(kex.public_key);
  return null;
}

/**
 * Build the single framed message a client sends in response to the challenge.
 * Mirrors `browser_cli.remote.auth.build_auth_message` + `signed_payload`.
 */
export async function buildAuthPayload(
  baseMsg: Record<string, unknown>,
  challenge: Challenge,
  privatePem: string | null,
): Promise<AuthPayload> {
  const nonceHex = challenge.type === 'challenge' ? challenge.nonce : undefined;
  if (!nonceHex || !privatePem) {
    // No-auth endpoint (loopback `serve --no-auth`): send the bare message.
    return { payload: baseMsg, pqSecret: null };
  }

  const clean = stripAuthFields(baseMsg);
  let secret: Buffer | null = null;
  const serverPub = pqPublicKey(challenge);
  if (serverPub) {
    const enc = await pqEncapsulate(serverPub);
    secret = enc.secret;
    clean.pq_kex = { alg: PQ_KEX_ALG, ciphertext: enc.ciphertextHex };
  }

  const sig = signAuth(privatePem, nonceHex, clean, secret);
  const pubkey = ed25519PublicKeyHex(privatePem);

  if (!secret) {
    return { payload: { ...clean, pubkey, sig }, pqSecret: null };
  }

  const encrypted = pqEncrypt(secret, 'request', Buffer.from(JSON.stringify(clean), 'utf8'));
  return {
    payload: {
      id: clean.id,
      user_agent: clean.user_agent,
      pubkey,
      sig,
      pq_kex: clean.pq_kex,
      encrypted,
    },
    pqSecret: secret,
  };
}

/** Decode a framed server response, decrypting the PQ envelope when present. */
export function decodeResponse(raw: Buffer, pqSecret: Buffer | null): any {
  const outer = JSON.parse(raw.toString('utf8'));
  if (pqSecret && outer && typeof outer === 'object' && 'encrypted' in outer) {
    return JSON.parse(pqDecrypt(pqSecret, 'response', outer.encrypted).toString('utf8'));
  }
  return outer;
}