// @pdsjs/spaces/kem - the shared secret a syncer and a repo host MAC one
// response under.
//
// ECDH on P-256 between the syncer's ephemeral key and the author's
// `#atproto_kem` key, with the secret derived as DHKEM(P-256, HKDF-SHA256)
// specifies (RFC 9180 §4.1). The syncer sends its ephemeral public key as a
// did:key, the `key_share` of Atproto-Space-Want-Mac.

import { compressPublicKey } from '@pdsjs/core/crypto';
import { publicKeyToDidKey } from '@pdsjs/core/plc';
import { parseDidKey } from '@pdsjs/core/verify';
import { concat, hkdfSha256, SpaceMacError, uint16be, utf8 } from './mac.js';

// HPKE-v1 || "KEM" || 0x0010, the suite id of DHKEM(P-256, HKDF-SHA256).
const KEM_PREFIX = concat(utf8('HPKE-v1'), utf8('KEM'), uint16be(0x0010));

const P256_P =
  0xffffffff00000001000000000000000000000000ffffffffffffffffffffffffn;
const P256_B =
  0x5ac635d8aa3a93e7b3ebbd55769886bc651d06b0cc53b0f63bce3c3e27d2604bn;

/**
 * @typedef {Object} Encapsulation
 * @property {string} keyShare - the syncer's ephemeral public key as a did:key
 * @property {Uint8Array} sharedSecret
 */

/**
 * Syncer side: a fresh ephemeral key and the secret it shares with the
 * author's KEM key. Use a new one for every request.
 *
 * @param {string} authorKemKey - the author's `#atproto_kem` did:key
 * @returns {Promise<Encapsulation>}
 */
export async function encapsulate(authorKemKey) {
  const authorKeyBytes = uncompressedKey(authorKemKey);
  const ephemeral = /** @type {CryptoKeyPair} */ (
    await crypto.subtle.generateKey(
      { name: 'ECDH', namedCurve: 'P-256' },
      false,
      ['deriveBits'],
    )
  );
  const syncerKeyBytes = new Uint8Array(
    await crypto.subtle.exportKey('raw', ephemeral.publicKey),
  );
  const dh = await ecdh(ephemeral.privateKey, authorKeyBytes);
  return {
    keyShare: publicKeyToDidKey(compressPublicKey(syncerKeyBytes), 'p256'),
    sharedSecret: await deriveSharedSecret(dh, syncerKeyBytes, authorKeyBytes),
  };
}

/**
 * Repo host side: the same secret, from the syncer's key share and the
 * author's KEM private key.
 *
 * @param {Uint8Array} authorPrivateKey - 32-byte P-256 scalar
 * @param {string} keyShare - the syncer's did:key
 * @returns {Promise<Uint8Array>}
 * @throws {SpaceMacError} when the key share is not a P-256 did:key
 */
export async function decapsulate(authorPrivateKey, keyShare) {
  const syncerKeyBytes = uncompressedKey(keyShare);
  const privateKey = await importEcdhPrivateKey(authorPrivateKey);
  const jwk = await crypto.subtle.exportKey('jwk', privateKey);
  const authorKeyBytes = concat(
    new Uint8Array([0x04]),
    base64UrlToBytes(/** @type {string} */ (jwk.x)),
    base64UrlToBytes(/** @type {string} */ (jwk.y)),
  );
  const dh = await ecdh(privateKey, syncerKeyBytes);
  return deriveSharedSecret(dh, syncerKeyBytes, authorKeyBytes);
}

/**
 * @param {Uint8Array} dh
 * @param {Uint8Array} syncerKeyBytes - uncompressed
 * @param {Uint8Array} authorKeyBytes - uncompressed
 * @returns {Promise<Uint8Array>}
 */
function deriveSharedSecret(dh, syncerKeyBytes, authorKeyBytes) {
  const ikm = concat(KEM_PREFIX, utf8('eae_prk'), dh);
  const info = concat(
    uint16be(32),
    KEM_PREFIX,
    utf8('shared_secret'),
    syncerKeyBytes,
    authorKeyBytes,
  );
  return hkdfSha256(ikm, new Uint8Array(0), info, 32);
}

/**
 * The x-coordinate of the shared point.
 * @param {CryptoKey} privateKey
 * @param {Uint8Array} publicKeyBytes - uncompressed
 * @returns {Promise<Uint8Array>}
 */
async function ecdh(privateKey, publicKeyBytes) {
  /** @type {CryptoKey} */
  let publicKey;
  try {
    publicKey = await crypto.subtle.importKey(
      'raw',
      /** @type {BufferSource} */ (publicKeyBytes),
      { name: 'ECDH', namedCurve: 'P-256' },
      true,
      [],
    );
  } catch {
    throw new SpaceMacError('invalid KEM key');
  }
  return new Uint8Array(
    await crypto.subtle.deriveBits(
      { name: 'ECDH', public: publicKey },
      privateKey,
      256,
    ),
  );
}

/**
 * @param {Uint8Array} scalar - 32 bytes
 * @returns {Promise<CryptoKey>}
 */
function importEcdhPrivateKey(scalar) {
  if (scalar.length !== 32) throw new Error('KEM key must be 32 bytes');
  const pkcs8 = concat(
    new Uint8Array([
      0x30, 0x41, 0x02, 0x01, 0x00, 0x30, 0x13, 0x06, 0x07, 0x2a, 0x86, 0x48,
      0xce, 0x3d, 0x02, 0x01, 0x06, 0x08, 0x2a, 0x86, 0x48, 0xce, 0x3d, 0x03,
      0x01, 0x07, 0x04, 0x27, 0x30, 0x25, 0x02, 0x01, 0x01, 0x04, 0x20,
    ]),
    scalar,
  );
  return crypto.subtle.importKey(
    'pkcs8',
    /** @type {BufferSource} */ (pkcs8),
    { name: 'ECDH', namedCurve: 'P-256' },
    true,
    ['deriveBits'],
  );
}

/**
 * The uncompressed point (0x04 || x || y) a P-256 did:key names.
 * @param {string} didKey
 * @returns {Uint8Array}
 * @throws {SpaceMacError}
 */
export function uncompressedKey(didKey) {
  /** @type {{curve: string, publicKey: Uint8Array}} */
  let parsed;
  try {
    parsed = parseDidKey(didKey);
  } catch {
    throw new SpaceMacError('invalid KEM key');
  }
  if (parsed.curve !== 'p256') {
    throw new SpaceMacError('KEM key must be a P-256 did:key');
  }
  const key = parsed.publicKey;
  if (key.length !== 33 || (key[0] !== 0x02 && key[0] !== 0x03)) {
    throw new SpaceMacError('invalid KEM key');
  }
  const x = bytesToBigInt(key.subarray(1));
  if (x >= P256_P) throw new SpaceMacError('invalid KEM key');
  // y² = x³ - 3x + b, and p ≡ 3 (mod 4), so y = rhs^((p + 1) / 4).
  const rhs = mod(x ** 3n - 3n * x + P256_B);
  let y = modPow(rhs, (P256_P + 1n) / 4n);
  if (mod(y * y) !== rhs) throw new SpaceMacError('invalid KEM key');
  if (Number(y & 1n) !== (key[0] & 1)) y = P256_P - y;
  return concat(new Uint8Array([0x04]), key.subarray(1), bigIntToBytes(y));
}

/** @param {bigint} n */
function mod(n) {
  const r = n % P256_P;
  return r < 0n ? r + P256_P : r;
}

/**
 * @param {bigint} base
 * @param {bigint} exp
 * @returns {bigint}
 */
function modPow(base, exp) {
  let result = 1n;
  let b = mod(base);
  let e = exp;
  while (e > 0n) {
    if (e & 1n) result = mod(result * b);
    b = mod(b * b);
    e >>= 1n;
  }
  return result;
}

/** @param {Uint8Array} bytes */
function bytesToBigInt(bytes) {
  let n = 0n;
  for (const b of bytes) n = (n << 8n) | BigInt(b);
  return n;
}

/** @param {bigint} n */
function bigIntToBytes(n) {
  const out = new Uint8Array(32);
  let v = n;
  for (let i = 31; i >= 0; i--) {
    out[i] = Number(v & 0xffn);
    v >>= 8n;
  }
  return out;
}

/** @param {string} b64 */
function base64UrlToBytes(b64) {
  const std = b64.replace(/-/g, '+').replace(/_/g, '/');
  return Uint8Array.from(
    atob(std.padEnd(Math.ceil(std.length / 4) * 4, '=')),
    (c) => c.charCodeAt(0),
  );
}