/**
 * FPE-FF1 (Format-preserving encryption algorithm) specified in
 * {@link https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-38G.pdf | NIST 800-38G}.
 * @module
 */
import { unsafe } from './aes.ts';
import {
  aarray,
  abytes,
  anumber,
  bytesToNumberBE,
  clean,
  numberToBytesBE,
  type Cipher,
  type TArg,
  type TRet,
} from './utils.ts';

// NIST SP 800-38G §4.3 / §5.1 Algorithm 7: FF1's designated CIPH_K here is AES, so this file
// reuses the reviewed AES key schedule and single-block encryption helpers.
// NOTE: no point in inlining encrypt instead of encryptBlock, since BigInt stuff will be slow
const { expandKeyLE, encryptBlock } = unsafe;

// Format-preserving encryption algorithm (FPE-FF1) specified in
// {@link https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-38G.pdf | NIST Special Publication 800-38G}.

const BLOCK_LEN = 16;

// FF1 step 6vi needs mathematical modulo in [0, b); JS `%` is remainder, so decrypt rounds must
// map negative intermediate values back into that range.
function mod(a: number, b: number): number;
function mod(a: bigint, b: bigint): bigint;
function mod(a: any, b: any): number | bigint {
  const result = a % b;
  return result >= 0 ? result : b + result;
}

function NUMradix(radix: number, data: number[]): bigint {
  let res = BigInt(0);
  for (let i of data) res = res * BigInt(radix) + BigInt(i);
  return res;
}

type IntInterval = { min: bigint; max: bigint; shift: number };

function bigintBits(n: bigint): number {
  return n.toString(2).length;
}

// Keep an outward-rounded interval while discarding low bits. This bounds large powers without
// asking runtimes to materialize them (Bun, for example, has a lower maximum BigInt size).
function trimInterval({ min, max, shift }: IntInterval, precision: number): IntInterval {
  const drop = Math.max(0, bigintBits(max) - precision);
  if (!drop) return { min, max, shift };
  const bits = BigInt(drop);
  const mask = (BigInt(1) << bits) - BigInt(1);
  return { min: min >> bits, max: (max + mask) >> bits, shift: shift + drop };
}

function mulInterval(a: IntInterval, b: IntInterval, precision: number): IntInterval {
  return trimInterval(
    { min: a.min * b.min, max: a.max * b.max, shift: a.shift + b.shift },
    precision
  );
}

function powInterval(radix: number, exp: number, precision: number): IntInterval {
  let res = { min: BigInt(1), max: BigInt(1), shift: 0 };
  let base = { min: BigInt(radix), max: BigInt(radix), shift: 0 };
  while (exp) {
    if (exp % 2) res = mulInterval(res, base, precision);
    exp = Math.floor(exp / 2);
    if (exp) base = mulInterval(base, base, precision);
  }
  return res;
}

// ceil(ceil(v * log2(radix)) / 8), calculated with exact integer bounds as required by the
// current SP 800-38G revision. radix**v - 1 is the largest value encoded in `b` bytes.
function getFF1RadixBytes(radix: number, v: number): number {
  if (!v) return 0;
  // radix**v is a power of two only when radix is, and subtracting one then changes its bit size.
  if ((radix & (radix - 1)) === 0) {
    const bits = bigintBits(BigInt(radix)) - 1;
    return Math.ceil((bits * v) / 8);
  }
  for (let precision = 64; ; precision *= 2) {
    const { min, max, shift } = powInterval(radix, v, precision);
    const minBytes = Math.ceil((bigintBits(min) + shift) / 8);
    const maxBytes = Math.ceil((bigintBits(max) + shift) / 8);
    if (minBytes === maxBytes) return minBytes;
  }
}

// Test-only hook for sizing edge cases that would require million-digit FF1 inputs.
export const __TESTS: {
  getFF1RadixBytes: typeof getFF1RadixBytes;
} = /* @__PURE__ */ Object.freeze({ getFF1RadixBytes });

function getRound(radix: number, key: TArg<Uint8Array>, tweak: TArg<Uint8Array>, x: number[]) {
  // This implementation writes [radix]3 as 0x00 || uint16_be(radix), so radix=2^16
  // needs a real 24-bit encoder before it can be supported.
  if (radix < 2 || radix > 2 ** 16 - 1) throw new Error('invalid radix ' + radix);
  // minLen must satisfy both radix**minlen ≥ 100 and minlen ≥ 2.
  let minLen = 0;
  for (let domain = BigInt(1); domain < BigInt(100); domain *= BigInt(radix)) minLen++;
  minLen = Math.max(2, minLen);
  const maxLen = 2 ** 32 - 1;
  // 2 ≤ minlen ≤ maxlen < 2**32
  if (2 > minLen || minLen > maxLen || maxLen >= 2 ** 32)
    throw new Error('Invalid radix: 2 ≤ minlen ≤ maxlen < 2**32');
  if (!Array.isArray(x)) throw new Error('invalid X');
  if (x.length < minLen || x.length > maxLen) throw new Error('X is outside minLen..maxLen bounds');
  // SP 800-38G defines FF1 over numeral strings in base `radix`; out-of-range digits must fail
  // before NUMradix(...) or round splitting can reinterpret them as a different numeral string.
  for (const i of x) {
    if (!Number.isSafeInteger(i) || i < 0 || i >= radix)
      throw new Error('invalid X: digit outside radix');
  }
  const u = Math.floor(x.length / 2);
  const v = x.length - u;
  const b = getFF1RadixBytes(radix, v);
  const d = 4 * Math.ceil(b / 4) + 4;
  const padding = mod(-tweak.length - b - 1, 16);
  // P = [1]1 || [2]1 || [1]1 || [radix]3 || [10]1 || [u mod 256]1 || [n]4 || [t]4.
  const P = Uint8Array.from([1, 2, 1, 0, 0, 0, 10, u, 0, 0, 0, 0, 0, 0, 0, 0]);
  const view = new DataView(P.buffer);
  // The radix guard above keeps [radix]3's high byte zero, so a 16-bit write is enough.
  view.setUint16(4, radix, false);
  view.setUint32(8, x.length, false);
  view.setUint32(12, tweak.length, false);
  // Q = T || [0](−t−b−1) mod 16 || [i]1 || [NUMradix(B)]b.
  const PQ = new Uint8Array(P.length + tweak.length + padding + 1 + b);
  PQ.set(P);
  clean(P);
  PQ.set(tweak, P.length);
  const xk = expandKeyLE(key);
  const round = (A: number[], B: number[], i: number, decrypt = false) => {
    // Q = ... || [i]1 || [NUMradix(B)]b.
    PQ[PQ.length - b - 1] = i;
    if (b) PQ.set(numberToBytesBE(NUMradix(radix, B), b), PQ.length - b);
    // NIST SP 800-38G Algorithm 6 PRF: Y_j = CIPH_K(Y_(j-1) xor X_j) starting from Y_0 = 0^128.
    let r = new Uint8Array(16);
    for (let j = 0; j < PQ.length / BLOCK_LEN; j++) {
      for (let i = 0; i < BLOCK_LEN; i++) r[i] ^= PQ[j * BLOCK_LEN + i];
      encryptBlock(xk, r);
    }
    // Let S be the first d bytes of the following string of ⎡d/16⎤ blocks:
    // R || CIPHK(R ⊕[1]16) || CIPHK(R ⊕[2]16) ...CIPHK(R ⊕[⎡d / 16⎤ – 1]16).
    let s = Array.from(r);
    for (let j = 1; s.length < d; j++) {
      const block = numberToBytesBE(BigInt(j), 16);
      for (let k = 0; k < BLOCK_LEN; k++) block[k] ^= r[k];
      s.push(...Array.from(encryptBlock(xk, block)));
    }
    let y = bytesToNumberBE(Uint8Array.from(s.slice(0, d)));
    s.fill(0);
    if (decrypt) y = -y;
    const m = i % 2 === 0 ? u : v;
    let c = mod(NUMradix(radix, A) + y, BigInt(radix) ** BigInt(m));
    // STR(radix, m, c)
    const C = Array(m).fill(0);
    for (let i = 0; i < m; i++, c /= BigInt(radix)) C[m - 1 - i] = Number(c % BigInt(radix));
    A.fill(0);
    A = B;
    B = C;
    return [A, B];
  };
  const destroy = () => {
    clean(xk, PQ);
  };
  return { u, round, destroy };
}

const EMPTY_BUF = /* @__PURE__ */ Uint8Array.of();

/**
 * FPE-FF1 format-preserving encryption.
 * @param radix - Alphabet size for each input digit.
 * @param key - AES key bytes.
 * @param tweak - Optional tweak bytes.
 * @returns Encrypt/decrypt helpers over digit arrays.
 * @example
 * Encrypts decimal digits without changing their format, using a fresh AES key.
 *
 * ```ts
 * import { FF1 } from '@noble/ciphers/ff1.js';
 * import { randomBytes } from '@noble/ciphers/utils.js';
 * const key = randomBytes(16);
 * const ff1 = FF1(10, key);
 * ff1.encrypt([1, 2, 3]);
 * ```
 */
export function FF1(
  radix: number,
  key: TArg<Uint8Array>,
  tweak: TArg<Uint8Array> = EMPTY_BUF
): { encrypt(x: number[]): number[]; decrypt(x: number[]): number[] } {
  anumber(radix, 'radix');
  abytes(key, undefined, 'key');
  abytes(tweak, undefined, 'tweak');
  // This borrows caller key/tweak buffers by reference through the bound closure; mutating them
  // after construction changes later encrypt/decrypt outputs.
  const PQ = getRound.bind(null, radix, key, tweak);
  return {
    encrypt(x: number[]): number[] {
      const { u, round, destroy } = PQ(x);
      let [A, B] = [x.slice(0, u), x.slice(u)];
      for (let i = 0; i < 10; i++) [A, B] = round(A, B, i);
      destroy();
      const res = A.concat(B);
      A.fill(0);
      B.fill(0);
      return res;
    },
    decrypt(x: number[]): number[] {
      const { u, round, destroy } = PQ(x);
      // The FF1.Decrypt algorithm is similar to the FF1.Encrypt algorithm;
      // the differences are in Step 6, where:
      // 1) the order of the indices is reversed,
      // 2) the roles of A and B are swapped
      // 3) modular addition is replaced by modular subtraction, in Step 6vi.
      let [B, A] = [x.slice(0, u), x.slice(u)];
      for (let i = 9; i >= 0; i--) [A, B] = round(A, B, i, true);
      destroy();
      const res = B.concat(A);
      A.fill(0);
      B.fill(0);
      return res;
    },
  };
}
// Binary wrapper uses little-endian bit order within each byte so bit 0 stays
// in the first numeral slot for this library-defined byte-array surface.
const binLE = {
  encode(bytes: TArg<Uint8Array> | number[]): number[] {
    const x = [];
    for (let i = 0; i < bytes.length; i++) {
      for (let j = 0, tmp = bytes[i]; j < 8; j++, tmp >>= 1) x.push(tmp & 1);
    }
    return x;
  },
  decode(b: number[]): TRet<Uint8Array> {
    if (!Array.isArray(b) || b.length % 8) throw new Error('Invalid binary string');
    const res = new Uint8Array(b.length / 8);
    for (let i = 0, j = 0; i < res.length; i++) {
      res[i] = b[j++] | (b[j++] << 1) | (b[j++] << 2) | (b[j++] << 3);
      res[i] |= (b[j++] << 4) | (b[j++] << 5) | (b[j++] << 6) | (b[j++] << 7);
    }
    return res as TRet<Uint8Array>;
  },
};

/**
 * Binary FPE-FF1 wrapper over byte arrays.
 * @param key - AES key bytes.
 * @param tweak - Optional tweak bytes.
 * @returns Encrypt/decrypt helpers over byte arrays.
 * @example
 * Encrypts raw bytes through FF1's binary alphabet wrapper with a fresh AES key.
 *
 * ```ts
 * import { BinaryFF1 } from '@noble/ciphers/ff1.js';
 * import { randomBytes } from '@noble/ciphers/utils.js';
 * const key = randomBytes(16);
 * const ff1 = BinaryFF1(key);
 * ff1.encrypt(new Uint8Array([1, 2, 3]));
 * ```
 */
export function BinaryFF1(
  key: TArg<Uint8Array>,
  tweak: TArg<Uint8Array> = EMPTY_BUF
): TRet<Cipher> {
  const ff1 = FF1(2, key, tweak);
  return {
    encrypt: (x: TArg<Uint8Array> | number[]): TRet<Uint8Array> => {
      if (Array.isArray(x)) {
        aarray<number>(x, 'x', (elm, title) => {
          anumber(elm, title);
          if (elm > 255) throw new RangeError(`"${title}" expected byte`);
        });
      } else {
        x = abytes(x, undefined, 'x');
      }
      return binLE.decode(ff1.encrypt(binLE.encode(x))) as TRet<Uint8Array>;
    },
    decrypt: (x: TArg<Uint8Array> | number[]): TRet<Uint8Array> => {
      if (Array.isArray(x)) {
        aarray<number>(x, 'x', (elm, title) => {
          anumber(elm, title);
          if (elm > 255) throw new RangeError(`"${title}" expected byte`);
        });
      } else {
        x = abytes(x, undefined, 'x');
      }
      return binLE.decode(ff1.decrypt(binLE.encode(x))) as TRet<Uint8Array>;
    },
  } as TRet<Cipher>;
}