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

// ---------------------------------------------------------------------------
// AerePQCAttestation adversarial test + seeded invariant fuzz.
//
// Hardhat has no PQC precompile, so we install MockPQCPrecompile at 0x0AE1..0x0AE4
// (hardhat_setCode). The mock parses the input at the precompiles' SPEC offsets and
// accepts iff commitment == keccak256(pk || message). A "genuine signature" is one
// whose commitment equals keccak256(pk || challenge); a forgery/replay/wrong-key/
// wrong-message signature has a mismatching commitment and MUST be rejected.
//
// The REAL precompile encoding is separately proven 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 fuzz is fully 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;
  }
  function int(n) {
    const b = bytes(4);
    return (((b[0] << 24) | (b[1] << 16) | (b[2] << 8) | b[3]) >>> 0) % n;
  }
  return { bytes, int };
}

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 the attester signature envelope with 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 envelope binding to `message` under `pubKey`.
function genuine(scheme, pubKey, message, rng) {
  return envelope(scheme, commitmentFor(pubKey, message), rng);
}

// Flip one byte inside the commitment region (Falcon: esig[1] at offset 41;
// fixed schemes: sig[0] at offset 0). This is the region the precompile binds
// the (key,message) pair to, so a flip here is a genuine forgery.
function tamperCommitment(scheme, sigHex) {
  const b = ethers.getBytes(sigHex);
  const idx = META[scheme].falcon ? 41 : 0;
  b[idx] ^= 0x01;
  return ethers.hexlify(b);
}

describe("AerePQCAttestation (adversarial + seeded fuzz)", function () {
  let att, 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 A = await ethers.getContractFactory("AerePQCAttestation");
    att = await A.deploy();
    await att.waitForDeployment();
  });

  // Helper: register a key, return its keyId (from the event).
  async function register(scheme, pubKey) {
    const tx = await att.registerKey(scheme, pubKey);
    const rc = await tx.wait();
    const log = rc.logs.map((l) => att.interface.parseLog(l)).find((p) => p && p.name === "KeyRegistered");
    return Number(log.args.keyId);
  }

  // -----------------------------------------------------------------------
  // Registration validation
  // -----------------------------------------------------------------------
  describe("registerKey", function () {
    it("accepts a well-formed key for every scheme and assigns sequential ids", async function () {
      const rng = makeRng("reg");
      let expectedId = 0;
      for (const scheme of [1, 2, 3, 4]) {
        const pk = makePubKey(scheme, rng);
        const id = await register(scheme, pk);
        expect(id).to.equal(expectedId++);
        const k = await att.getKey(id);
        expect(Number(k.scheme)).to.equal(scheme);
        expect(k.pubKey).to.equal(pk);
        expect(Number(k.nonce)).to.equal(0);
      }
      expect(await att.keyCount()).to.equal(4n);
    });

    it("rejects an unknown scheme", async function () {
      await expect(att.registerKey(5, "0x" + "00".repeat(897))).to.be.revertedWithCustomError(att, "InvalidScheme");
      await expect(att.registerKey(0, "0x" + "00".repeat(897))).to.be.revertedWithCustomError(att, "InvalidScheme");
    });

    it("rejects a wrong-length or wrong-header public key", async function () {
      // Falcon-1024 wrong length
      await expect(att.registerKey(2, "0x" + "0a" + "00".repeat(1791))).to.be.revertedWithCustomError(att, "InvalidPubKey");
      // Falcon-1024 correct length but wrong header (0x09 instead of 0x0A)
      await expect(att.registerKey(2, "0x" + "09" + "00".repeat(1792))).to.be.revertedWithCustomError(att, "InvalidPubKey");
      // Falcon-512 correct length but wrong header
      await expect(att.registerKey(1, "0x" + "0a" + "00".repeat(896))).to.be.revertedWithCustomError(att, "InvalidPubKey");
      // ML-DSA-44 wrong length
      await expect(att.registerKey(3, "0x" + "00".repeat(1311))).to.be.revertedWithCustomError(att, "InvalidPubKey");
      // SLH-DSA wrong length
      await expect(att.registerKey(4, "0x" + "00".repeat(31))).to.be.revertedWithCustomError(att, "InvalidPubKey");
    });
  });

  // -----------------------------------------------------------------------
  // Happy path per scheme
  // -----------------------------------------------------------------------
  describe("attest happy path", function () {
    for (const scheme of [1, 2, 3, 4]) {
      it(`records a genuine attestation for scheme ${scheme}`, async function () {
        const rng = makeRng("happy-" + scheme);
        const pk = makePubKey(scheme, rng);
        const keyId = await register(scheme, pk);
        const messageHash = ethers.keccak256(ethers.toUtf8Bytes("payload-" + scheme));

        const challenge = await att.attestChallenge(keyId, 0, messageHash);
        const sig = genuine(scheme, pk, challenge, rng);

        const before = await att.attestationCount();
        await expect(att.attest(keyId, messageHash, sig)).to.emit(att, "Attested");
        expect(await att.attestationCount()).to.equal(before + 1n);
        expect(Number(await att.nonceOf(keyId))).to.equal(1);

        const a = await att.getAttestation(keyId, 0);
        expect(a.exists).to.equal(true);
        expect(a.messageHash).to.equal(messageHash);
        expect(a.challenge).to.equal(challenge);
        expect(await att.isValidAttestation(keyId, 0, messageHash)).to.equal(true);
        expect(await att.isValidAttestation(keyId, 0, ethers.ZeroHash)).to.equal(false);
      });
    }

    it("verifySignature view returns true for genuine, false for tampered (raw messageHash binding)", async function () {
      const rng = makeRng("view");
      const scheme = 2;
      const pk = makePubKey(scheme, rng);
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("direct-verify"));
      const good = genuine(scheme, pk, messageHash, rng);
      expect(await att.verifySignature(scheme, pk, messageHash, good)).to.equal(true);

      const bad = ethers.getBytes(good);
      bad[bad.length - 1] ^= 0xff;
      expect(await att.verifySignature(scheme, pk, messageHash, ethers.hexlify(bad))).to.equal(false);
    });
  });

  // -----------------------------------------------------------------------
  // Adversarial: forgery, wrong key, wrong message, replay
  // -----------------------------------------------------------------------
  describe("adversarial rejection (nothing recorded)", function () {
    async function setup(scheme) {
      const rng = makeRng("adv-" + scheme);
      const pk = makePubKey(scheme, rng);
      const keyId = await register(scheme, pk);
      const messageHash = ethers.keccak256(ethers.toUtf8Bytes("adv-msg"));
      const challenge = await att.attestChallenge(keyId, 0, messageHash);
      return { rng, pk, keyId, messageHash, challenge };
    }

    async function expectRejected(promise) {
      await expect(promise).to.be.revertedWithCustomError(att, "PQCVerificationFailed");
    }

    for (const scheme of [1, 2, 3, 4]) {
      it(`rejects a tampered signature for scheme ${scheme}`, async function () {
        const { rng, pk, keyId, messageHash, challenge } = await setup(scheme);
        const good = tamperCommitment(scheme, genuine(scheme, pk, challenge, rng));
        await expectRejected(att.attest(keyId, messageHash, good));
        expect(await att.attestationCount()).to.equal(0n);
        expect(Number(await att.nonceOf(keyId))).to.equal(0);
        expect((await att.getAttestation(keyId, 0)).exists).to.equal(false);
      });

      it(`rejects a signature over a DIFFERENT message for scheme ${scheme}`, async function () {
        const { rng, pk, keyId, messageHash } = await setup(scheme);
        const otherMessage = ethers.keccak256(ethers.toUtf8Bytes("other-msg"));
        const otherChallenge = await att.attestChallenge(keyId, 0, otherMessage);
        const sigForOther = genuine(scheme, pk, otherChallenge, rng);
        // signature is genuine for otherMessage, but we attest messageHash -> reject
        await expectRejected(att.attest(keyId, messageHash, sigForOther));
        expect(await att.attestationCount()).to.equal(0n);
      });

      it(`rejects a signature by a DIFFERENT key for scheme ${scheme}`, async function () {
        const { rng, keyId, messageHash, challenge } = await setup(scheme);
        const otherPk = makePubKey(scheme, rng);
        const sigOtherKey = genuine(scheme, otherPk, challenge, rng); // commitment uses other pk
        await expectRejected(att.attest(keyId, messageHash, sigOtherKey));
        expect(await att.attestationCount()).to.equal(0n);
      });
    }

    it("rejects replay of a consumed (message,signature) after the nonce advances", async function () {
      const scheme = 2;
      const { rng, pk, keyId, messageHash } = await setup(scheme);
      const challenge0 = await att.attestChallenge(keyId, 0, messageHash);
      const sig0 = genuine(scheme, pk, challenge0, rng);

      await att.attest(keyId, messageHash, sig0); // consumes nonce 0
      expect(Number(await att.nonceOf(keyId))).to.equal(1);

      // Exact replay of the same (message, signature) now fails: challenge uses nonce 1.
      await expect(att.attest(keyId, messageHash, sig0)).to.be.revertedWithCustomError(att, "PQCVerificationFailed");
      expect(await att.attestationCount()).to.equal(1n);
      expect(Number(await att.nonceOf(keyId))).to.equal(1);
    });

    it("rejects a stale-nonce signature (signed for an old nonce)", async function () {
      const scheme = 2;
      const { rng, pk, keyId, messageHash } = await setup(scheme);
      // advance nonce to 1 with a genuine attestation
      const c0 = await att.attestChallenge(keyId, 0, messageHash);
      await att.attest(keyId, messageHash, genuine(scheme, pk, c0, rng));
      // now craft a signature for the OLD nonce 0 and try again -> reject
      const staleSig = genuine(scheme, pk, c0, rng);
      await expect(att.attest(keyId, messageHash, staleSig)).to.be.revertedWithCustomError(att, "PQCVerificationFailed");
    });

    it("reverts on unknown keyId", async function () {
      await expect(att.attest(99, ethers.ZeroHash, "0x" + "00".repeat(80))).to.be.revertedWithCustomError(att, "UnknownKey");
    });

    it("treats an empty precompile return (pre-fork / stale) as invalid", async function () {
      const scheme = 2;
      const { rng, pk, keyId, messageHash, challenge } = await setup(scheme);
      const sig = genuine(scheme, pk, challenge, rng);
      // Wipe the precompile: an empty account returns success + empty returndata.
      await network.provider.send("hardhat_setCode", [PRECOMPILE[scheme], "0x"]);
      expect(await att.verifySignature(scheme, pk, challenge, sig)).to.equal(false);
      await expect(att.attest(keyId, messageHash, sig)).to.be.revertedWithCustomError(att, "PQCVerificationFailed");
      expect(await att.attestationCount()).to.equal(0n);
    });
  });

  // -----------------------------------------------------------------------
  // Explicit encoding check: the input the contract sends matches the SPEC layout.
  // -----------------------------------------------------------------------
  describe("encoding", function () {
    it("Falcon-1024: mock accepts spec-encoded input and rejects a shifted message", async function () {
      const rng = makeRng("enc");
      const scheme = 2;
      const m = META[scheme];
      const pk = ethers.getBytes(makePubKey(scheme, rng));
      const message = ethers.getBytes(ethers.keccak256(ethers.toUtf8Bytes("enc-msg")));
      const commitment = ethers.getBytes(ethers.keccak256(ethers.concat([pk, message])));
      const nonce = rng.bytes(40);
      const esig = ethers.concat([new Uint8Array([m.esigHeader]), commitment]);
      const sigLen = ethers.getBytes(esig).length;
      const sm = ethers.concat([new Uint8Array([(sigLen >> 8) & 0xff, sigLen & 0xff]), nonce, message, esig]);
      const input = ethers.hexlify(ethers.concat([pk, sm]));

      // direct call to the installed precompile mock
      let ret = await ethers.provider.call({ to: PRECOMPILE[scheme], data: input });
      expect(BigInt(ret)).to.equal(1n);

      // shift the message by one byte -> commitment no longer matches -> invalid
      const badMsg = Uint8Array.from(message);
      badMsg[0] ^= 0xff;
      const smBad = ethers.concat([new Uint8Array([(sigLen >> 8) & 0xff, sigLen & 0xff]), nonce, badMsg, esig]);
      ret = await ethers.provider.call({ to: PRECOMPILE[scheme], data: ethers.hexlify(ethers.concat([pk, smBad])) });
      expect(BigInt(ret)).to.equal(0n);
    });
  });

  // -----------------------------------------------------------------------
  // Seeded invariant fuzz
  // -----------------------------------------------------------------------
  describe("seeded invariant fuzz", function () {
    it("holds forgery / replay / key-binding invariants over many random rounds", async function () {
      const rng = makeRng("fuzz-seed-v1");
      const ROUNDS = 48;
      let expectedCount = 0n;

      for (let r = 0; r < ROUNDS; r++) {
        const scheme = [1, 2, 3, 4][rng.int(4)];
        const pk = makePubKey(scheme, rng);
        const keyId = await register(scheme, pk);
        const messageHash = ethers.hexlify(rng.bytes(32));

        // --- genuine attestation succeeds; state advances by exactly one ---
        const nonce0 = Number(await att.nonceOf(keyId));
        const challenge0 = await att.attestChallenge(keyId, nonce0, messageHash);
        const good = genuine(scheme, pk, challenge0, rng);
        await att.attest(keyId, messageHash, good);
        expectedCount += 1n;
        expect(await att.attestationCount()).to.equal(expectedCount);
        expect(Number(await att.nonceOf(keyId))).to.equal(nonce0 + 1);
        expect(await att.isValidAttestation(keyId, nonce0, messageHash)).to.equal(true);

        // --- a battery of forgeries at the new nonce: ALL rejected, state frozen ---
        const nonce1 = Number(await att.nonceOf(keyId));
        const challenge1 = await att.attestChallenge(keyId, nonce1, messageHash);
        const forgeries = [];

        // (a) tampered commitment
        forgeries.push(tamperCommitment(scheme, genuine(scheme, pk, challenge1, rng)));
        // (b) genuine-but-for-a-different-message
        {
          const otherMsg = ethers.hexlify(rng.bytes(32));
          const otherCh = await att.attestChallenge(keyId, nonce1, otherMsg);
          forgeries.push(genuine(scheme, pk, otherCh, rng));
        }
        // (c) genuine-but-by-a-different-key
        {
          const otherPk = makePubKey(scheme, rng);
          forgeries.push(genuine(scheme, otherPk, challenge1, rng));
        }
        // (d) replay of the just-consumed signature (challenge was nonce0)
        forgeries.push(good);
        // (e) stale: signature for the previous nonce
        forgeries.push(genuine(scheme, pk, challenge0, rng));

        for (const f of forgeries) {
          await expect(att.attest(keyId, messageHash, f)).to.be.revertedWithCustomError(att, "PQCVerificationFailed");
        }
        // invariants unchanged after the whole battery
        expect(await att.attestationCount()).to.equal(expectedCount);
        expect(Number(await att.nonceOf(keyId))).to.equal(nonce1);
        expect((await att.getAttestation(keyId, nonce1)).exists).to.equal(false);

        // --- a second genuine attestation at the new nonce still works ---
        const good1 = genuine(scheme, pk, challenge1, rng);
        await att.attest(keyId, messageHash, good1);
        expectedCount += 1n;
        expect(await att.attestationCount()).to.equal(expectedCount);
        expect(Number(await att.nonceOf(keyId))).to.equal(nonce1 + 1);
      }
    });
  });
});