const { expect } = require("chai");
const { ethers } = require("hardhat");
const kat = require("./fixtures/xmss-sha2_10_256-kat.json");

// ---------------------------------------------------------------------------
// Independent in-repo oracle: RFC 8391 XMSS-SHA2_10_256 verify path in pure JS
// (ethers sha256), reimplemented from the spec, to cross-check the Solidity
// verifier against the OFFICIAL reference vector without trusting either side.
// ---------------------------------------------------------------------------
const N = 32, W = 16, LOGW = 4, LEN1 = 64, LEN2 = 3, LEN = 67, H = 10;
const cat = (...a) => ethers.concat(a);
const sha = (b) => ethers.sha256(b);
const tob = (x, l) => ethers.toBeHex(BigInt(x), l);
const xor = (a, b) => {
  const A = ethers.getBytes(a), B = ethers.getBytes(b), O = new Uint8Array(A.length);
  for (let i = 0; i < A.length; i++) O[i] = A[i] ^ B[i];
  return ethers.hexlify(O);
};
const addr = (typ, a4, a5, a6, km) =>
  tob((BigInt(typ) << 128n) | (BigInt(a4) << 96n) | (BigInt(a5) << 64n) | (BigInt(a6) << 32n) | BigInt(km), 32);
const prf = (seed, a) => sha(cat(tob(3, 32), seed, a));
const F = (seed, typ, a4, a5, a6, x) => {
  const key = prf(seed, addr(typ, a4, a5, a6, 0));
  const mask = prf(seed, addr(typ, a4, a5, a6, 1));
  return sha(cat(tob(0, 32), key, xor(x, mask)));
};
const Hh = (seed, typ, a4, a5, a6, l, r) => {
  const key = prf(seed, addr(typ, a4, a5, a6, 0));
  const m0 = prf(seed, addr(typ, a4, a5, a6, 1));
  const m1 = prf(seed, addr(typ, a4, a5, a6, 2));
  return sha(cat(tob(1, 32), key, xor(l, m0), xor(r, m1)));
};
function chainLengths(mHex) {
  const m = ethers.getBytes(mHex), d = new Array(LEN);
  let csum = 0;
  for (let i = 0; i < 32; i++) { const b = m[i], hi = b >> 4, lo = b & 15; d[2 * i] = hi; d[2 * i + 1] = lo; csum += (15 - hi) + (15 - lo); }
  const c = csum << 4;
  d[64] = (c >> 12) & 15; d[65] = (c >> 8) & 15; d[66] = (c >> 4) & 15;
  return d;
}
function wotsPkFromSig(seed, idxLeaf, mhash, wots) {
  const lens = chainLengths(mhash), pk = [];
  for (let i = 0; i < LEN; i++) { let x = wots[i]; for (let s = lens[i]; s < W - 1; s++) x = F(seed, 0, idxLeaf, i, s, x); pk.push(x); }
  return pk;
}
function lTree(seed, idxLeaf, pk) {
  let nodes = pk.slice(), l = LEN, height = 0;
  while (l > 1) {
    const parent = l >> 1;
    for (let i = 0; i < parent; i++) nodes[i] = Hh(seed, 1, idxLeaf, height, i, nodes[2 * i], nodes[2 * i + 1]);
    if (l & 1) { nodes[parent] = nodes[l - 1]; l = parent + 1; } else { l = parent; }
    height++;
  }
  return nodes[0];
}
function computeRoot(seed, leaf, leafIdx, auth) {
  let left, right;
  if (leafIdx & 1) { left = auth[0]; right = leaf; } else { left = leaf; right = auth[0]; }
  let li = leafIdx;
  for (let i = 0; i < H - 1; i++) {
    li >>= 1;
    const node = Hh(seed, 2, 0, i, li, left, right);
    if (li & 1) { left = auth[i + 1]; right = node; } else { left = node; right = auth[i + 1]; }
  }
  li >>= 1;
  return Hh(seed, 2, 0, H - 1, li, left, right);
}
function jsVerify(pubRoot, pubSeed, idx, R, msg, wots, auth) {
  const mhash = sha(cat(tob(2, 32), R, pubRoot, tob(idx, 32), msg));
  const idxLeaf = idx & ((1 << H) - 1);
  const wpk = wotsPkFromSig(pubSeed, idxLeaf, mhash, wots);
  const leaf = lTree(pubSeed, idxLeaf, wpk);
  return computeRoot(pubSeed, leaf, idxLeaf, auth);
}

describe("AereXmssVerifier (RFC 8391 XMSS-SHA2_10_256)", function () {
  async function deploy() {
    const V = await ethers.getContractFactory("AereXmssVerifier");
    const v = await V.deploy();
    await v.waitForDeployment();
    return v;
  }
  const V = () => [kat.pubRoot, kat.pubSeed, kat.idx, kat.R, kat.msg, kat.wots, kat.auth];

  it("independent JS oracle recomputes the official reference root", function () {
    const root = jsVerify(kat.pubRoot, kat.pubSeed, kat.idx, kat.R, kat.msg, kat.wots, kat.auth);
    expect(root).to.equal(kat.pubRoot);
  });

  it("accepts the OFFICIAL reference KAT signature (verify -> true)", async function () {
    const v = await deploy();
    expect(await v.verify(...V())).to.equal(true);
  });

  it("chainLengths matches the WOTS+ checksum spec for the KAT message hash", async function () {
    const v = await deploy();
    // mhash as computed in verify()
    const mhash = ethers.sha256(ethers.concat([ethers.toBeHex(2n, 32), kat.R, kat.pubRoot, ethers.toBeHex(BigInt(kat.idx), 32), kat.msg]));
    const onchain = (await v.chainLengths(mhash)).map((x) => Number(x));
    expect(onchain).to.deep.equal(chainLengths(mhash));
  });

  it("rejects a tampered WOTS+ signature (verify -> false)", async function () {
    const v = await deploy();
    const wots = kat.wots.slice();
    wots[0] = xor(wots[0], "0x01" + "00".repeat(31));
    expect(await v.verify(kat.pubRoot, kat.pubSeed, kat.idx, kat.R, kat.msg, wots, kat.auth)).to.equal(false);
  });

  it("rejects a tampered authentication path (verify -> false)", async function () {
    const v = await deploy();
    const auth = kat.auth.slice();
    auth[9] = xor(auth[9], "0x01" + "00".repeat(31));
    expect(await v.verify(kat.pubRoot, kat.pubSeed, kat.idx, kat.R, kat.msg, kat.wots, auth)).to.equal(false);
  });

  it("rejects a tampered message (verify -> false)", async function () {
    const v = await deploy();
    expect(await v.verify(kat.pubRoot, kat.pubSeed, kat.idx, kat.R, "0x26", kat.wots, kat.auth)).to.equal(false);
  });

  it("rejects a wrong leaf index (verify -> false)", async function () {
    const v = await deploy();
    expect(await v.verify(kat.pubRoot, kat.pubSeed, kat.idx + 1, kat.R, kat.msg, kat.wots, kat.auth)).to.equal(false);
  });

  it("rejects a wrong public root (verify -> false)", async function () {
    const v = await deploy();
    const badRoot = xor(kat.pubRoot, "0x01" + "00".repeat(31));
    expect(await v.verify(badRoot, kat.pubSeed, kat.idx, kat.R, kat.msg, kat.wots, kat.auth)).to.equal(false);
  });

  it("verifyAndRecord records the result and stays under the EIP-7825 cap", async function () {
    const v = await deploy();
    const est = await v.verifyAndRecord.estimateGas(...V());
    expect(est).to.be.lessThan(16777216n); // 2^24 Fusaka per-tx cap
    await (await v.verifyAndRecord(...V())).wait();
    expect(await v.lastResult()).to.equal(true);
    expect(await v.verifyCount()).to.equal(1n);
  });
});