const { expect } = require("chai");
const { ethers } = require("hardhat");

// Off-chain WOTS+ signer matching AerePQCVerifier exactly.
function chain(x, count) { for (let i = 0; i < count; i++) x = ethers.keccak256(x); return x; }
function digits(mHex) {
  const m = ethers.getBytes(mHex); const d = new Array(67); let csum = 0;
  for (let i = 0; i < 32; i++) { const b = m[i], hi = b >> 4, lo = b & 0x0f; d[2 * i] = hi; d[2 * i + 1] = lo; csum += (15 - hi) + (15 - lo); }
  d[64] = (csum >> 8) & 0x0f; d[65] = (csum >> 4) & 0x0f; d[66] = csum & 0x0f; return d;
}
function keygen() {
  const sk = []; for (let i = 0; i < 67; i++) sk.push(ethers.hexlify(ethers.randomBytes(32)));
  const pkEnds = sk.map(s => chain(s, 15));
  const compressedPK = ethers.keccak256(ethers.concat(pkEnds));
  return { sk, compressedPK };
}
function sign(sk, msgHash) { const d = digits(msgHash); return sk.map((s, i) => chain(s, d[i])); }

describe("AerePQCVerifier (WOTS+ post-quantum)", function () {
  async function deploy() {
    const V = await ethers.getContractFactory("AerePQCVerifier");
    const v = await V.deploy(); await v.waitForDeployment(); return v;
  }

  it("accepts a valid WOTS+ signature", async function () {
    const v = await deploy();
    const { sk, compressedPK } = keygen();
    const msgHash = ethers.keccak256(ethers.toUtf8Bytes("AERE is post-quantum ready"));
    const sig = sign(sk, msgHash);
    expect(await v.verify(msgHash, sig, compressedPK)).to.equal(true);
    // derivePublicKey helper agrees
    expect(await v.derivePublicKey(sk)).to.equal(compressedPK);
  });

  it("rejects a wrong public key", async function () {
    const v = await deploy();
    const { sk } = keygen();
    const msgHash = ethers.keccak256(ethers.toUtf8Bytes("m"));
    const sig = sign(sk, msgHash);
    expect(await v.verify(msgHash, sig, ethers.keccak256("0xdead"))).to.equal(false);
  });

  it("rejects a signature reused on a DIFFERENT message (forgery-resistant)", async function () {
    const v = await deploy();
    const { sk, compressedPK } = keygen();
    const m1 = ethers.keccak256(ethers.toUtf8Bytes("original"));
    const sig = sign(sk, m1);
    const m2 = ethers.keccak256(ethers.toUtf8Bytes("forged"));
    expect(await v.verify(m2, sig, compressedPK)).to.equal(false);
  });

  it("checksum blocks a naive forge (raising a digit without lowering checksum)", async function () {
    const v = await deploy();
    const { sk, compressedPK } = keygen();
    const m = ethers.keccak256(ethers.toUtf8Bytes("x"));
    const sig = sign(sk, m);
    // advance one message-chain by 1 (attacker can hash forward) — checksum digits stay, so PK mismatch
    const tampered = [...sig];
    tampered[0] = ethers.keccak256(tampered[0]);
    expect(await v.verify(m, tampered, compressedPK)).to.equal(false);
  });
});
