const { expect } = require("chai");
const { ethers, network } = require("hardhat");
const kat = require("./falcon512_kat0.json");

// ---------------------------------------------------------------------------
// AerePQCMessageVerifier: on-chain PQC message verification for cross-chain
// bridging / oracle attestations, on the LIVE mainnet precompiles.
//
// Two precompile backends are installed at the precompile addresses via
// hardhat_setCode (the Hardhat EVM has no PQC precompile):
//
//   * Falcon512PrecompileShim at 0x0AE1 - delegates to AereFalcon512Verifier, the
//     KAT-proven pure-Solidity Falcon-512 verifier. Used for the GENUINE accept +
//     tamper-reject over the OFFICIAL NIST Falcon-512 KAT vector, so a positive here
//     is a real Falcon-512 signature verified end to end through the contract's real
//     precompile-calling path (input = pk || sm, raw 32-byte return).
//
//   * MockPQCPrecompile at 0x0AE1 / 0x0AE3 - a faithful mock (parses the precompiles'
//     SPEC offsets, accepts iff commitment == keccak256(pk || message)). Used for the
//     signer-bound / domain-separation / rotation / consumer tests, where genuine
//     signatures must be minted for arbitrary domain-separated digests that no offline
//     signer is available to sign. The real wire encoding is proven by the KAT path
//     above and (separately) against the live mainnet precompile via eth_call.
// ---------------------------------------------------------------------------

const AE1 = ethers.getAddress("0x0000000000000000000000000000000000000ae1"); // Falcon-512
const AE3 = ethers.getAddress("0x0000000000000000000000000000000000000ae3"); // ML-DSA-44

const SCHEME_FALCON512 = 1;
const SCHEME_MLDSA44 = 3;

const MESSAGE_DOMAIN = ethers.keccak256(ethers.toUtf8Bytes("AerePQCMessageVerifier.v1.message"));
const ACTION_DOMAIN = ethers.keccak256(ethers.toUtf8Bytes("AerePQCGatedAction.v1.action"));

// Deterministic seeded RNG (keccak hash-chain) so the tests are fully reproducible.
function makeRng(seedText) {
  let seed = ethers.keccak256(ethers.toUtf8Bytes(seedText));
  return {
    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;
    },
  };
}

function falconPubKey(rng) {
  const raw = rng.bytes(897);
  raw[0] = 0x09; // Falcon-512 header
  return ethers.hexlify(raw);
}

function mldsaPubKey(rng) {
  return ethers.hexlify(rng.bytes(1312));
}

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

// Genuine Falcon-512 envelope for the mock: nonce(40) || esig ; esig = 0x29 || commitment(32).
function falconMockEnvelope(pubKey, message, rng) {
  const commitment = commitmentFor(pubKey, message);
  const nonce = rng.bytes(40);
  const esig = ethers.concat([new Uint8Array([0x29]), commitment]);
  return ethers.hexlify(ethers.concat([nonce, esig]));
}

// Genuine ML-DSA-44 envelope for the mock: 2420-byte sig with commitment in the first 32 bytes.
function mldsaMockEnvelope(pubKey, message) {
  const sig = new Uint8Array(2420);
  sig.set(ethers.getBytes(commitmentFor(pubKey, message)), 0);
  return ethers.hexlify(sig);
}

// Reproduce signerMessageDigest for an ARBITRARY chainId (to test cross-chain replay).
function foreignSignerDigest(chainId, mvAddr, signerId, epoch, domainTag, payload) {
  const enc = ethers.AbiCoder.defaultAbiCoder().encode(
    ["bytes32", "uint256", "address", "uint256", "uint32", "bytes32", "bytes"],
    [MESSAGE_DOMAIN, chainId, mvAddr, signerId, epoch, domainTag, payload]
  );
  return ethers.keccak256(enc);
}

async function setCodeFrom(factoryName, addr, ...ctorArgs) {
  const F = await ethers.getContractFactory(factoryName);
  const c = await F.deploy(...ctorArgs);
  await c.waitForDeployment();
  const code = await ethers.provider.getCode(await c.getAddress());
  await network.provider.send("hardhat_setCode", [addr, code]);
}

async function deployVerifier() {
  const F = await ethers.getContractFactory("AerePQCMessageVerifier");
  const mv = await F.deploy();
  await mv.waitForDeployment();
  return mv;
}

async function registerSigner(mv, scheme, pubKey) {
  const tx = await mv.registerSigner(scheme, pubKey);
  const rc = await tx.wait();
  const log = rc.logs.map((l) => mv.interface.parseLog(l)).find((p) => p && p.name === "SignerRegistered");
  return Number(log.args.signerId);
}

describe("AerePQCMessageVerifier", function () {
  // =========================================================================
  // 1. Signer registry: register / rotate / revoke / transfer
  // =========================================================================
  describe("signer registry", function () {
    let mv;
    beforeEach(async function () {
      mv = await deployVerifier();
    });

    it("registers well-formed keys with sequential ids, epoch 0, ACTIVE", async function () {
      const rng = makeRng("reg");
      const [a] = await ethers.getSigners();

      const id0 = await registerSigner(mv, SCHEME_FALCON512, falconPubKey(rng));
      const id1 = await registerSigner(mv, SCHEME_MLDSA44, mldsaPubKey(rng));
      expect(id0).to.equal(0);
      expect(id1).to.equal(1);
      expect(await mv.signerCount()).to.equal(2n);

      const s0 = await mv.getSigner(0);
      expect(s0.owner).to.equal(a.address);
      expect(Number(s0.scheme)).to.equal(SCHEME_FALCON512);
      expect(Number(s0.status)).to.equal(1); // ACTIVE
      expect(Number(s0.epoch)).to.equal(0);
      expect(await mv.isActiveSigner(0)).to.equal(true);
      expect(Number(await mv.epochOf(1))).to.equal(0);
    });

    it("rejects an unknown scheme and malformed keys", async function () {
      await expect(mv.registerSigner(5, "0x" + "00".repeat(897))).to.be.revertedWithCustomError(mv, "InvalidScheme");
      // Falcon-512 wrong length
      await expect(mv.registerSigner(1, "0x09" + "00".repeat(895))).to.be.revertedWithCustomError(mv, "InvalidPubKey");
      // Falcon-512 correct length, wrong header
      await expect(mv.registerSigner(1, "0x0a" + "00".repeat(896))).to.be.revertedWithCustomError(mv, "InvalidPubKey");
      // ML-DSA-44 wrong length
      await expect(mv.registerSigner(3, "0x" + "00".repeat(1311))).to.be.revertedWithCustomError(mv, "InvalidPubKey");
    });

    it("rotates in place: owner-only, epoch advances, key/scheme update", async function () {
      const rng = makeRng("rot");
      const [owner, other] = await ethers.getSigners();
      const id = await registerSigner(mv, SCHEME_FALCON512, falconPubKey(rng));

      // A non-owner cannot rotate.
      await expect(
        mv.connect(other).rotateSigner(id, SCHEME_FALCON512, falconPubKey(rng))
      ).to.be.revertedWithCustomError(mv, "NotSignerOwner");

      const newKey = mldsaPubKey(rng); // rotate to a different scheme too
      await expect(mv.connect(owner).rotateSigner(id, SCHEME_MLDSA44, newKey)).to.emit(mv, "SignerRotated");
      const s = await mv.getSigner(id);
      expect(Number(s.epoch)).to.equal(1);
      expect(Number(s.scheme)).to.equal(SCHEME_MLDSA44);
      expect(s.pubKey).to.equal(newKey);
    });

    it("revokes (owner-only, terminal) and blocks rotation afterwards", async function () {
      const rng = makeRng("rev");
      const [owner, other] = await ethers.getSigners();
      const id = await registerSigner(mv, SCHEME_FALCON512, falconPubKey(rng));

      await expect(mv.connect(other).revokeSigner(id)).to.be.revertedWithCustomError(mv, "NotSignerOwner");
      await mv.connect(owner).revokeSigner(id);
      expect(await mv.isActiveSigner(id)).to.equal(false);
      expect(Number(await mv.statusOf(id))).to.equal(2); // REVOKED

      await expect(mv.connect(owner).rotateSigner(id, SCHEME_FALCON512, falconPubKey(rng)))
        .to.be.revertedWithCustomError(mv, "SignerNotActive");
    });

    it("transfers signer ownership", async function () {
      const rng = makeRng("xfer");
      const [owner, next] = await ethers.getSigners();
      const id = await registerSigner(mv, SCHEME_FALCON512, falconPubKey(rng));

      await expect(mv.connect(owner).transferSignerOwner(id, ethers.ZeroAddress))
        .to.be.revertedWithCustomError(mv, "ZeroAddress");
      await mv.connect(owner).transferSignerOwner(id, next.address);
      expect((await mv.getSigner(id)).owner).to.equal(next.address);
      // Old owner can no longer rotate; new owner can.
      await expect(mv.connect(owner).rotateSigner(id, SCHEME_FALCON512, falconPubKey(rng)))
        .to.be.revertedWithCustomError(mv, "NotSignerOwner");
      await expect(mv.connect(next).rotateSigner(id, SCHEME_FALCON512, falconPubKey(rng))).to.emit(mv, "SignerRotated");
    });

    it("reverts on unknown signerId views", async function () {
      await expect(mv.getSigner(0)).to.be.revertedWithCustomError(mv, "UnknownSigner");
      expect(Number(await mv.statusOf(0))).to.equal(0); // STATUS_NONE, does not revert
      expect(await mv.isActiveSigner(0)).to.equal(false);
    });
  });

  // =========================================================================
  // 2. GENUINE Falcon-512 accept via the KAT-proven verifier path (real crypto)
  // =========================================================================
  describe("genuine Falcon-512 accept (KAT-proven verifier path)", function () {
    let mv;

    before(async function () {
      // Deploy the real KAT-proven verifier, then install a shim over it at 0x0AE1.
      const V = await ethers.getContractFactory("AereFalcon512Verifier");
      const verifier = await V.deploy();
      await verifier.waitForDeployment();
      await setCodeFrom("Falcon512PrecompileShim", AE1, await verifier.getAddress());
    });

    beforeEach(async function () {
      mv = await deployVerifier();
    });

    // Envelope handed to verifyMessage: nonce(40) || esig ; esig = 0x29 || compsig.
    function katEnvelope() {
      return ethers.hexlify(ethers.concat([kat.nonce, new Uint8Array([0x29]), kat.compsig]));
    }

    it("verifyMessage ACCEPTS the official NIST Falcon-512 KAT signature", async function () {
      const ok = await mv.verifyMessage(SCHEME_FALCON512, kat.pk, kat.msg, katEnvelope());
      expect(ok).to.equal(true);
    });

    it("REJECTS a tampered signature (flip one byte of the compressed sig)", async function () {
      const env = ethers.getBytes(katEnvelope());
      env[100] ^= 0x01; // deep inside the compressed signature
      const ok = await mv.verifyMessage(SCHEME_FALCON512, kat.pk, kat.msg, ethers.hexlify(env));
      expect(ok).to.equal(false);
    });

    it("REJECTS a tampered message (flip one byte of the message)", async function () {
      const msg = ethers.getBytes(kat.msg);
      msg[0] ^= 0x01;
      const ok = await mv.verifyMessage(SCHEME_FALCON512, kat.pk, ethers.hexlify(msg), katEnvelope());
      expect(ok).to.equal(false);
    });

    it("REJECTS the genuine signature under the WRONG public key", async function () {
      const ok = await mv.verifyMessage(SCHEME_FALCON512, kat.pk1, kat.msg, katEnvelope());
      expect(ok).to.equal(false);
    });
  });

  // =========================================================================
  // 3. Domain separation, signer-bound verify, rotation, consumer gate (mock)
  // =========================================================================
  describe("signer-bound verification, domain separation, rotation, consumer", function () {
    let mv, falconMockCode, mldsaMockCode;

    before(async function () {
      const M = await ethers.getContractFactory("MockPQCPrecompile");
      const f = await M.deploy(1);
      await f.waitForDeployment();
      falconMockCode = await ethers.provider.getCode(await f.getAddress());
      const d = await M.deploy(3);
      await d.waitForDeployment();
      mldsaMockCode = await ethers.provider.getCode(await d.getAddress());
    });

    beforeEach(async function () {
      await network.provider.send("hardhat_setCode", [AE1, falconMockCode]);
      await network.provider.send("hardhat_setCode", [AE3, mldsaMockCode]);
      mv = await deployVerifier();
    });

    it("verifyBySigner ACCEPTS a genuine signer signature (Falcon-512 and ML-DSA-44)", async function () {
      const rng = makeRng("bysigner");
      const domain = ethers.keccak256(ethers.toUtf8Bytes("bridge-route-7"));
      const payload = ethers.toUtf8Bytes("cross-chain-packet-A");

      // Falcon-512 signer
      const fpk = falconPubKey(rng);
      const fid = await registerSigner(mv, SCHEME_FALCON512, fpk);
      const fdigest = await mv.signerMessageDigest(fid, domain, payload);
      const fsig = falconMockEnvelope(fpk, fdigest, rng);
      expect(await mv.verifyBySigner(fid, domain, payload, fsig)).to.equal(true);

      // ML-DSA-44 signer
      const dpk = mldsaPubKey(rng);
      const did = await registerSigner(mv, SCHEME_MLDSA44, dpk);
      const ddigest = await mv.signerMessageDigest(did, domain, payload);
      const dsig = mldsaMockEnvelope(dpk, ddigest);
      expect(await mv.verifyBySigner(did, domain, payload, dsig)).to.equal(true);
    });

    it("verifyBySigner returns false for unknown / revoked signer", async function () {
      const rng = makeRng("revoked");
      const domain = ethers.keccak256(ethers.toUtf8Bytes("d"));
      const payload = ethers.toUtf8Bytes("p");
      expect(await mv.verifyBySigner(99, domain, payload, "0x" + "00".repeat(73))).to.equal(false);

      const pk = falconPubKey(rng);
      const id = await registerSigner(mv, SCHEME_FALCON512, pk);
      const digest = await mv.signerMessageDigest(id, domain, payload);
      const sig = falconMockEnvelope(pk, digest, rng);
      expect(await mv.verifyBySigner(id, domain, payload, sig)).to.equal(true);
      await mv.revokeSigner(id);
      expect(await mv.verifyBySigner(id, domain, payload, sig)).to.equal(false);
    });

    it("DOMAIN SEPARATION: a signature for one domain tag is invalid under another", async function () {
      const rng = makeRng("domainsep");
      const domainA = ethers.keccak256(ethers.toUtf8Bytes("oracle-feed-ETHUSD"));
      const domainB = ethers.keccak256(ethers.toUtf8Bytes("oracle-feed-BTCUSD"));
      const payload = ethers.toUtf8Bytes("price=123456");

      const pk = falconPubKey(rng);
      const id = await registerSigner(mv, SCHEME_FALCON512, pk);
      const digestA = await mv.signerMessageDigest(id, domainA, payload);
      const sigA = falconMockEnvelope(pk, digestA, rng);

      expect(await mv.verifyBySigner(id, domainA, payload, sigA)).to.equal(true);
      // Same key, same payload, DIFFERENT domain -> different digest -> rejected.
      expect(await mv.verifyBySigner(id, domainB, payload, sigA)).to.equal(false);
    });

    it("CROSS-CHAIN SEPARATION: a signature bound to a foreign chainId is invalid here", async function () {
      const rng = makeRng("crosschain");
      const domain = ethers.keccak256(ethers.toUtf8Bytes("bridge"));
      const payload = ethers.toUtf8Bytes("withdraw-1000");
      const pk = falconPubKey(rng);
      const id = await registerSigner(mv, SCHEME_FALCON512, pk);
      const mvAddr = await mv.getAddress();

      const localChainId = (await ethers.provider.getNetwork()).chainId; // 31337
      const foreignChainId = 1n; // Ethereum mainnet

      // A signature genuine for the SAME (domain, payload) but bound to chainId=1.
      const foreignDigest = foreignSignerDigest(foreignChainId, mvAddr, id, 0, domain, payload);
      const foreignSig = falconMockEnvelope(pk, foreignDigest, rng);
      // On THIS chain (31337) the digest differs -> the foreign-bound signature is rejected.
      expect(await mv.verifyBySigner(id, domain, payload, foreignSig)).to.equal(false);

      // Sanity: the same envelope built for the local chainId verifies.
      const localDigest = foreignSignerDigest(localChainId, mvAddr, id, 0, domain, payload);
      expect(localDigest).to.equal(await mv.signerMessageDigest(id, domain, payload));
      const localSig = falconMockEnvelope(pk, localDigest, rng);
      expect(await mv.verifyBySigner(id, domain, payload, localSig)).to.equal(true);
    });

    it("KEY ROTATION: a pre-rotation signature is rejected; a post-rotation one is accepted", async function () {
      const rng = makeRng("rotate-verify");
      const domain = ethers.keccak256(ethers.toUtf8Bytes("route"));
      const payload = ethers.toUtf8Bytes("msg");

      const pk0 = falconPubKey(rng);
      const id = await registerSigner(mv, SCHEME_FALCON512, pk0);

      // Genuine under epoch 0 / key0.
      const digest0 = await mv.signerMessageDigest(id, domain, payload);
      const sig0 = falconMockEnvelope(pk0, digest0, rng);
      expect(await mv.verifyBySigner(id, domain, payload, sig0)).to.equal(true);

      // Rotate to key1 (epoch 1). The old signature must no longer verify: the digest now binds
      // epoch 1 AND the current key is key1, so both the message and the key have changed.
      const pk1 = falconPubKey(rng);
      await mv.rotateSigner(id, SCHEME_FALCON512, pk1);
      expect(Number(await mv.epochOf(id))).to.equal(1);
      expect(await mv.verifyBySigner(id, domain, payload, sig0)).to.equal(false);

      // A fresh signature under key1 / epoch 1 verifies.
      const digest1 = await mv.signerMessageDigest(id, domain, payload);
      expect(digest1).to.not.equal(digest0);
      const sig1 = falconMockEnvelope(pk1, digest1, rng);
      expect(await mv.verifyBySigner(id, domain, payload, sig1)).to.equal(true);
    });

    it("CONSUMER GATE: executes once on a valid PQC signature, rejects tampered and replays", async function () {
      const rng = makeRng("consumer");
      const G = await ethers.getContractFactory("AerePQCGatedAction");
      const gate = await G.deploy(await mv.getAddress());
      await gate.waitForDeployment();

      const pk = falconPubKey(rng);
      const id = await registerSigner(mv, SCHEME_FALCON512, pk);
      const payload = ethers.toUtf8Bytes("bridge-deliver:0xdeadbeef");

      // Sign the consumer's ACTION_DOMAIN + payload digest.
      const digest = await mv.signerMessageDigest(id, ACTION_DOMAIN, payload);
      const sig = falconMockEnvelope(pk, digest, rng);

      // Tampered signature -> gate reverts, nothing recorded.
      const bad = ethers.getBytes(sig);
      bad[41] ^= 0x01; // flip inside the commitment
      await expect(gate.executeIfAttested(id, payload, ethers.hexlify(bad)))
        .to.be.revertedWithCustomError(gate, "NotAttested");
      expect(await gate.actionCount()).to.equal(0n);

      // Genuine signature -> action executes exactly once.
      await expect(gate.executeIfAttested(id, payload, sig)).to.emit(gate, "ActionExecuted");
      expect(await gate.actionCount()).to.equal(1n);
      const actionId = await gate.actionId(id, payload);
      expect(await gate.consumed(actionId)).to.equal(true);

      // Replay of the same (signer, payload) -> rejected.
      await expect(gate.executeIfAttested(id, payload, sig))
        .to.be.revertedWithCustomError(gate, "AlreadyConsumed");
      expect(await gate.actionCount()).to.equal(1n);
    });
  });

  // =========================================================================
  // 4. Fail-closed: empty/stale precompile return is INVALID
  // =========================================================================
  describe("fail-closed", function () {
    it("treats an empty precompile return (pre-fork / stale) as invalid", async function () {
      const rng = makeRng("failclosed");
      // Install the faithful Falcon mock, register + build a genuine sig, then wipe the precompile.
      const M = await ethers.getContractFactory("MockPQCPrecompile");
      const f = await M.deploy(1);
      await f.waitForDeployment();
      await network.provider.send("hardhat_setCode", [AE1, await ethers.provider.getCode(await f.getAddress())]);

      const mv = await deployVerifier();
      const domain = ethers.keccak256(ethers.toUtf8Bytes("d"));
      const payload = ethers.toUtf8Bytes("p");
      const pk = falconPubKey(rng);
      const id = await registerSigner(mv, SCHEME_FALCON512, pk);
      const digest = await mv.signerMessageDigest(id, domain, payload);
      const sig = falconMockEnvelope(pk, digest, rng);
      expect(await mv.verifyBySigner(id, domain, payload, sig)).to.equal(true);

      // An address with no precompile returns success + empty returndata -> must be INVALID.
      await network.provider.send("hardhat_setCode", [AE1, "0x"]);
      expect(await mv.verifyBySigner(id, domain, payload, sig)).to.equal(false);
      expect(await mv.verifyMessage(SCHEME_FALCON512, pk, digest, sig)).to.equal(false);
    });
  });
});