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

// ---------------------------------------------------------------------------
// AereThresholdPQCRegistry -- on-chain t-of-n POST-QUANTUM threshold authorization.
//
// A committee of n PQC signers each hold their own NIST PQC key. A valid authorization
// requires >= t DISTINCT members to each PQC-sign the SAME domain-separated challenge,
// with every leg verified on-chain by the scheme's live precompile (0x0AE1..0x0AE4).
//
// HONEST label: this is a threshold MULTISIG of independent PQC signatures, NOT a single
// aggregate threshold-PQC signature (that is open research; see docs/THRESHOLD-PQC-*.md).
//
// Precompile modelling: Hardhat has no PQC precompile, so we install the repo's faithful
// MockPQCPrecompile at 0x0AE1..0x0AE4 (hardhat_setCode). It parses the input at the
// precompiles' SPEC offsets and accepts iff commitment == keccak256(pk || message) -- the
// same model the shipped AerePQCAttestation adversarial suite uses. Because the mock parses
// at the spec offsets (not offsets copied from the contract), a passing authorize() proves
// BOTH the threshold-counting logic AND that the contract built the precompile envelope at
// the correct offsets. The REAL envelope is separately proven bit-for-bit against the LIVE
// mainnet precompile via eth_call (scripts/pqc/verify-live-precompile.js).
// ---------------------------------------------------------------------------

const PRECOMPILE = {
  1: ethers.getAddress("0x0000000000000000000000000000000000000ae1"), // Falcon-512
  2: ethers.getAddress("0x0000000000000000000000000000000000000ae2"), // Falcon-1024
  3: ethers.getAddress("0x0000000000000000000000000000000000000ae3"), // ML-DSA-44
  4: ethers.getAddress("0x0000000000000000000000000000000000000ae4"), // SLH-DSA-128s
};

const META = {
  1: { pkLen: 897, pkHeader: 0x09, esigHeader: 0x29, falcon: true },
  2: { pkLen: 1793, pkHeader: 0x0a, esigHeader: 0x2a, falcon: true },
  3: { pkLen: 1312, sigLen: 2420, falcon: false },
  4: { pkLen: 32, sigLen: 7856, falcon: false },
};

// Deterministic seeded RNG (keccak hash-chain) so the whole suite is reproducible.
function makeRng(seedText) {
  let seed = ethers.keccak256(ethers.toUtf8Bytes(seedText));
  function bytes(n) {
    const out = new Uint8Array(n);
    let i = 0;
    while (i < n) {
      seed = ethers.keccak256(seed);
      const chunk = ethers.getBytes(seed);
      for (let j = 0; j < chunk.length && i < n; j++, i++) out[i] = chunk[j];
    }
    return out;
  }
  return { bytes };
}

function makePubKey(scheme, rng) {
  const m = META[scheme];
  const raw = rng.bytes(m.pkLen);
  if (m.falcon) raw[0] = m.pkHeader; // Falcon requires the header byte
  return ethers.hexlify(raw);
}

// commitment the mock expects for a genuine signature over `message` by `pubKey`.
function commitmentFor(pubKey, message) {
  return ethers.keccak256(ethers.concat([pubKey, message]));
}

// Build a signature envelope carrying a chosen commitment (genuine or forged).
function envelope(scheme, commitment, rng) {
  const m = META[scheme];
  if (m.falcon) {
    const nonce = rng.bytes(40);
    const esig = ethers.concat([new Uint8Array([m.esigHeader]), commitment]); // 33 bytes
    return ethers.hexlify(ethers.concat([nonce, esig]));
  } else {
    const sig = new Uint8Array(m.sigLen);
    sig.set(ethers.getBytes(commitment), 0);
    return ethers.hexlify(sig);
  }
}

// A genuine signature by `pubKey` over `message`.
function genuine(scheme, pubKey, message, rng) {
  return envelope(scheme, commitmentFor(pubKey, message), rng);
}

// Flip one byte inside the commitment region -> a genuine forgery the precompile rejects.
function tamperCommitment(scheme, sigHex) {
  const b = ethers.getBytes(sigHex);
  const idx = META[scheme].falcon ? 41 : 0;
  b[idx] ^= 0x01;
  return ethers.hexlify(b);
}

describe("AereThresholdPQCRegistry (post-quantum threshold multisig)", function () {
  let reg, mockCode;

  before(async function () {
    // Deploy one mock per scheme; capture its runtime bytecode (immutable scheme baked in).
    mockCode = {};
    const F = await ethers.getContractFactory("MockPQCPrecompile");
    for (const scheme of [1, 2, 3, 4]) {
      const mock = await F.deploy(scheme);
      await mock.waitForDeployment();
      mockCode[scheme] = await ethers.provider.getCode(await mock.getAddress());
    }
  });

  async function installMocks() {
    for (const scheme of [1, 2, 3, 4]) {
      await network.provider.send("hardhat_setCode", [PRECOMPILE[scheme], mockCode[scheme]]);
    }
  }

  beforeEach(async function () {
    await installMocks();
    const R = await ethers.getContractFactory("AereThresholdPQCRegistry");
    reg = await R.deploy();
    await reg.waitForDeployment();
  });

  // Register a fresh n-member committee of `scheme`, threshold t. Returns the committee id,
  // its member public keys, and a shared rng.
  async function registerCommittee(scheme, n, t, seed) {
    const rng = makeRng(seed || `c-${scheme}-${n}-${t}`);
    const pubKeys = [];
    for (let i = 0; i < n; i++) pubKeys.push(makePubKey(scheme, rng));
    const tx = await reg.registerCommittee(scheme, t, pubKeys);
    const rc = await tx.wait();
    const log = rc.logs.map((l) => reg.interface.parseLog(l)).find((p) => p && p.name === "CommitteeRegistered");
    return { id: Number(log.args.committeeId), pubKeys, rng, scheme, n, t };
  }

  // Build the `sigs` array of {memberIndex, signature} for the given member indices, each a
  // genuine PQC signature over the committee's current authorization challenge.
  async function legsFor(c, messageHash, memberIndices) {
    const nonce = await reg.nonceOf(c.id);
    const challenge = await reg.authChallenge(c.id, nonce, messageHash);
    return memberIndices.map((idx) => ({
      memberIndex: idx,
      signature: genuine(c.scheme, c.pubKeys[idx], challenge, c.rng),
    }));
  }

  // -----------------------------------------------------------------------
  // Registration validation
  // -----------------------------------------------------------------------
  describe("registerCommittee", function () {
    it("registers a committee for every scheme and stores scheme/t/n + members commitment", async function () {
      for (const scheme of [1, 2, 3, 4]) {
        const c = await registerCommittee(scheme, 3, 2, `reg-${scheme}`);
        const got = await reg.getCommittee(c.id);
        expect(got.exists).to.equal(true);
        expect(Number(got.scheme)).to.equal(scheme);
        expect(Number(got.threshold)).to.equal(2);
        expect(Number(got.size)).to.equal(3);
        expect(Number(got.authNonce)).to.equal(0);
        // members commitment matches abi.encode(pubKeys)
        const expectHash = ethers.keccak256(
          ethers.AbiCoder.defaultAbiCoder().encode(["bytes[]"], [c.pubKeys]),
        );
        expect(got.membersHash).to.equal(expectHash);
        // stored keys are retrievable
        expect(await reg.getMemberPubKey(c.id, 0)).to.equal(c.pubKeys[0]);
        expect(await reg.precompileFor(scheme)).to.equal(PRECOMPILE[scheme]);
      }
    });

    it("rejects bad committee params (t=0, t>n, empty, n>255)", async function () {
      const rng = makeRng("bad");
      const three = [0, 1, 2].map(() => makePubKey(1, rng));
      await expect(reg.registerCommittee(1, 0, three)).to.be.revertedWithCustomError(reg, "BadCommitteeParams");
      await expect(reg.registerCommittee(1, 4, three)).to.be.revertedWithCustomError(reg, "BadCommitteeParams");
      await expect(reg.registerCommittee(1, 1, [])).to.be.revertedWithCustomError(reg, "BadCommitteeParams");
    });

    it("rejects an unknown scheme and malformed member keys", async function () {
      const rng = makeRng("mal");
      const goodF = makePubKey(1, rng);
      // unknown scheme
      await expect(reg.registerCommittee(5, 1, [goodF])).to.be.revertedWithCustomError(reg, "InvalidScheme");
      // Falcon-512 wrong header
      const badHeader = ethers.getBytes(makePubKey(1, rng));
      badHeader[0] = 0x0a;
      await expect(reg.registerCommittee(1, 1, [ethers.hexlify(badHeader)])).to.be.revertedWithCustomError(
        reg,
        "InvalidPubKey",
      );
      // ML-DSA-44 wrong length
      await expect(reg.registerCommittee(3, 1, ["0x" + "00".repeat(1311)])).to.be.revertedWithCustomError(
        reg,
        "InvalidPubKey",
      );
      // the offending member index is reported
      await expect(reg.registerCommittee(1, 2, [goodF, "0x" + "09" + "00".repeat(10)]))
        .to.be.revertedWithCustomError(reg, "InvalidPubKey")
        .withArgs(1);
    });

    it("rejects a committee with duplicate keys (distinctness is by KEY, not index)", async function () {
      // Authorization counts distinct signers by member index, so a repeated key would let one
      // keyholder fill multiple slots and reach threshold alone. registerCommittee must reject it.
      const rng = makeRng("dup-key-registry");
      const keys = [];
      for (let i = 0; i < 4; i++) keys.push(makePubKey(1, rng));
      keys[2] = keys[0]; // member 2 reuses member 0's key
      await expect(reg.registerCommittee(1, 2, keys))
        .to.be.revertedWithCustomError(reg, "DuplicatePubKey")
        .withArgs(2);
    });
  });

  // -----------------------------------------------------------------------
  // t-of-n happy path + threshold gate
  // -----------------------------------------------------------------------
  describe("authorize: t-of-n happy path and threshold gate", function () {
    it("authorizes with EXACTLY t distinct valid PQC legs (2-of-3, Falcon-512)", async function () {
      const c = await registerCommittee(1, 3, 2, "happy-f512");
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("release vault payment #1"));
      const sigs = await legsFor(c, messageHash, [0, 2]); // any 2 distinct members

      const before = await reg.authorizationCount();
      await expect(reg.authorize(c.id, messageHash, sigs))
        .to.emit(reg, "Authorized")
        .withArgs(c.id, 0, messageHash, await reg.authChallenge(c.id, 0, messageHash), 2, anyBlock());
      expect(await reg.authorizationCount()).to.equal(before + 1n);
      expect(Number(await reg.nonceOf(c.id))).to.equal(1);

      const a = await reg.getAuthorization(c.id, 0);
      expect(a.exists).to.equal(true);
      expect(a.messageHash).to.equal(messageHash);
      expect(Number(a.quorum)).to.equal(2);
      expect(await reg.isAuthorized(c.id, 0, messageHash)).to.equal(true);
      expect(await reg.isAuthorized(c.id, 0, ethers.ZeroHash)).to.equal(false);
    });

    it("authorizes for EVERY scheme with a t-of-n quorum", async function () {
      for (const scheme of [1, 2, 3, 4]) {
        const c = await registerCommittee(scheme, 4, 3, `all-${scheme}`);
        const messageHash = ethers.keccak256(ethers.toUtf8Bytes(`msg-${scheme}`));
        const sigs = await legsFor(c, messageHash, [0, 1, 3]);
        await expect(reg.authorize(c.id, messageHash, sigs)).to.emit(reg, "Authorized");
        expect(Number(await reg.nonceOf(c.id))).to.equal(1);
      }
    });

    it("REJECTS a below-threshold quorum (t-1 valid legs)", async function () {
      const c = await registerCommittee(1, 3, 2, "below");
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("only one signer"));
      const sigs = await legsFor(c, messageHash, [1]); // 1 < t=2
      await expect(reg.authorize(c.id, messageHash, sigs))
        .to.be.revertedWithCustomError(reg, "BelowThreshold")
        .withArgs(1, 2);
      // nothing recorded, nonce frozen
      expect(await reg.authorizationCount()).to.equal(0n);
      expect(Number(await reg.nonceOf(c.id))).to.equal(0);
      expect((await reg.getAuthorization(c.id, 0)).exists).to.equal(false);
    });

    it("previewAuthorize reports quorum without recording (lenient view)", async function () {
      const c = await registerCommittee(2, 5, 3, "preview");
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("preview-msg"));
      const nonce = await reg.nonceOf(c.id);
      const three = await legsFor(c, messageHash, [0, 2, 4]);
      const two = three.slice(0, 2);
      expect((await reg.previewAuthorize(c.id, nonce, messageHash, three)).would).to.equal(true);
      expect((await reg.previewAuthorize(c.id, nonce, messageHash, two)).would).to.equal(false);
      // view did not mutate state
      expect(Number(await reg.nonceOf(c.id))).to.equal(0);
    });
  });

  // -----------------------------------------------------------------------
  // Distinctness, membership, message binding, replay
  // -----------------------------------------------------------------------
  describe("authorize: distinctness / membership / binding / replay", function () {
    it("does NOT double-count a duplicate member (reverts DuplicateMember)", async function () {
      const c = await registerCommittee(1, 3, 2, "dup");
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("dup-msg"));
      const nonce = await reg.nonceOf(c.id);
      const challenge = await reg.authChallenge(c.id, nonce, messageHash);
      // member 0 signs twice: two legs, both valid, but same index -> not a 2-of-3 quorum
      const m0a = { memberIndex: 0, signature: genuine(c.scheme, c.pubKeys[0], challenge, c.rng) };
      const m0b = { memberIndex: 0, signature: genuine(c.scheme, c.pubKeys[0], challenge, c.rng) };
      await expect(reg.authorize(c.id, messageHash, [m0a, m0b]))
        .to.be.revertedWithCustomError(reg, "DuplicateMember")
        .withArgs(0);
      expect(await reg.authorizationCount()).to.equal(0n);
    });

    it("rejects an out-of-range member index (non-member)", async function () {
      const c = await registerCommittee(1, 3, 2, "oor");
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("oor-msg"));
      const good = await legsFor(c, messageHash, [0]);
      const bogus = { memberIndex: 9, signature: good[0].signature };
      await expect(reg.authorize(c.id, messageHash, [good[0], bogus]))
        .to.be.revertedWithCustomError(reg, "MemberOutOfRange")
        .withArgs(9);
    });

    it("rejects a leg signed by a NON-member key at a valid index", async function () {
      const c = await registerCommittee(1, 3, 2, "foreign");
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("foreign-msg"));
      const nonce = await reg.nonceOf(c.id);
      const challenge = await reg.authChallenge(c.id, nonce, messageHash);
      const outsiderKey = makePubKey(1, c.rng); // never registered
      const legGood = { memberIndex: 0, signature: genuine(c.scheme, c.pubKeys[0], challenge, c.rng) };
      const legForeign = { memberIndex: 1, signature: genuine(c.scheme, outsiderKey, challenge, c.rng) };
      await expect(reg.authorize(c.id, messageHash, [legGood, legForeign]))
        .to.be.revertedWithCustomError(reg, "LegVerificationFailed")
        .withArgs(1);
    });

    it("rejects a leg that is a tampered / invalid PQC signature", async function () {
      const c = await registerCommittee(2, 3, 2, "tamper");
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("tamper-msg"));
      const sigs = await legsFor(c, messageHash, [0, 1]);
      sigs[1].signature = tamperCommitment(c.scheme, sigs[1].signature);
      await expect(reg.authorize(c.id, messageHash, sigs))
        .to.be.revertedWithCustomError(reg, "LegVerificationFailed")
        .withArgs(1);
    });

    it("rejects legs that are genuine over a DIFFERENT message", async function () {
      const c = await registerCommittee(1, 3, 2, "wrongmsg");
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("intended"));
      const otherHash = ethers.keccak256(ethers.toUtf8Bytes("attacker-substituted"));
      // sign the challenge for `otherHash`, then submit against `messageHash`
      const nonce = await reg.nonceOf(c.id);
      const otherChallenge = await reg.authChallenge(c.id, nonce, otherHash);
      const sigs = [0, 1].map((idx) => ({
        memberIndex: idx,
        signature: genuine(c.scheme, c.pubKeys[idx], otherChallenge, c.rng),
      }));
      await expect(reg.authorize(c.id, messageHash, sigs)).to.be.revertedWithCustomError(reg, "LegVerificationFailed");
      expect(await reg.authorizationCount()).to.equal(0n);
    });

    it("prevents REPLAY of a used authorization (nonce advances, same sigs no longer verify)", async function () {
      const c = await registerCommittee(1, 3, 2, "replay");
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("pay once"));
      const sigs = await legsFor(c, messageHash, [0, 1]); // legs bound to nonce 0
      await reg.authorize(c.id, messageHash, sigs); // consumes nonce 0
      expect(Number(await reg.nonceOf(c.id))).to.equal(1);
      // exact replay: challenge now uses nonce 1, so the nonce-0 legs fail verification
      await expect(reg.authorize(c.id, messageHash, sigs)).to.be.revertedWithCustomError(reg, "LegVerificationFailed");
      expect(await reg.authorizationCount()).to.equal(1n);
      expect(Number(await reg.nonceOf(c.id))).to.equal(1);
    });

    it("allows a fresh authorization at the new nonce after a successful one", async function () {
      const c = await registerCommittee(3, 4, 3, "second");
      const msg1 = ethers.keccak256(ethers.toUtf8Bytes("first"));
      await reg.authorize(c.id, msg1, await legsFor(c, msg1, [0, 1, 2]));
      expect(Number(await reg.nonceOf(c.id))).to.equal(1);
      const msg2 = ethers.keccak256(ethers.toUtf8Bytes("second"));
      await reg.authorize(c.id, msg2, await legsFor(c, msg2, [1, 2, 3]));
      expect(Number(await reg.nonceOf(c.id))).to.equal(2);
      expect(await reg.isAuthorized(c.id, 0, msg1)).to.equal(true);
      expect(await reg.isAuthorized(c.id, 1, msg2)).to.equal(true);
    });

    it("reverts on an unknown committee", async function () {
      await expect(reg.authorize(999, ethers.ZeroHash, [])).to.be.revertedWithCustomError(reg, "UnknownCommittee");
    });

    it("treats an empty (pre-fork / wiped) precompile as an invalid leg, not an authorization", async function () {
      const c = await registerCommittee(2, 3, 2, "wiped");
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("wiped-msg"));
      const sigs = await legsFor(c, messageHash, [0, 1]);
      await network.provider.send("hardhat_setCode", [PRECOMPILE[2], "0x"]); // wipe Falcon-1024 precompile
      await expect(reg.authorize(c.id, messageHash, sigs)).to.be.revertedWithCustomError(reg, "LegVerificationFailed");
      expect(await reg.authorizationCount()).to.equal(0n);
    });
  });

  // -----------------------------------------------------------------------
  // Gas ceiling under EIP-7825: measure the marginal cost per additional leg.
  // NB: the mock verify is cheaper than the live precompile's FIXED verify gas
  // (Falcon-512 40k, Falcon-1024 75k, ML-DSA-44 55k, SLH-DSA-128s 350k). So the
  // REAL per-leg cost = (measured marginal here, which is calldata + envelope build +
  // distinct-member bookkeeping + cheap mock verify) + the precompile's fixed verify gas.
  // This test reports the Solidity/calldata overhead per leg so the ceiling in
  // docs/THRESHOLD-PQC-2026-07-12.md is grounded in a measurement, not a guess.
  // -----------------------------------------------------------------------
  describe("gas: per-leg marginal cost (informs the EIP-7825 (scheme, t) ceiling)", function () {
    const CAP = 16_777_216n;
    const PRECOMPILE_VERIFY_GAS = { 1: 40000n, 2: 75000n, 3: 55000n, 4: 350000n };

    async function authGas(scheme, t, seed) {
      const n = t; // exactly-t quorum
      const c = await registerCommittee(scheme, n, t, seed);
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes(seed));
      const sigs = await legsFor(
        c,
        messageHash,
        Array.from({ length: t }, (_, i) => i),
      );
      const tx = await reg.authorize(c.id, messageHash, sigs);
      const rc = await tx.wait();
      return rc.gasUsed;
    }

    it("reports marginal Solidity+calldata overhead per leg and the projected real ceiling", async function () {
      for (const scheme of [1, 2, 3, 4]) {
        const g2 = await authGas(scheme, 2, `gas-${scheme}-2`);
        const g6 = await authGas(scheme, 6, `gas-${scheme}-6`);
        const overheadPerLeg = (g6 - g2) / 4n; // 4 extra legs between t=2 and t=6

        // Real per-leg = measured Solidity/calldata overhead (which already includes the
        // cheap mock verify) + the precompile's fixed verify gas. Slightly conservative
        // (double-counts the tiny mock verify, so the ceiling below is an under-estimate).
        const realPerLeg = overheadPerLeg + PRECOMPILE_VERIFY_GAS[scheme];
        const fixedOverhead = g2 - overheadPerLeg * 2n; // tx base + committee load + record
        const tMax = (CAP - fixedOverhead) / realPerLeg;

        console.log(
          `      scheme ${scheme}: mock authorize t=2 ${g2} gas, t=6 ${g6} gas; ` +
            `Solidity/calldata overhead/leg ~${overheadPerLeg}; ` +
            `+precompile ${PRECOMPILE_VERIFY_GAS[scheme]} => real/leg ~${realPerLeg}; ` +
            `projected t_max under 2^24 ~= ${tMax}`,
        );

        // Sanity: the per-leg cost is positive and the projected ceiling is a real bound.
        expect(overheadPerLeg > 0n).to.equal(true);
        expect(tMax > 0n).to.equal(true);
        // Falcon-512 (cheapest verify) must admit a larger quorum than SLH-DSA-128s.
        if (scheme === 1) expect(tMax > 20n).to.equal(true);
      }
    });
  });
});

// chai matcher helper: accept any block number in withArgs.
function anyBlock() {
  return (x) => typeof x === "bigint" || typeof x === "number";
}