// @pdsjs/spaces/blake3 - vendored BLAKE3 with extendable output.
//
// This is the only hand-written cryptographic primitive in pds.js. WebCrypto
// has no BLAKE3, and LtHash (proposal 0016) is defined in terms of it. Verified
// against the official BLAKE3 test vectors in test/fixtures/blake3-vectors.json;
// do not modify without re-running packages/spaces/test/blake3.test.js.
//
// Transliterated from the BLAKE3 reference implementation
// (https://github.com/BLAKE3-team/BLAKE3, reference_impl/reference_impl.rs),
// per the spec at https://github.com/BLAKE3-team/BLAKE3-specs. Only the
// unkeyed, extendable-output mode is implemented — no keyed hash, no
// derive-key mode.

const IV = new Uint32Array([
  0x6a09e667, 0xbb67ae85, 0x3c6ef372, 0xa54ff53a, 0x510e527f, 0x9b05688c,
  0x1f83d9ab, 0x5be0cd19,
]);

const MSG_PERMUTATION = [2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8];

const CHUNK_START = 1 << 0;
const CHUNK_END = 1 << 1;
const PARENT = 1 << 2;
const ROOT = 1 << 3;

const BLOCK_LEN = 64;
const CHUNK_LEN = 1024;

/**
 * A deferred compression: the final block of a chunk or parent, not yet
 * compressed, so the ROOT flag can still be applied to the last one.
 * @typedef {Object} Output
 * @property {Uint32Array} cv
 * @property {Uint32Array} block
 * @property {number} counter
 * @property {number} blockLen
 * @property {number} flags
 */

/**
 * @param {number} x
 * @param {number} n
 * @returns {number}
 */
const rotr = (x, n) => ((x >>> n) | (x << (32 - n))) >>> 0;

/**
 * The quarter-round. Rotations are 16, 12, 8, 7.
 * @param {Uint32Array} s - 16-word compression state, mutated in place
 * @param {number} a
 * @param {number} b
 * @param {number} c
 * @param {number} d
 * @param {number} mx
 * @param {number} my
 * @returns {void}
 */
function g(s, a, b, c, d, mx, my) {
  s[a] = (s[a] + s[b] + mx) >>> 0;
  s[d] = rotr(s[d] ^ s[a], 16);
  s[c] = (s[c] + s[d]) >>> 0;
  s[b] = rotr(s[b] ^ s[c], 12);
  s[a] = (s[a] + s[b] + my) >>> 0;
  s[d] = rotr(s[d] ^ s[a], 8);
  s[c] = (s[c] + s[d]) >>> 0;
  s[b] = rotr(s[b] ^ s[c], 7);
}

/**
 * Four column mixes then four diagonal mixes.
 * @param {Uint32Array} s - 16-word compression state, mutated in place
 * @param {Uint32Array} m - 16 message words
 * @returns {void}
 */
function round(s, m) {
  g(s, 0, 4, 8, 12, m[0], m[1]);
  g(s, 1, 5, 9, 13, m[2], m[3]);
  g(s, 2, 6, 10, 14, m[4], m[5]);
  g(s, 3, 7, 11, 15, m[6], m[7]);
  g(s, 0, 5, 10, 15, m[8], m[9]);
  g(s, 1, 6, 11, 12, m[10], m[11]);
  g(s, 2, 7, 8, 13, m[12], m[13]);
  g(s, 3, 4, 9, 14, m[14], m[15]);
}

/**
 * @param {Uint32Array} m - 16 message words
 * @returns {Uint32Array} the message words reordered by MSG_PERMUTATION
 */
function permute(m) {
  const out = new Uint32Array(16);
  for (let i = 0; i < 16; i++) out[i] = m[MSG_PERMUTATION[i]];
  return out;
}

/**
 * @param {Uint32Array} cv - 8-word chaining value
 * @param {Uint32Array} block - 16 message words
 * @param {number} counter
 * @param {number} blockLen
 * @param {number} flags
 * @returns {Uint32Array} all 16 state words — the first 8 are the chaining
 * value, all 16 are needed for extended output.
 */
function compress(cv, block, counter, blockLen, flags) {
  const s = new Uint32Array(16);
  s.set(cv, 0);
  s[8] = IV[0];
  s[9] = IV[1];
  s[10] = IV[2];
  s[11] = IV[3];
  s[12] = counter >>> 0;
  s[13] = Math.floor(counter / 2 ** 32) >>> 0;
  s[14] = blockLen;
  s[15] = flags;

  let m = block;
  for (let i = 0; i < 7; i++) {
    round(s, m);
    if (i < 6) m = permute(m);
  }
  for (let i = 0; i < 8; i++) {
    s[i] = (s[i] ^ s[i + 8]) >>> 0;
    s[i + 8] = (s[i + 8] ^ cv[i]) >>> 0;
  }
  return s;
}

/**
 * Read 16 little-endian uint32s from a 64-byte block.
 * @param {Uint8Array} bytes
 * @param {number} offset
 * @returns {Uint32Array}
 */
function wordsFromBlock(bytes, offset) {
  const m = new Uint32Array(16);
  for (let i = 0; i < 16; i++) {
    const o = offset + i * 4;
    m[i] =
      (bytes[o] |
        (bytes[o + 1] << 8) |
        (bytes[o + 2] << 16) |
        (bytes[o + 3] << 24)) >>>
      0;
  }
  return m;
}

/**
 * @param {Output} o
 * @returns {Uint32Array} the 8-word chaining value
 */
function outputChainingValue(o) {
  return compress(o.cv, o.block, o.counter, o.blockLen, o.flags).slice(0, 8);
}

/**
 * @param {Output} o
 * @param {number} outLen - output length in bytes
 * @returns {Uint8Array}
 */
function rootOutputBytes(o, outLen) {
  const out = new Uint8Array(outLen);
  let i = 0;
  let counter = 0;
  while (i < outLen) {
    const words = compress(o.cv, o.block, counter, o.blockLen, o.flags | ROOT);
    for (let w = 0; w < 16 && i < outLen; w++) {
      for (let b = 0; b < 4 && i < outLen; b++) {
        out[i++] = (words[w] >>> (8 * b)) & 0xff;
      }
    }
    counter++;
  }
  return out;
}

/**
 * @param {Uint8Array} input
 * @param {number} start - byte offset of the chunk within input
 * @param {number} len - chunk length in bytes
 * @param {number} chunkCounter
 * @returns {Output}
 */
function chunkOutput(input, start, len, chunkCounter) {
  let cv = IV.slice(0, 8);
  let blockFlags = CHUNK_START;
  let pos = 0;
  while (len - pos > BLOCK_LEN) {
    const m = wordsFromBlock(input, start + pos);
    cv = compress(cv, m, chunkCounter, BLOCK_LEN, blockFlags).slice(0, 8);
    pos += BLOCK_LEN;
    blockFlags = 0;
  }
  const blockLen = len - pos;
  const padded = new Uint8Array(BLOCK_LEN);
  padded.set(input.subarray(start + pos, start + len));
  return {
    cv,
    block: wordsFromBlock(padded, 0),
    counter: chunkCounter,
    blockLen,
    flags: blockFlags | CHUNK_END,
  };
}

/**
 * @param {Uint32Array} leftCv - 8-word chaining value
 * @param {Uint32Array} rightCv - 8-word chaining value
 * @returns {Output}
 */
function parentOutput(leftCv, rightCv) {
  const block = new Uint32Array(16);
  block.set(leftCv, 0);
  block.set(rightCv, 8);
  return {
    cv: IV.slice(0, 8),
    block,
    counter: 0,
    blockLen: BLOCK_LEN,
    flags: PARENT,
  };
}

/**
 * Unkeyed BLAKE3 with extendable output.
 * @param {Uint8Array} input
 * @param {number} [outLen] - output length in bytes, default 32
 * @returns {Uint8Array}
 */
export function blake3(input, outLen = 32) {
  const chunks = Math.max(1, Math.ceil(input.length / CHUNK_LEN));
  /** @type {Uint32Array[]} */
  const stack = [];
  /** @type {Output|undefined} */
  let output;

  for (let i = 0; i < chunks; i++) {
    const start = i * CHUNK_LEN;
    const len = Math.min(CHUNK_LEN, Math.max(0, input.length - start));
    const chunk = chunkOutput(input, start, len, i);

    // The last chunk stays an Output so the ROOT flag can reach it.
    if (i === chunks - 1) {
      output = chunk;
      break;
    }

    // Merge while the completed-chunk count has a trailing zero bit.
    let cv = outputChainingValue(chunk);
    let n = i + 1;
    while ((n & 1) === 0) {
      // Guaranteed non-empty: n's trailing zero count never exceeds the
      // number of merges already pushed onto the stack.
      cv = outputChainingValue(
        parentOutput(/** @type {Uint32Array} */ (stack.pop()), cv),
      );
      n >>= 1;
    }
    stack.push(cv);
  }

  // At least one chunk always runs, so `output` is set by the time either
  // loop below reads it.
  while (stack.length > 0) {
    output = parentOutput(
      /** @type {Uint32Array} */ (stack.pop()),
      outputChainingValue(/** @type {Output} */ (output)),
    );
  }
  return rootOutputBytes(/** @type {Output} */ (output), outLen);
}