import { blake2b } from '@noble/hashes/blake2.js';

// AEZ v5 — the wide-block cipher LND's aezeed uses to encrypt a cipher seed.
//
// SCOPE, deliberately narrow. This is a DECRYPT-ONLY implementation of the SHORT-INPUT path:
//
//   - No encrypt. Officer never mints aezeed seeds — it only imports ones LND/Zeus already created.
//     A new Officer wallet gets a BIP39 mnemonic, which is the portable format.
//   - No aezCore. An aezeed ciphertext is a FIXED 23 bytes (33 total, minus the 1-byte version, the
//     5-byte salt and the 4-byte CRC), and AEZ dispatches on length: anything under 32 bytes takes the
//     `aezTiny` Feistel path. `aezCore` is unreachable here, so shipping it would mean carrying ~120
//     lines of untested cipher in a seed-recovery path. `decrypt` throws rather than guess.
//   - No aezPRF. Only reached when the ciphertext is exactly `tau` bytes (empty plaintext).
//
// Every line below is exercised by the tests. That is the point of the narrowness — a subtle bug in
// unreachable crypto is a bug nobody finds until it eats someone's wallet.
//
// Ported from the reference implementation (Yawning/aez, as used by lnd/aezeed). Faithful to the point
// of preserving its integer-truncation quirks; where the original relied on a JavaScript accident, the
// comment says so.

const BLOCK_SIZE = 16;
const EXTRACTED_KEY_SIZE = 48;
const ZERO = new Uint8Array(BLOCK_SIZE);

// ── AES round tables ─────────────────────────────────────────────────────────────────────────────
//
// AEZ is built on reduced-round AES (4 and 10 rounds, always with MixColumns and no key schedule —
// the "keys" are just the extracted key material). The four 256-entry T-tables are pure functions of
// the AES S-box: TE0[i] = [2·S[i], S[i], S[i], 3·S[i]] big-endian, and TE1..TE3 are TE0 rotated right
// by 8, 16 and 24 bits. Deriving them costs microseconds once and removes 1024 hand-copied constants
// from a file where a single wrong digit silently produces the wrong wallet.

// prettier-ignore
const SBOX_HEX =
  '637c777bf26b6fc53001672bfed7ab76ca82c97dfa5947f0add4a2af9ca472c0' +
  'b7fd9326363ff7cc34a5e5f171d8311504c723c31896059a071280e2eb27b275' +
  '09832c1a1b6e5aa0523bd6b329e32f8453d100ed20fcb15b6acbbe394a4c58cf' +
  'd0efaafb434d338545f9027f503c9fa851a3408f929d38f5bcb6da2110fff3d2' +
  'cd0c13ec5f974417c4a77e3d645d197360814fdc222a908846eeb814de5e0bdb' +
  'e0323a0a4906245cc2d3ac629195e479e7c8376d8dd54ea96c56f4ea657aae08' +
  'ba78252e1ca6b4c6e8dd741f4bbd8b8a703eb5664803f60e613557b986c11d9e' +
  'e1f8981169d98e949b1e87e9ce5528df8ca1890dbfe6426841992d0fb054bb16';

const SBOX = Uint8Array.from(SBOX_HEX.match(/../g)!, (h) => parseInt(h, 16));

/** GF(2^8) doubling — the `xtime` of the AES spec. */
const xtime = (b: number): number => ((b << 1) ^ (b & 0x80 ? 0x1b : 0)) & 0xff;

function buildTables(): [Uint32Array, Uint32Array, Uint32Array, Uint32Array] {
  const te0 = new Uint32Array(256);
  for (let i = 0; i < 256; i++) {
    const s = SBOX[i]!;
    te0[i] = ((xtime(s) << 24) | (s << 16) | (s << 8) | (xtime(s) ^ s)) >>> 0;
  }
  const ror8 = (x: number): number => ((x >>> 8) | (x << 24)) >>> 0;
  const te1 = te0.map(ror8);
  const te2 = te1.map(ror8);
  const te3 = te2.map(ror8);
  return [te0, te1, te2, te3];
}

const [TE0, TE1, TE2, TE3] = buildTables();

// ── block helpers ────────────────────────────────────────────────────────────────────────────────

const block = (size = BLOCK_SIZE): Uint8Array => new Uint8Array(size);

function xor16(a: Uint8Array, b: Uint8Array, dst: Uint8Array): void {
  for (let i = 0; i < BLOCK_SIZE; i++) dst[i] = a[i]! ^ b[i]!;
}

function xor4x16(a: Uint8Array, b: Uint8Array, c: Uint8Array, d: Uint8Array, dst: Uint8Array): void {
  for (let i = 0; i < BLOCK_SIZE; i++) dst[i] = a[i]! ^ b[i]! ^ c[i]! ^ d[i]!;
}

/** Multiply a block by x in GF(2^128), in place. */
function doubleBlock(p: Uint8Array): void {
  const carry = p[0]!;
  for (let i = 0; i < 15; i++) p[i] = ((p[i]! << 1) | (p[i + 1]! >> 7)) & 0xff;
  p[15] = ((p[15]! << 1) ^ (carry >> 7 ? 135 : 0)) & 0xff;
}

/** Multiply a block by the integer `x` in GF(2^128) — repeated doubling, accumulating odd bits. */
function multBlock(x: number, src: Uint8Array, dst: Uint8Array): void {
  const t = block();
  const r = block();
  t.set(src);
  let n = x;
  while (n !== 0) {
    if (n & 1) xor16(r, t, r);
    doubleBlock(t);
    n >>= 1;
  }
  dst.set(r);
}

const readU32BE = (b: Uint8Array, off: number): number =>
  ((b[off]! << 24) | (b[off + 1]! << 16) | (b[off + 2]! << 8) | b[off + 3]!) >>> 0;

function writeU32BE(b: Uint8Array, off: number, v: number): void {
  b[off] = (v >>> 24) & 0xff;
  b[off + 1] = (v >>> 16) & 0xff;
  b[off + 2] = (v >>> 8) & 0xff;
  b[off + 3] = v & 0xff;
}

// ── reduced-round AES ────────────────────────────────────────────────────────────────────────────

type AesKeys = { aes4: Uint32Array; aes10: Uint32Array };

/** Lay the 48-byte extracted key out into the fixed round-key schedules AEZ prescribes. */
function aesKeys(extracted: Uint8Array): AesKeys {
  const k = new Uint32Array(12);
  for (let i = 0; i < 12; i++) k[i] = readU32BE(extracted, 4 * i);

  const aes10 = new Uint32Array(40);
  aes10.set(k, 0);
  aes10.set(k, 12);
  aes10.set(k, 24);
  aes10.set(k.subarray(0, 4), 36);

  const aes4 = new Uint32Array(16);
  aes4.set(k.subarray(4, 8), 0);
  aes4.set(k.subarray(0, 4), 4);
  aes4.set(k.subarray(8, 12), 8);

  return { aes4, aes10 };
}

/** AES rounds with MixColumns applied on every round, including the last (as AEZ specifies). */
function rounds(keys: Uint32Array, count: number, b: Uint8Array): void {
  let s0 = readU32BE(b, 0);
  let s1 = readU32BE(b, 4);
  let s2 = readU32BE(b, 8);
  let s3 = readU32BE(b, 12);

  for (let r = 0; r < count; r++) {
    const o = r * 4;
    const t0 =
      TE0[(s0 >>> 24) & 0xff]! ^ TE1[(s1 >>> 16) & 0xff]! ^ TE2[(s2 >>> 8) & 0xff]! ^ TE3[s3 & 0xff]! ^ keys[o]!;
    const t1 =
      TE0[(s1 >>> 24) & 0xff]! ^ TE1[(s2 >>> 16) & 0xff]! ^ TE2[(s3 >>> 8) & 0xff]! ^ TE3[s0 & 0xff]! ^ keys[o + 1]!;
    const t2 =
      TE0[(s2 >>> 24) & 0xff]! ^ TE1[(s3 >>> 16) & 0xff]! ^ TE2[(s0 >>> 8) & 0xff]! ^ TE3[s1 & 0xff]! ^ keys[o + 2]!;
    const t3 =
      TE0[(s3 >>> 24) & 0xff]! ^ TE1[(s0 >>> 16) & 0xff]! ^ TE2[(s1 >>> 8) & 0xff]! ^ TE3[s2 & 0xff]! ^ keys[o + 3]!;
    s0 = t0 >>> 0;
    s1 = t1 >>> 0;
    s2 = t2 >>> 0;
    s3 = t3 >>> 0;
  }

  writeU32BE(b, 0, s0);
  writeU32BE(b, 4, s1);
  writeU32BE(b, 8, s2);
  writeU32BE(b, 12, s3);
}

function aes4(k: AesKeys, j: Uint8Array, i: Uint8Array, l: Uint8Array, src: Uint8Array, dst: Uint8Array): void {
  xor4x16(j, i, l, src, dst);
  rounds(k.aes4, 4, dst);
}

function aes10(k: AesKeys, l: Uint8Array, src: Uint8Array, dst: Uint8Array): void {
  xor16(src, l, dst);
  rounds(k.aes10, 10, dst);
}

// ── AEZ state ────────────────────────────────────────────────────────────────────────────────────

type AezState = { I: Uint8Array[]; J: Uint8Array[]; L: Uint8Array[]; aes: AesKeys };

/** AEZ tolerates any key length by hashing it to 48 bytes; a 48-byte key is used verbatim. */
function extractKey(key: Uint8Array): Uint8Array {
  return key.length === EXTRACTED_KEY_SIZE ? key : blake2b(key, { dkLen: EXTRACTED_KEY_SIZE });
}

function init(key: Uint8Array): AezState {
  const ext = extractKey(key);
  const I = [block(), block()];
  const J = [block(), block(), block()];
  const L = Array.from({ length: 8 }, () => block());

  I[0]!.set(ext.subarray(0, 16));
  multBlock(2, I[0]!, I[1]!);

  J[0]!.set(ext.subarray(16, 32));
  multBlock(2, J[0]!, J[1]!);
  multBlock(2, J[1]!, J[2]!);

  // L[0] is deliberately left as the zero block — the reference never assigns it.
  L[1]!.set(ext.subarray(32, 48));
  multBlock(2, L[1]!, L[2]!);
  xor16(L[2]!, L[1]!, L[3]!);
  multBlock(2, L[2]!, L[4]!);
  xor16(L[4]!, L[1]!, L[5]!);
  multBlock(2, L[3]!, L[6]!);
  xor16(L[6]!, L[1]!, L[7]!);

  return { I, J, L, aes: aesKeys(ext) };
}

/** The AXU hash binding the tag length, the nonce and each associated-data string into `delta`. */
function aezHash(st: AezState, nonce: Uint8Array | null, ad: Uint8Array[], tauBits: number): Uint8Array {
  const buf = block();
  const sum = block();
  const I = block();
  const J = block();

  writeU32BE(buf, 12, tauBits >>> 0);
  xor16(st.J[0]!, st.J[1]!, J);
  aes4(st.aes, J, st.I[1]!, st.L[1]!, buf, sum);

  const nonceEmpty = !nonce || nonce.length === 0;
  let n = nonce ?? new Uint8Array(0);
  let nBytes = n.length;
  I.set(st.I[1]!);
  for (let i = 1; nBytes >= BLOCK_SIZE; i++, nBytes -= BLOCK_SIZE) {
    aes4(st.aes, st.J[2]!, I, st.L[i % 8]!, n.subarray(0, BLOCK_SIZE), buf);
    xor16(sum, buf, sum);
    n = n.subarray(BLOCK_SIZE);
    if (i % 8 === 0) doubleBlock(I);
  }
  if (nBytes > 0 || nonceEmpty) {
    buf.fill(0);
    if (!nonceEmpty) buf.set(n.subarray(0, nBytes));
    buf[nBytes] = 0x80;
    aes4(st.aes, st.J[2]!, st.I[0]!, st.L[0]!, buf, buf);
    xor16(sum, buf, sum);
  }

  ad.forEach((entry, k) => {
    let p = entry ?? new Uint8Array(0);
    const adEmpty = p.length === 0;
    let bytes = p.length;
    I.set(st.I[1]!);
    multBlock(5 + k, st.J[0]!, J);
    for (let i = 1; bytes >= BLOCK_SIZE; i++, bytes -= BLOCK_SIZE) {
      aes4(st.aes, J, I, st.L[i % 8]!, p.subarray(0, BLOCK_SIZE), buf);
      xor16(sum, buf, sum);
      p = p.subarray(BLOCK_SIZE);
      if (i % 8 === 0) doubleBlock(I);
    }
    if (bytes > 0 || adEmpty) {
      buf.fill(0);
      if (!adEmpty) buf.set(p.subarray(0, bytes));
      buf[bytes] = 0x80;
      aes4(st.aes, J, st.I[0]!, st.L[0]!, buf, buf);
      xor16(sum, buf, sum);
    }
  });

  return sum;
}

/**
 * The short-input (< 32 byte) balanced Feistel network. `d` is 0 to encipher, 1 to decipher — the only
 * difference is which end of the round counter it starts from.
 */
function aezTiny(st: AezState, delta: Uint8Array, input: Uint8Array, d: 0 | 1, dst: Uint8Array): void {
  const inBytes = input.length;
  const buf = block(2 * BLOCK_SIZE);
  const L = block();
  const R = block();
  const tmp = block();
  let mask = 0x00;
  let pad = 0x80;

  // Integer halves. The reference indexes with `inBytes/2` and leans on truncation; spelling the two
  // halves out makes the odd-length case readable instead of accidental.
  const half = inBytes >> 1;
  const halfUp = (inBytes + 1) >> 1;

  let i = 7;
  let roundCount: number;
  if (inBytes === 1) roundCount = 24;
  else if (inBytes === 2) roundCount = 16;
  else if (inBytes < 16) roundCount = 10;
  else {
    i = 6;
    roundCount = 8;
  }

  L.set(input.subarray(0, halfUp));
  R.set(input.subarray(half, half + halfUp));

  if (inBytes & 1) {
    // Odd length: the two halves overlap by a nibble, so shift R left by 4 bits and switch to
    // nibble-granular padding.
    for (let k = 0; k < half; k++) R[k] = ((R[k]! << 4) | (R[k + 1]! >> 4)) & 0xff;
    R[half] = (R[half]! << 4) & 0xff;
    pad = 0x08;
    mask = 0xf0;
  }

  let j: number;
  let step: number;
  if (d !== 0) {
    if (inBytes < 16) {
      buf.set(input.subarray(0, BLOCK_SIZE));
      buf[0] = buf[0]! | 0x80;
      xor16(delta, buf, buf);
      aes4(st.aes, ZERO, st.I[1]!, st.L[3]!, buf, tmp);
      L[0] = L[0]! ^ (tmp[0]! & 0x80);
    }
    j = roundCount - 1;
    step = -1;
  } else {
    j = 0;
    step = 1;
  }

  for (let k = 0; k < roundCount / 2; k++, j += step * 2) {
    buf.fill(0, 0, BLOCK_SIZE);
    buf.set(R.subarray(0, halfUp));
    buf[half] = (buf[half]! & mask) | pad;
    xor16(buf, delta, buf);
    buf[15] = buf[15]! ^ (j & 0xff);
    aes4(st.aes, ZERO, st.I[1]!, st.L[i]!, buf, tmp);
    xor16(L, tmp, L);

    buf.fill(0, 0, BLOCK_SIZE);
    buf.set(L.subarray(0, halfUp));
    buf[half] = (buf[half]! & mask) | pad;
    xor16(buf, delta, buf);
    buf[15] = buf[15]! ^ ((j + step) & 0xff);
    aes4(st.aes, ZERO, st.I[1]!, st.L[i]!, buf, tmp);
    xor16(R, tmp, R);
  }

  buf.set(R.subarray(0, half), 0);
  buf.set(L.subarray(0, halfUp), half);
  if (inBytes & 1) {
    for (let k = inBytes - 1; k > half; k--) buf[k] = ((buf[k]! >> 4) | (buf[k - 1]! << 4)) & 0xff;
    buf[half] = (L[0]! >> 4) | (R[half]! & 0xf0);
  }

  dst.set(buf.subarray(0, inBytes));

  if (inBytes < 16 && d === 0) {
    buf.fill(0, inBytes, BLOCK_SIZE);
    buf[0] = buf[0]! | 0x80;
    xor16(delta, buf, buf);
    aes4(st.aes, ZERO, st.I[1]!, st.L[3]!, buf, tmp);
    dst[0] = dst[0]! ^ (tmp[0]! & 0x80);
  }
}

/**
 * AEZ decrypt-and-verify. Returns the plaintext, or `null` when the `tau`-byte authenticator does not
 * check out — which is the signal that the key (and therefore the seed passphrase) is wrong.
 *
 * Throws for inputs outside the aezeed envelope rather than silently taking an untested path; see the
 * scope note at the top of this file.
 */
export function aezDecrypt(
  key: Uint8Array,
  nonce: Uint8Array | null,
  ad: Uint8Array[],
  tau: number,
  ciphertext: Uint8Array,
): Uint8Array | null {
  if (ciphertext.length >= 32 || ciphertext.length <= tau) {
    throw new Error(`aez: unsupported ciphertext length ${ciphertext.length} (decrypt is short-input only)`);
  }

  const st = init(key);
  const delta = aezHash(st, nonce, ad, tau * 8);
  const x = new Uint8Array(ciphertext.length);
  aezTiny(st, delta, ciphertext, 1, x);

  // Constant-time-ish: fold every authenticator byte before deciding, so the comparison does not
  // short-circuit on the first mismatch.
  let sum = 0;
  for (let i = 0; i < tau; i++) sum |= x[ciphertext.length - tau + i]!;
  if (sum !== 0) return null;

  return x.slice(0, ciphertext.length - tau);
}