/**
 * Reed-Solomon codec — GF(2^8) with primitive polynomial 0x11D
 * (x^8 + x^4 + x^3 + x^2 + 1, the standard CCSDS / NASA polynomial).
 *
 * Two API surfaces:
 *
 * 1. Per-block primitives (Mayan v2): rsEncodeBlock / rsDecodeBlock with
 *    arbitrary parity count. Used for RS(255,204) — 204 data + 51 parity —
 *    where the Mayan layer does its own blocking and interleaving.
 *
 * 2. Legacy block-mode RS(255,178): reedSolomonEncode / reedSolomonDecode
 *    auto-block a payload into 178/77 blocks. Used by Mayan v1 codes only.
 *    New v2 codes do not go through this path.
 */

const FIELD_SIZE = 256;
const PRIMITIVE_POLY = 0x11d;

const BLOCK_DATA = 178;
const BLOCK_PARITY = 77;
const BLOCK_TOTAL = BLOCK_DATA + BLOCK_PARITY; // 255

// ── GF(2^8) tables ──────────────────────────────────────────────────────────

const EXP_TABLE = new Uint8Array(FIELD_SIZE * 2);
const LOG_TABLE = new Uint8Array(FIELD_SIZE);

(function initTables() {
  let x = 1;
  for (let i = 0; i < 255; i++) {
    EXP_TABLE[i] = x;
    LOG_TABLE[x] = i;
    x <<= 1;
    if (x & 0x100) x ^= PRIMITIVE_POLY;
  }
  // Double-length exp table avoids modulo on the hot path.
  for (let i = 255; i < 512; i++) EXP_TABLE[i] = EXP_TABLE[i - 255];
})();

function gfMul(a: number, b: number): number {
  if (a === 0 || b === 0) return 0;
  return EXP_TABLE[LOG_TABLE[a] + LOG_TABLE[b]];
}

function gfDiv(a: number, b: number): number {
  if (b === 0) throw new Error("rs: divide by zero");
  if (a === 0) return 0;
  return EXP_TABLE[(LOG_TABLE[a] + 255 - LOG_TABLE[b]) % 255];
}

function gfPow(x: number, power: number): number {
  if (x === 0) return 0;
  return EXP_TABLE[(LOG_TABLE[x] * power) % 255];
}

function gfInv(x: number): number {
  if (x === 0) throw new Error("rs: inverse of zero");
  return EXP_TABLE[(255 - LOG_TABLE[x]) % 255];
}

// ── Polynomial helpers (coefficients high→low; out[0] = x^(n-1)) ────────────

function polyScale(p: Uint8Array, x: number): Uint8Array {
  const r = new Uint8Array(p.length);
  for (let i = 0; i < p.length; i++) r[i] = gfMul(p[i], x);
  return r;
}

function polyAdd(p: Uint8Array, q: Uint8Array): Uint8Array {
  const len = Math.max(p.length, q.length);
  const r = new Uint8Array(len);
  for (let i = 0; i < p.length; i++) r[i + len - p.length] = p[i];
  for (let i = 0; i < q.length; i++) r[i + len - q.length] ^= q[i];
  return r;
}

function polyMul(p: Uint8Array, q: Uint8Array): Uint8Array {
  const r = new Uint8Array(p.length + q.length - 1);
  for (let j = 0; j < q.length; j++) {
    if (q[j] === 0) continue;
    for (let i = 0; i < p.length; i++) {
      if (p[i] === 0) continue;
      r[i + j] ^= gfMul(p[i], q[j]);
    }
  }
  return r;
}

function polyEval(p: Uint8Array, x: number): number {
  let y = p[0];
  for (let i = 1; i < p.length; i++) y = gfMul(y, x) ^ p[i];
  return y;
}

function generatorPoly(parity: number): Uint8Array {
  let g: Uint8Array = new Uint8Array([1]);
  for (let i = 0; i < parity; i++) {
    g = polyMul(g, new Uint8Array([1, EXP_TABLE[i]]));
  }
  return g;
}

// Cache generators per parity length — only one is used in practice.
const generatorCache = new Map<number, Uint8Array>();
function getGenerator(parity: number): Uint8Array {
  let g = generatorCache.get(parity);
  if (!g) {
    g = generatorPoly(parity);
    generatorCache.set(parity, g);
  }
  return g;
}

// ── Block encode/decode ─────────────────────────────────────────────────────

function encodeBlock(data: Uint8Array, parity: number): Uint8Array {
  const gen = getGenerator(parity);
  const out = new Uint8Array(data.length + parity);
  out.set(data);
  for (let i = 0; i < data.length; i++) {
    const coef = out[i];
    if (coef !== 0) {
      for (let j = 1; j < gen.length; j++) {
        out[i + j] ^= gfMul(gen[j], coef);
      }
    }
  }
  // First `data.length` bytes still hold the original message; replace with
  // it (it was used as scratch above) and copy the parity tail unchanged.
  out.set(data);
  return out;
}

function calcSyndromes(received: Uint8Array, parity: number): Uint8Array {
  const synd = new Uint8Array(parity + 1);
  synd[0] = 0;
  for (let i = 0; i < parity; i++) {
    synd[i + 1] = polyEval(received, EXP_TABLE[i]);
  }
  return synd;
}

function findErrorLocator(synd: Uint8Array, parity: number): Uint8Array | null {
  let errLoc: Uint8Array = new Uint8Array([1]);
  let oldLoc: Uint8Array = new Uint8Array([1]);

  for (let i = 0; i < parity; i++) {
    oldLoc = new Uint8Array([...oldLoc, 0]);
    let delta = synd[i + 1];
    for (let j = 1; j < errLoc.length; j++) {
      delta ^= gfMul(errLoc[errLoc.length - 1 - j], synd[i + 1 - j]);
    }
    if (delta !== 0) {
      if (oldLoc.length > errLoc.length) {
        const newLoc = polyScale(oldLoc, delta);
        oldLoc = polyScale(errLoc, gfInv(delta));
        errLoc = newLoc;
      }
      errLoc = polyAdd(errLoc, polyScale(oldLoc, delta));
    }
  }

  // Strip leading zeros
  let lead = 0;
  while (lead < errLoc.length && errLoc[lead] === 0) lead++;
  errLoc = errLoc.slice(lead);

  const errs = errLoc.length - 1;
  if (errs * 2 > parity) return null; // too many errors
  return errLoc;
}

function findErrors(errLoc: Uint8Array, length: number): number[] | null {
  // Λ has roots at α^(-poly_degree). Chien finds Λ(α^j)=0 at j = (b+1) mod 255
  // for byte index b in a length-255 codeword. So b = (j - 1) mod 255.
  // For our use this is always full RS(255,178) blocks; the modulus is 255.
  const errs = errLoc.length - 1;
  const positions: number[] = [];
  const N = 255;
  for (let i = 0; i < N; i++) {
    if (polyEval(errLoc, gfPow(2, i)) === 0) {
      const byteIdx = (i - 1 + N) % N;
      if (byteIdx < length) positions.push(byteIdx);
    }
  }
  if (positions.length !== errs) return null;
  return positions;
}

function correctErrors(
  received: Uint8Array,
  synd: Uint8Array,
  errPositions: number[],
): Uint8Array {
  const coefPos = errPositions.map((p) => received.length - 1 - p);
  // Forney syndrome
  const errLocPoly = (() => {
    let p: Uint8Array = new Uint8Array([1]);
    for (const pos of coefPos) {
      p = polyMul(p, polyAdd(new Uint8Array([1]), new Uint8Array([gfPow(2, pos), 0])));
    }
    return p;
  })();

  const syndRev = new Uint8Array(synd.length - 1);
  for (let i = 0; i < syndRev.length; i++) syndRev[i] = synd[synd.length - 1 - i];
  let omega: Uint8Array = polyMul(syndRev, errLocPoly);
  omega = omega.slice(omega.length - errLocPoly.length + 1);

  // Formal derivative of errLocPoly (skip even-degree terms)
  const errLocDeriv = new Uint8Array(errLocPoly.length - 1);
  for (let i = 0; i < errLocDeriv.length; i++) {
    const deg = errLocPoly.length - 1 - i;
    if (deg % 2 === 1) errLocDeriv[i] = errLocPoly[i];
  }

  const out = new Uint8Array(received);
  for (let i = 0; i < coefPos.length; i++) {
    const pos = coefPos[i];
    const xi = gfPow(2, pos);
    const xiInv = gfInv(xi);
    const num = polyEval(omega, xiInv);
    const den = polyEval(errLocDeriv, xiInv);
    if (den === 0) throw new Error("rs: denominator zero in Forney");
    const magnitude = gfMul(xi, gfDiv(num, den));
    out[received.length - 1 - pos] ^= magnitude;
  }
  return out;
}

function decodeBlock(received: Uint8Array, parity: number): Uint8Array | null {
  const synd = calcSyndromes(received, parity);
  let hasError = false;
  for (let i = 1; i < synd.length; i++) {
    if (synd[i] !== 0) { hasError = true; break; }
  }
  if (!hasError) return received.slice(0, received.length - parity);

  const errLoc = findErrorLocator(synd, parity);
  if (!errLoc) return null;

  const positions = findErrors(errLoc, received.length);
  if (!positions) return null;

  let corrected: Uint8Array;
  try {
    corrected = correctErrors(received, synd, positions);
  } catch {
    return null;
  }

  // Verify by recomputing syndromes
  const synd2 = calcSyndromes(corrected, parity);
  for (let i = 1; i < synd2.length; i++) {
    if (synd2[i] !== 0) return null;
  }
  return corrected.slice(0, received.length - parity);
}

// ── Public block-mode API ───────────────────────────────────────────────────

/**
 * Encode `data` with Reed-Solomon. Splits into BLOCK_DATA-byte data blocks
 * (zero-padded), appends BLOCK_PARITY parity bytes per block, returns the
 * concatenation. `_advisoryParityBytes` is accepted for spec compatibility
 * and ignored — the actual parity volume is fixed by the block geometry.
 */
export function reedSolomonEncode(
  data: Uint8Array,
  _advisoryParityBytes?: number,
): Uint8Array {
  if (data.length === 0) return new Uint8Array(0);
  const blocks = Math.ceil(data.length / BLOCK_DATA);
  const out = new Uint8Array(blocks * BLOCK_TOTAL);
  for (let b = 0; b < blocks; b++) {
    const start = b * BLOCK_DATA;
    const slice = data.slice(start, start + BLOCK_DATA);
    const padded = slice.length === BLOCK_DATA
      ? slice
      : (() => { const p = new Uint8Array(BLOCK_DATA); p.set(slice); return p; })();
    const enc = encodeBlock(padded, BLOCK_PARITY);
    out.set(enc, b * BLOCK_TOTAL);
  }
  return out;
}

/**
 * Decode an RS-encoded stream. Splits into BLOCK_TOTAL-byte received blocks
 * and decodes each. `dataBytes` is the original payload length; the result
 * is trimmed accordingly. Throws if any block fails (too many errors).
 *
 * `_parityBytes` is accepted for spec compatibility; the per-block parity
 * count is fixed by the block geometry.
 */
export function reedSolomonDecode(
  received: Uint8Array,
  dataBytes: number,
  _parityBytes?: number,
): Uint8Array {
  if (dataBytes === 0) return new Uint8Array(0);
  if (received.length % BLOCK_TOTAL !== 0) {
    throw new Error(
      `rs: received length ${received.length} not a multiple of ${BLOCK_TOTAL}`,
    );
  }
  const blocks = received.length / BLOCK_TOTAL;
  const out = new Uint8Array(blocks * BLOCK_DATA);
  for (let b = 0; b < blocks; b++) {
    const slice = received.slice(b * BLOCK_TOTAL, (b + 1) * BLOCK_TOTAL);
    const decoded = decodeBlock(slice, BLOCK_PARITY);
    if (!decoded) throw new Error(`rs: block ${b} uncorrectable`);
    out.set(decoded, b * BLOCK_DATA);
  }
  return out.slice(0, dataBytes);
}

/** Constants the Mayan layer needs for capacity math. */
export const RS_BLOCK_DATA = BLOCK_DATA;
export const RS_BLOCK_PARITY = BLOCK_PARITY;
export const RS_BLOCK_TOTAL = BLOCK_TOTAL;

// ── Per-block API (Mayan v2 uses these directly with parity=51) ─────────────

/**
 * Encode a single shortened RS codeword: append `parity` parity bytes to
 * `data`, returning a `data.length + parity`-byte block. Caller controls
 * the block geometry; for RS(255,204) pass 204-byte data + parity=51.
 */
export function rsEncodeBlock(data: Uint8Array, parity: number): Uint8Array {
  return encodeBlock(data, parity);
}

/**
 * Decode a single shortened RS codeword. Returns the data portion (without
 * parity) on success, or null if uncorrectable. Block length and parity must
 * match what was used at encode time.
 */
export function rsDecodeBlock(received: Uint8Array, parity: number): Uint8Array | null {
  return decodeBlock(received, parity);
}