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

// ---------------------------------------------------------------------------
// AerePQFinalityCertificate: the ADDITIVE Post-Quantum Finality Certificate.
//
// WHAT IS UNDER TEST (real, on-chain): AerePQAttestationKeyRegistry (per-validator
// hash-based attestation keys + a deterministic validator-set root) and
// AereFinalityCertificateVerifier (binds ONE succinct SP1 proof to the registry root,
// enforces the ceil(2N/3) quorum, records PQ finality, exposes isFinalPQ).
//
// 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" here means the gateway asserted the
// exact bound public values.
//
// HONEST BOUNDARY: the OFF-CHAIN SP1 aggregation guest circuit (verify N hash-based
// XMSS/leanSig signatures against the validator-set root inside the zkVM and emit these
// public values) is NOT implemented and is [MEASURE]. These tests prove the on-chain
// verifier + registry only, against a mock gateway, exactly like AerePQAggregate and the
// zk light clients. Aere consensus stays classical ECDSA QBFT; this is purely additive.
// ---------------------------------------------------------------------------

const SCHEME_XMSS = 1;
const SCHEME_LEANSIG = 3;

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

// The fixed 192-byte SP1 public-values layout the aggregation circuit must emit.
function encodePV({ blockHash, height, root, size, quorum, domain }) {
  return abicoder().encode(
    ["bytes32", "uint256", "bytes32", "uint256", "uint256", "bytes32"],
    [blockHash, height, root, size, quorum, domain]
  );
}

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

// A 48-byte hash-based (XMSS-style) attestation public key: root(32) || seed(16).
function xmssKey(label) {
  const root = ethers.keccak256(ethers.toUtf8Bytes("xmss-root:" + label));
  const seed = ethers.keccak256(ethers.toUtf8Bytes("xmss-seed:" + label));
  return ethers.concat([root, ethers.dataSlice(seed, 0, 16)]);
}

describe("AerePQFinalityCertificate (attestation-key registry + certificate verifier)", function () {
  let sp1, registry, verifier;
  let deployer, relayer, other;
  let validators; // 7 signers
  const VKEY = ethers.keccak256(ethers.toUtf8Bytes("sp1-program:aere-pq-finality-cert-v1"));

  const BLOCK_HASH = ethers.keccak256(ethers.toUtf8Bytes("aere-block:12345678"));
  const BLOCK_HEIGHT = 12345678n;

  beforeEach(async function () {
    const signers = await ethers.getSigners();
    deployer = signers[0];
    relayer = signers[8];
    other = signers[9];
    validators = signers.slice(1, 8); // 7 validators

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

    registry = await (await ethers.getContractFactory("AerePQAttestationKeyRegistry")).deploy();
    await registry.waitForDeployment();

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

    // Enroll the 7-validator set (owner), then each validator self-registers its PQ key.
    for (const v of validators) {
      await registry.connect(deployer).enrollValidator(v.address);
    }
    for (const v of validators) {
      await registry.connect(v).registerAttestationKey(SCHEME_XMSS, xmssKey(v.address));
    }
  });

  // Build public values bound to the CURRENT registry state (unless overridden), and
  // (optionally) register the proof valid in the mock gateway.
  async function makeCert({
    blockHash = BLOCK_HASH,
    height = BLOCK_HEIGHT,
    quorum,
    proof = "0xabcd01",
    overrideRoot = null,
    overrideSize = null,
    overrideDomain = null,
    register = true,
  }) {
    const root = overrideRoot !== null ? overrideRoot : await registry.validatorSetRoot();
    const size = overrideSize !== null ? overrideSize : await registry.validatorCount();
    const domain = overrideDomain !== null ? overrideDomain : await verifier.attestationDomain();
    const pv = encodePV({ blockHash, height, root, size, quorum, domain });
    if (register) {
      await sp1.setValid(proofMarker(VKEY, pv, proof), true);
    }
    return { pv, proof };
  }

  // =========================================================================
  // 1. Registry: 7-validator set with hash-based attestation keys
  // =========================================================================
  describe("attestation-key registry", function () {
    it("registers a 7-validator set with hash-based keys, quorum threshold is 5", async function () {
      expect(await registry.validatorCount()).to.equal(7n);
      expect(await registry.isSetComplete()).to.equal(true);
      expect(await registry.activeKeyCount()).to.equal(7n);
      expect(await verifier.currentQuorumThreshold()).to.equal(5n); // ceil(2*7/3) = 5

      for (const v of validators) {
        expect(await registry.isEnrolled(v.address)).to.equal(true);
        expect(await registry.keyEpochOf(v.address)).to.equal(1n);
        const [keyId, scheme, epoch] = await registry.currentKey(v.address);
        expect(scheme).to.equal(SCHEME_XMSS);
        expect(epoch).to.equal(1n);
        expect(keyId).to.not.equal(ethers.MaxUint256);
      }
      expect(await registry.validatorSetRoot()).to.not.equal(ethers.ZeroHash);
    });

    it("is append-only and fail-closed: only enrolled validators register; enroll is idempotent-guarded", async function () {
      // A non-enrolled address cannot register a key.
      await expect(
        registry.connect(other).registerAttestationKey(SCHEME_XMSS, xmssKey("intruder"))
      ).to.be.revertedWithCustomError(registry, "NotEnrolled");

      // Re-enrolling an existing validator reverts (no silent re-ordering).
      await expect(
        registry.connect(deployer).enrollValidator(validators[0].address)
      ).to.be.revertedWithCustomError(registry, "AlreadyEnrolled");

      // Only the owner may enroll.
      await expect(
        registry.connect(other).enrollValidator(other.address)
      ).to.be.revertedWith("Ownable: caller is not the owner");

      // Bad scheme / bad key length are rejected.
      await expect(
        registry.connect(validators[0]).registerAttestationKey(99, xmssKey("x"))
      ).to.be.revertedWithCustomError(registry, "InvalidScheme");
      await expect(
        registry.connect(validators[0]).registerAttestationKey(SCHEME_XMSS, "0x1234")
      ).to.be.revertedWithCustomError(registry, "InvalidPubKey");
    });
  });

  // =========================================================================
  // 2. A valid certificate with quorum 5 records and isFinalPQ returns true
  // =========================================================================
  describe("valid certificate records PQ finality", function () {
    it("quorum 5 verifies, records, and isFinalPQ(blockHash) is true", async function () {
      const { pv, proof } = await makeCert({ quorum: 5 });

      // Strict view path returns the finalized block.
      const [bh, height] = await verifier.verifyCertificate(pv, proof);
      expect(bh).to.equal(BLOCK_HASH);
      expect(height).to.equal(BLOCK_HEIGHT);

      // Non-reverting path is true.
      expect(await verifier.tryVerifyCertificate(pv, proof)).to.equal(true);

      // Consumer confirm-a-specific-block path.
      expect(await verifier.assertFinalizes(BLOCK_HASH, BLOCK_HEIGHT, pv, proof)).to.equal(true);

      // isFinalPQ is false BEFORE recording.
      expect(await verifier.isFinalPQ(BLOCK_HASH)).to.equal(false);

      // Recording emits and stores provenance.
      await expect(verifier.connect(relayer).recordCertificate(pv, proof))
        .to.emit(verifier, "CertificateRecorded")
        .withArgs(BLOCK_HASH, BLOCK_HEIGHT, 5, 7, await registry.validatorSetRoot(), relayer.address);

      expect(await verifier.isFinalPQ(BLOCK_HASH)).to.equal(true);
      expect(await verifier.isFinalPQAtHeight(BLOCK_HASH, BLOCK_HEIGHT)).to.equal(true);
      expect(await verifier.highestFinalizedHeight()).to.equal(BLOCK_HEIGHT);
      expect(await verifier.highestFinalizedBlockHash()).to.equal(BLOCK_HASH);
      expect(await verifier.certificateCount()).to.equal(1n);

      const cert = await verifier.certificateOf(BLOCK_HASH);
      expect(cert.exists).to.equal(true);
      expect(cert.quorumCount).to.equal(5);
      expect(cert.validatorSetSize).to.equal(7);
      expect(cert.prover).to.equal(relayer.address);
    });

    it("full quorum (7) also records", async function () {
      const { pv, proof } = await makeCert({ quorum: 7, proof: "0xf011" });
      await verifier.recordCertificate(pv, proof);
      expect(await verifier.isFinalPQ(BLOCK_HASH)).to.equal(true);
    });
  });

  // =========================================================================
  // 3. A certificate below quorum rejects
  // =========================================================================
  describe("below-quorum rejects fail-closed", function () {
    it("quorum 4 (< 5) is rejected even with a gateway-valid proof", async function () {
      const { pv, proof } = await makeCert({ quorum: 4 });

      await expect(verifier.verifyCertificate(pv, proof)).to.be.revertedWithCustomError(verifier, "BelowQuorum");
      expect(await verifier.tryVerifyCertificate(pv, proof)).to.equal(false);
      await expect(verifier.recordCertificate(pv, proof)).to.be.revertedWithCustomError(verifier, "BelowQuorum");
      expect(await verifier.isFinalPQ(BLOCK_HASH)).to.equal(false);
    });

    it("an invalid proof (not registered in the gateway) is rejected fail-closed", async function () {
      const { pv, proof } = await makeCert({ quorum: 5, register: false });
      await expect(verifier.verifyCertificate(pv, proof)).to.be.revertedWithCustomError(sp1, "MockInvalidProof");
      expect(await verifier.tryVerifyCertificate(pv, proof)).to.equal(false);
      expect(await verifier.isFinalPQ(BLOCK_HASH)).to.equal(false);
    });
  });

  // =========================================================================
  // 4. A wrong validator-set root rejects
  // =========================================================================
  describe("wrong validator-set root rejects", function () {
    it("a proof carrying a forged root is rejected before the gateway call", async function () {
      const forgedRoot = ethers.keccak256(ethers.toUtf8Bytes("attacker-validator-set"));
      const { pv, proof } = await makeCert({ quorum: 5, overrideRoot: forgedRoot });

      await expect(verifier.verifyCertificate(pv, proof)).to.be.revertedWithCustomError(
        verifier,
        "ValidatorSetRootMismatch"
      );
      expect(await verifier.tryVerifyCertificate(pv, proof)).to.equal(false);
    });

    it("a proof carrying the wrong set size or domain is rejected", async function () {
      const wrongSize = await makeCert({ quorum: 5, overrideSize: 9, proof: "0x51" });
      await expect(verifier.verifyCertificate(wrongSize.pv, wrongSize.proof)).to.be.revertedWithCustomError(
        verifier,
        "SizeMismatch"
      );

      const wrongDomain = await makeCert({
        quorum: 5,
        overrideDomain: ethers.keccak256(ethers.toUtf8Bytes("wrong-domain")),
        proof: "0x52",
      });
      await expect(verifier.verifyCertificate(wrongDomain.pv, wrongDomain.proof)).to.be.revertedWithCustomError(
        verifier,
        "DomainMismatch"
      );

      // Malformed (not 192 bytes) public values reject.
      await expect(verifier.verifyCertificate("0x1234", "0x53")).to.be.revertedWithCustomError(
        verifier,
        "BadPublicValues"
      );
    });
  });

  // =========================================================================
  // 5. A wrong block rejects
  // =========================================================================
  describe("wrong block rejects", function () {
    it("assertFinalizes rejects a certificate substituted for a different block", async function () {
      const { pv, proof } = await makeCert({ quorum: 5 });
      const otherBlock = ethers.keccak256(ethers.toUtf8Bytes("aere-block:99999999"));

      // The proof genuinely finalizes BLOCK_HASH...
      expect(await verifier.assertFinalizes(BLOCK_HASH, BLOCK_HEIGHT, pv, proof)).to.equal(true);
      // ...but cannot be claimed for a different block.
      await expect(
        verifier.assertFinalizes(otherBlock, BLOCK_HEIGHT, pv, proof)
      ).to.be.revertedWithCustomError(verifier, "BlockHashMismatch");
      // ...nor a different height.
      await expect(
        verifier.assertFinalizes(BLOCK_HASH, BLOCK_HEIGHT + 1n, pv, proof)
      ).to.be.revertedWithCustomError(verifier, "BlockHeightMismatch");
      expect(await verifier.tryAssertFinalizes(otherBlock, BLOCK_HEIGHT, pv, proof)).to.equal(false);
    });

    it("a zero block hash is rejected", async function () {
      const { pv, proof } = await makeCert({ blockHash: ethers.ZeroHash, quorum: 5, proof: "0x2000" });
      await expect(verifier.verifyCertificate(pv, proof)).to.be.revertedWithCustomError(verifier, "ZeroBlockHash");
    });
  });

  // =========================================================================
  // 6. Key rotation updates the set root; an old-root certificate no longer matches
  // =========================================================================
  describe("key rotation invalidates old-root certificates", function () {
    it("rotating one validator's key changes the root and rejects the pre-rotation certificate", async function () {
      const oldRoot = await registry.validatorSetRoot();

      // A certificate valid against the CURRENT (pre-rotation) root.
      const { pv: oldPV, proof: oldProof } = await makeCert({ quorum: 5, proof: "0x01d0" });
      // It verifies now.
      expect(await verifier.tryVerifyCertificate(oldPV, oldProof)).to.equal(true);

      // Validator 0 rotates its one-time/few-time key (epoch 1 -> 2), scheme change allowed.
      await registry.connect(validators[0]).registerAttestationKey(SCHEME_LEANSIG, xmssKey("rotated-key-0"));
      expect(await registry.keyEpochOf(validators[0].address)).to.equal(2n);

      const newRoot = await registry.validatorSetRoot();
      expect(newRoot).to.not.equal(oldRoot);

      // The pre-rotation certificate (bound to oldRoot) no longer matches the live registry.
      await expect(verifier.verifyCertificate(oldPV, oldProof)).to.be.revertedWithCustomError(
        verifier,
        "ValidatorSetRootMismatch"
      );
      expect(await verifier.tryVerifyCertificate(oldPV, oldProof)).to.equal(false);

      // A fresh certificate bound to the NEW root verifies and records.
      const { pv: newPV, proof: newProof } = await makeCert({ quorum: 5, proof: "0x0ec0" });
      await verifier.recordCertificate(newPV, newProof);
      expect(await verifier.isFinalPQ(BLOCK_HASH)).to.equal(true);

      // History is append-only: validator 0 has two keys on record.
      const history = await registry.keyHistoryOf(validators[0].address);
      expect(history.length).to.equal(2);
      const first = await registry.getKey(history[0]);
      expect(first.superseded).to.equal(true);
      const second = await registry.getKey(history[1]);
      expect(second.superseded).to.equal(false);
      expect(second.scheme).to.equal(SCHEME_LEANSIG);
    });
  });

  // =========================================================================
  // 7. Empty validator set fails closed
  // =========================================================================
  describe("empty validator set fails closed", function () {
    it("a verifier over an empty registry rejects any certificate", async function () {
      const emptyRegistry = await (await ethers.getContractFactory("AerePQAttestationKeyRegistry")).deploy();
      await emptyRegistry.waitForDeployment();
      const emptyVerifier = await (await ethers.getContractFactory("AereFinalityCertificateVerifier")).deploy(
        await sp1.getAddress(),
        VKEY,
        await emptyRegistry.getAddress()
      );
      await emptyVerifier.waitForDeployment();

      const domain = await emptyVerifier.attestationDomain();
      const pv = encodePV({
        blockHash: BLOCK_HASH,
        height: BLOCK_HEIGHT,
        root: await emptyRegistry.validatorSetRoot(),
        size: 0,
        quorum: 0,
        domain,
      });
      await sp1.setValid(proofMarker(VKEY, pv, "0xee"), true);
      await expect(emptyVerifier.verifyCertificate(pv, "0xee")).to.be.revertedWithCustomError(
        emptyVerifier,
        "EmptyValidatorSet"
      );
    });
  });
});