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

// ---------------------------------------------------------------------------
// AerePQAggregate: constant-cost on-chain verification of a t-of-n POST-QUANTUM
// committee authorization via ONE succinct SP1 proof.
//
// WHAT IS UNDER TEST (real, on-chain): AerePQAggregateVerifier and the ERC-7579
// validator module AerePQAggregateModule. The verifier binds a committee root, t,
// n, scheme and the message digest into the SP1 public values and gates on a single
// gateway verifyProof call whose cost does not grow with n.
//
// TEST DOUBLE: MockSp1Verifier is the deployed SP1 gateway stand-in. Its verifyProof
// REVERTS on an invalid proof (exact production gateway semantics), so it faithfully
// exercises the fail-closed path. It accepts only (vkey, publicValues, proof) triples
// pre-registered via setValid(), so a "valid proof" in these tests means the gateway
// asserted the exact bound public values.
//
// HONEST BOUNDARY: the OFF-CHAIN SP1 aggregation guest circuit (verify n ML-DSA/Falcon
// signatures against the committee root inside the zkVM and emit these public values)
// is NOT implemented and is [MEASURE]. These tests prove the on-chain verifier + module
// only, against a mock gateway, exactly like AereComputeMarketV3 and the zk light clients.
// ---------------------------------------------------------------------------

const SCHEME_MLDSA65 = 5;
const VALIDATION_SUCCESS = 0n;
const SIG_VALIDATION_FAILED = 1n;
const EIP1271_MAGIC = "0x1626ba7e";
const EIP1271_FAIL = "0xffffffff";

function abi() {
  return ethers.AbiCoder.defaultAbiCoder();
}

// The fixed 192-byte SP1 public-values layout the aggregation circuit must emit.
function encodePublicValues({ root, threshold, size, scheme, digest, validCount }) {
  return abi().encode(
    ["bytes32", "uint256", "uint256", "uint256", "bytes32", "uint256"],
    [root, threshold, size, scheme, digest, validCount]
  );
}

// Marker the MockSp1Verifier keys a valid proof on: keccak256(abi.encode(vkey, publicValues, proof)).
function proofMarker(vkey, publicValues, proof) {
  return ethers.keccak256(abi().encode(["bytes32", "bytes", "bytes"], [vkey, publicValues, proof]));
}

// A committee root: a commitment to the n authorized PQC public keys. The contract only
// compares a bytes32, so any binding commitment works; we use keccak over a label + n.
function committeeRoot(label, n) {
  return ethers.keccak256(ethers.toUtf8Bytes(`committee:${label}:n=${n}`));
}

describe("AerePQAggregate (verifier + ERC-7579 module)", function () {
  let sp1, verifier, module;
  let deployer, account, relayer, other;
  const VKEY = ethers.keccak256(ethers.toUtf8Bytes("sp1-program:aere-pq-aggregate-mldsa65-v1"));

  beforeEach(async function () {
    [deployer, account, relayer, other] = await ethers.getSigners();

    sp1 = await (await ethers.getContractFactory("MockSp1Verifier")).deploy();
    await sp1.waitForDeployment();

    verifier = await (await ethers.getContractFactory("AerePQAggregateVerifier")).deploy(
      await sp1.getAddress(),
      VKEY
    );
    await verifier.waitForDeployment();

    module = await (await ethers.getContractFactory("AerePQAggregateModule")).deploy(
      await verifier.getAddress()
    );
    await module.waitForDeployment();
  });

  // Register a committee and return {id, root, threshold, size, scheme}.
  async function registerCommittee({ label = "A", threshold = 3, size = 5, scheme = SCHEME_MLDSA65 } = {}) {
    const root = committeeRoot(label, size);
    const idBefore = await verifier.committeeCount();
    await verifier.connect(deployer).registerCommittee(root, threshold, size, scheme);
    const id = idBefore + 1n;
    return { id, root, threshold, size, scheme };
  }

  // Build public values + a proof, and (optionally) register it valid in the mock gateway.
  async function makeProof(c, { digest, validCount, proof = "0xa1b2c3d4", register = true } = {}) {
    const publicValues = encodePublicValues({
      root: c.root,
      threshold: c.threshold,
      size: c.size,
      scheme: c.scheme,
      digest,
      validCount,
    });
    if (register) {
      await sp1.setValid(proofMarker(VKEY, publicValues, proof), true);
    }
    return { publicValues, proof };
  }

  // =========================================================================
  // 1. Registration
  // =========================================================================
  describe("registration", function () {
    it("registers a committee and rejects bad params", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      const info = await verifier.getCommittee(c.id);
      expect(info.exists).to.equal(true);
      expect(info.threshold).to.equal(3);
      expect(info.size).to.equal(5);
      expect(info.scheme).to.equal(SCHEME_MLDSA65);
      expect(info.root).to.equal(c.root);
      expect(await verifier.committeeExists(c.id)).to.equal(true);

      await expect(
        verifier.registerCommittee(ethers.ZeroHash, 3, 5, SCHEME_MLDSA65)
      ).to.be.revertedWithCustomError(verifier, "ZeroRoot");
      await expect(
        verifier.registerCommittee(committeeRoot("bad", 5), 0, 5, SCHEME_MLDSA65)
      ).to.be.revertedWithCustomError(verifier, "BadCommitteeParams");
      await expect(
        verifier.registerCommittee(committeeRoot("bad", 5), 6, 5, SCHEME_MLDSA65)
      ).to.be.revertedWithCustomError(verifier, "BadCommitteeParams");
      await expect(
        verifier.registerCommittee(committeeRoot("bad", 5), 3, 5, 0)
      ).to.be.revertedWithCustomError(verifier, "BadScheme");
    });

    // L1 (LOW) fix: verifyProof returns void, so a code-less gateway makes the call succeed silently
    // (fail-OPEN = post-quantum authorization bypass). The constructor must reject a code-less gateway,
    // exactly like the sibling AereFinalityCertificateVerifier.
    it("rejects a code-less SP1 gateway at construction (L1 fail-open guard)", async function () {
      const F = await ethers.getContractFactory("AerePQAggregateVerifier");
      // A non-zero EOA / empty address has no code.
      await expect(F.deploy(other.address, VKEY)).to.be.revertedWithCustomError(verifier, "VerifierHasNoCode");
      // A zero gateway still reverts on the earlier ZeroGateway check.
      await expect(F.deploy(ethers.ZeroAddress, VKEY)).to.be.revertedWithCustomError(verifier, "ZeroGateway");
      // Control: the real (code-bearing) mock gateway deploys fine.
      const ok = await F.deploy(await sp1.getAddress(), VKEY);
      await ok.waitForDeployment();
      expect(await ok.SP1_GATEWAY()).to.equal(await sp1.getAddress());
    });
  });

  // =========================================================================
  // 2. Verifier: valid proof authorizes; invalid / wrong-digest / wrong-root reject
  // =========================================================================
  describe("verifier fail-closed semantics", function () {
    it("a valid aggregate proof authorizes (view + recorded)", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      const digest = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:approve-transfer"));
      const { publicValues, proof } = await makeProof(c, { digest, validCount: 3 });

      // Strict view path does not revert.
      expect(await verifier.verifyAggregate(c.id, digest, publicValues, proof)).to.equal(true);
      // Non-reverting path returns true.
      expect(await verifier.tryVerifyAggregate(c.id, digest, publicValues, proof)).to.equal(true);

      // State path records provenance and emits.
      await expect(verifier.connect(relayer).recordAggregate(c.id, digest, publicValues, proof))
        .to.emit(verifier, "AggregateVerified")
        .withArgs(c.id, digest, c.threshold, 3, anyUint());
      const att = await verifier.getAttestation(c.id, digest);
      expect(att.exists).to.equal(true);
      expect(att.validCount).to.equal(3);
      expect(await verifier.isAttested(c.id, digest)).to.equal(true);
    });

    it("rejects an invalid proof fail-closed (proof not registered in the gateway)", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      const digest = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:1"));
      // register=false => the mock gateway will revert MockInvalidProof.
      const { publicValues, proof } = await makeProof(c, { digest, validCount: 3, register: false });

      await expect(
        verifier.verifyAggregate(c.id, digest, publicValues, proof)
      ).to.be.revertedWithCustomError(sp1, "MockInvalidProof");
      expect(await verifier.tryVerifyAggregate(c.id, digest, publicValues, proof)).to.equal(false);
      await expect(
        verifier.recordAggregate(c.id, digest, publicValues, proof)
      ).to.be.revertedWithCustomError(sp1, "MockInvalidProof");
      expect(await verifier.isAttested(c.id, digest)).to.equal(false);
    });

    it("rejects a proof for the WRONG message digest", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      const signedDigest = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:the-committee-signed-this"));
      const claimedDigest = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:but-caller-claims-this"));
      // The proof binds signedDigest and is even registered valid in the gateway...
      const { publicValues, proof } = await makeProof(c, { digest: signedDigest, validCount: 3 });

      // ...but verifying it against a different claimed digest is rejected BEFORE the gateway call.
      await expect(
        verifier.verifyAggregate(c.id, claimedDigest, publicValues, proof)
      ).to.be.revertedWithCustomError(verifier, "MessageDigestMismatch");
      expect(await verifier.tryVerifyAggregate(c.id, claimedDigest, publicValues, proof)).to.equal(false);
    });

    it("rejects a proof whose committee root does NOT match the registered committee", async function () {
      const c = await registerCommittee({ label: "real", threshold: 3, size: 5 });
      const digest = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:x"));
      // Public values carry a DIFFERENT root than the registered committee, and are registered valid.
      const forgedRoot = committeeRoot("attacker-key-set", 5);
      const publicValues = encodePublicValues({
        root: forgedRoot,
        threshold: c.threshold,
        size: c.size,
        scheme: c.scheme,
        digest,
        validCount: 3,
      });
      const proof = "0xdeadbeef";
      await sp1.setValid(proofMarker(VKEY, publicValues, proof), true);

      await expect(
        verifier.verifyAggregate(c.id, digest, publicValues, proof)
      ).to.be.revertedWithCustomError(verifier, "CommitteeRootMismatch");
      expect(await verifier.tryVerifyAggregate(c.id, digest, publicValues, proof)).to.equal(false);
    });

    it("rejects when the proven valid-signature count is below threshold", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      const digest = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:subquorum"));
      // Only 2 valid signatures proven, threshold is 3. Even a gateway-valid proof must be rejected.
      const { publicValues, proof } = await makeProof(c, { digest, validCount: 2 });

      await expect(
        verifier.verifyAggregate(c.id, digest, publicValues, proof)
      ).to.be.revertedWithCustomError(verifier, "BelowThreshold");
      expect(await verifier.tryVerifyAggregate(c.id, digest, publicValues, proof)).to.equal(false);
    });

    it("rejects on threshold / size / scheme mismatch and malformed public values", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5, scheme: SCHEME_MLDSA65 });
      const digest = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:mismatch"));

      const wrongT = encodePublicValues({ root: c.root, threshold: 2, size: 5, scheme: SCHEME_MLDSA65, digest, validCount: 5 });
      await sp1.setValid(proofMarker(VKEY, wrongT, "0x01"), true);
      await expect(verifier.verifyAggregate(c.id, digest, wrongT, "0x01")).to.be.revertedWithCustomError(verifier, "ThresholdMismatch");

      const wrongN = encodePublicValues({ root: c.root, threshold: 3, size: 9, scheme: SCHEME_MLDSA65, digest, validCount: 3 });
      await sp1.setValid(proofMarker(VKEY, wrongN, "0x02"), true);
      await expect(verifier.verifyAggregate(c.id, digest, wrongN, "0x02")).to.be.revertedWithCustomError(verifier, "SizeMismatch");

      const wrongScheme = encodePublicValues({ root: c.root, threshold: 3, size: 5, scheme: 3, digest, validCount: 3 });
      await sp1.setValid(proofMarker(VKEY, wrongScheme, "0x03"), true);
      await expect(verifier.verifyAggregate(c.id, digest, wrongScheme, "0x03")).to.be.revertedWithCustomError(verifier, "SchemeMismatch");

      // Malformed (not 192 bytes) public values: strict reverts, non-reverting returns false.
      const malformed = "0x1234";
      await expect(verifier.verifyAggregate(c.id, digest, malformed, "0x04")).to.be.revertedWithCustomError(verifier, "BadPublicValues");
      expect(await verifier.tryVerifyAggregate(c.id, digest, malformed, "0x04")).to.equal(false);
    });

    it("rejects an unknown committee", async function () {
      const digest = ethers.keccak256(ethers.toUtf8Bytes("x"));
      const publicValues = encodePublicValues({
        root: committeeRoot("ghost", 5), threshold: 3, size: 5, scheme: SCHEME_MLDSA65, digest, validCount: 3,
      });
      await expect(
        verifier.verifyAggregate(999, digest, publicValues, "0x05")
      ).to.be.revertedWithCustomError(verifier, "UnknownCommittee");
      expect(await verifier.tryVerifyAggregate(999, digest, publicValues, "0x05")).to.equal(false);
    });
  });

  // =========================================================================
  // 3. Gas is INDEPENDENT of the committee size n (the whole point)
  // =========================================================================
  describe("constant on-chain cost regardless of committee size n", function () {
    it("recordAggregate costs the same for a tiny and a huge committee", async function () {
      const small = await registerCommittee({ label: "small", threshold: 3, size: 5 });
      const huge = await registerCommittee({ label: "huge", threshold: 3, size: 1_000_000 });

      const digest = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:gas-parity"));
      const pSmall = await makeProof(small, { digest, validCount: 3, proof: "0xaa" });
      const pHuge = await makeProof(huge, { digest, validCount: 3, proof: "0xbb" });

      const rcSmall = await (await verifier.recordAggregate(small.id, digest, pSmall.publicValues, pSmall.proof)).wait();
      const rcHuge = await (await verifier.recordAggregate(huge.id, digest, pHuge.publicValues, pHuge.proof)).wait();

      // The only residual difference is the calldata encoding of the integer n itself (a few bytes),
      // which is O(log n), NOT the O(n) tens-of-thousands-of-gas-per-signature a linear verifier costs.
      const diff = rcSmall.gasUsed > rcHuge.gasUsed ? rcSmall.gasUsed - rcHuge.gasUsed : rcHuge.gasUsed - rcSmall.gasUsed;
      expect(diff).to.be.lessThan(500n);
      // Sanity: n grew 200,000x yet gas barely moved and stayed well under any per-member budget.
      expect(rcHuge.gasUsed).to.be.lessThan(200_000n);
      console.log(`        gas n=5: ${rcSmall.gasUsed}  gas n=1,000,000: ${rcHuge.gasUsed}  diff: ${diff}`);
    });
  });

  // =========================================================================
  // 4. ERC-7579 validator module
  // =========================================================================
  describe("AerePQAggregateModule (ERC-7579 validator, type 1)", function () {
    function makeUserOp(sender, signature) {
      return {
        sender,
        nonce: 0,
        initCode: "0x",
        callData: "0x",
        accountGasLimits: ethers.ZeroHash,
        preVerificationGas: 0,
        gasFees: ethers.ZeroHash,
        paymasterAndData: "0x",
        signature,
      };
    }

    // module signature = abi.encode(bytes publicValues, bytes proofBytes)
    function moduleSig(publicValues, proof) {
      return abi().encode(["bytes", "bytes"], [publicValues, proof]);
    }

    it("declares module type 1 (validator)", async function () {
      expect(await module.isModuleType(1)).to.equal(true);
      expect(await module.isModuleType(2)).to.equal(false);
      expect(await module.name()).to.equal("aere.pq-aggregate-validator.1.0.0");
    });

    it("install binds the account to a registered committee; unknown committee reverts", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      await expect(
        module.connect(account).onInstall(abi().encode(["uint256"], [999]))
      ).to.be.revertedWithCustomError(module, "UnknownCommittee");

      await module.connect(account).onInstall(abi().encode(["uint256"], [c.id]));
      expect(await module.committeeOf(account.address)).to.equal(c.id);
    });

    it("validateUserOp returns SUCCESS on a valid aggregate proof over the userOpHash", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      await module.connect(account).onInstall(abi().encode(["uint256"], [c.id]));

      const userOpHash = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:committee-approved"));
      const { publicValues, proof } = await makeProof(c, { digest: userOpHash, validCount: 4 });
      const userOp = makeUserOp(account.address, moduleSig(publicValues, proof));

      const res = await module.connect(account).validateUserOp.staticCall(userOp, userOpHash);
      expect(res).to.equal(VALIDATION_SUCCESS);
    });

    it("validateUserOp returns FAILED (no revert) on an invalid proof", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      await module.connect(account).onInstall(abi().encode(["uint256"], [c.id]));

      const userOpHash = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:forged"));
      // Not registered in the gateway => the proof is invalid.
      const { publicValues, proof } = await makeProof(c, { digest: userOpHash, validCount: 4, register: false });
      const userOp = makeUserOp(account.address, moduleSig(publicValues, proof));

      const res = await module.connect(account).validateUserOp.staticCall(userOp, userOpHash);
      expect(res).to.equal(SIG_VALIDATION_FAILED);
    });

    it("validateUserOp returns FAILED on a wrong-digest proof and on a malformed signature", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      await module.connect(account).onInstall(abi().encode(["uint256"], [c.id]));

      // Valid proof, but bound to a DIFFERENT digest than the userOpHash being validated.
      const otherDigest = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:other-op"));
      const { publicValues, proof } = await makeProof(c, { digest: otherDigest, validCount: 4 });
      const userOpHash = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:this-op"));
      const userOp = makeUserOp(account.address, moduleSig(publicValues, proof));
      expect(await module.connect(account).validateUserOp.staticCall(userOp, userOpHash)).to.equal(SIG_VALIDATION_FAILED);

      // Malformed signature blob must NOT revert; it returns FAILED.
      const junkOp = makeUserOp(account.address, "0x1234");
      expect(await module.connect(account).validateUserOp.staticCall(junkOp, userOpHash)).to.equal(SIG_VALIDATION_FAILED);
    });

    it("validateUserOp returns FAILED for an account with no committee configured", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      const userOpHash = ethers.keccak256(ethers.toUtf8Bytes("userOpHash:unconfigured"));
      const { publicValues, proof } = await makeProof(c, { digest: userOpHash, validCount: 3 });
      // `other` never installed the module.
      const userOp = makeUserOp(other.address, moduleSig(publicValues, proof));
      expect(await module.connect(other).validateUserOp.staticCall(userOp, userOpHash)).to.equal(SIG_VALIDATION_FAILED);
    });

    it("isValidSignatureWithSender returns the ERC-1271 magic on valid, fail on invalid", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      await module.connect(account).onInstall(abi().encode(["uint256"], [c.id]));

      const hash = ethers.keccak256(ethers.toUtf8Bytes("erc1271:committee-approved-message"));
      const good = await makeProof(c, { digest: hash, validCount: 3, proof: "0xc0ffee" });
      const okMagic = await module
        .connect(account)
        .isValidSignatureWithSender(other.address, hash, moduleSig(good.publicValues, good.proof));
      expect(okMagic).to.equal(EIP1271_MAGIC);

      // Invalid proof => fail magic, no revert.
      const bad = await makeProof(c, { digest: hash, validCount: 3, proof: "0xbadbad", register: false });
      const failMagic = await module
        .connect(account)
        .isValidSignatureWithSender(other.address, hash, moduleSig(bad.publicValues, bad.proof));
      expect(failMagic).to.equal(EIP1271_FAIL);

      // Malformed signature => fail magic, no revert.
      const junkMagic = await module.connect(account).isValidSignatureWithSender(other.address, hash, "0xabcd");
      expect(junkMagic).to.equal(EIP1271_FAIL);
    });

    it("uninstall clears the binding; validation then fails closed", async function () {
      const c = await registerCommittee({ threshold: 3, size: 5 });
      await module.connect(account).onInstall(abi().encode(["uint256"], [c.id]));
      await module.connect(account).onUninstall("0x");
      expect(await module.committeeOf(account.address)).to.equal(0n);

      const hash = ethers.keccak256(ethers.toUtf8Bytes("after-uninstall"));
      const p = await makeProof(c, { digest: hash, validCount: 3 });
      expect(
        await module.connect(account).isValidSignatureWithSender(other.address, hash, moduleSig(p.publicValues, p.proof))
      ).to.equal(EIP1271_FAIL);
    });
  });
});

// Matches any uint (for the block.number arg in the AggregateVerified event).
function anyUint() {
  const { anyValue } = require("@nomicfoundation/hardhat-chai-matchers/withArgs");
  return anyValue;
}