// AerePQCKeyExchange - on-chain registry + confirmed handshake for ML-KEM-768 key establishment.
//
// A recipient registers an ML-KEM-768 encapsulation key; an initiator encapsulates a shared secret K
// to it OFF-CHAIN (via the ML-KEM precompile 0x0AE6 where activated, or any FIPS-203 library) and
// posts the KEM ciphertext hash + a commitment to a one-way, K-derived confirmation opening kc. The
// recipient decapsulates, recomputes kc, and confirms by revealing it; the contract checks the commit
// and records the confirmation. Unconfirmed handshakes resolve via timeout-cancel.
//
// This suite drives the on-chain rail directly. It does NOT run real ML-KEM: the contract only checks
// byte LENGTHS (ek = 1184, ciphertext = 1088) and HASHES / commit-reveal, so deterministic filler
// bytes and a simulated 32-byte shared secret exercise every on-chain branch faithfully. Coverage:
//   - register / rotate / retire (happy path + every validation revert)
//   - initiateHandshake (happy path + every validation revert)
//   - confirmHandshake (happy path, only-responder, mismatch-then-correct, expiry, wrong-state)
//   - cancel / timeout-cancel (initiator abort, permissionless timeout reaper, cannot cancel confirmed)
//
// Run: npx hardhat test test/AerePQCKeyExchange.test.js
const { expect } = require("chai");
const { ethers } = require("hardhat");
const { time } = require("@nomicfoundation/hardhat-toolbox/network-helpers");
const { anyValue } = require("@nomicfoundation/hardhat-chai-matchers/withArgs");

const coder = ethers.AbiCoder.defaultAbiCoder();

const EK_LEN = 1184; // ML-KEM-768 encapsulation key
const CT_LEN = 1088; // ML-KEM-768 ciphertext
const HOUR = 3600;
const DAY = 24 * HOUR;
const CONFIRM_DOMAIN = ethers.keccak256(ethers.toUtf8Bytes("AerePQCKeyExchange.v1.confirm"));

// deterministic filler bytes (keccak hash-chain), so tests are reproducible.
function fill(n, seedText) {
  const out = new Uint8Array(n);
  let seed = ethers.keccak256(ethers.toUtf8Bytes(seedText));
  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 ethers.hexlify(out);
}

describe("AerePQCKeyExchange (ML-KEM-768 registry + confirmed handshake)", function () {
  let c, cAddr, chainId;
  let alice, bob, relayer, stranger; // alice = initiator, bob = recipient/responder

  const WINDOW = 2 * HOUR;

  // A canonical recipient key and one KEM ciphertext blob.
  const EK = fill(EK_LEN, "bob-ek-0");
  const EK_HASH = ethers.keccak256(EK);
  const KEM_CT = fill(CT_LEN, "kem-ct-0");
  const KEM_CT_HASH = ethers.keccak256(KEM_CT);
  // A simulated 32-byte ML-KEM shared secret (in production: output of Encaps/Decaps).
  const K = ethers.keccak256(ethers.toUtf8Bytes("shared-secret-0"));

  // Off-chain derivation of the confirmation opening kc, mirroring the recommended NatSpec formula:
  // kc = keccak256(abi.encode(CONFIRM_DOMAIN, chainid, contract, initiator, responder, ekHash, ctHash, K)).
  function deriveKc(initiator, responder, ekHash, ctHash, k) {
    return ethers.keccak256(
      coder.encode(
        ["bytes32", "uint256", "address", "address", "address", "bytes32", "bytes32", "bytes32"],
        [CONFIRM_DOMAIN, chainId, cAddr, initiator, responder, ekHash, ctHash, k]
      )
    );
  }
  const commitFor = (kc) => ethers.keccak256(coder.encode(["bytes32"], [kc]));

  beforeEach(async function () {
    [alice, bob, relayer, stranger] = await ethers.getSigners();
    const F = await ethers.getContractFactory("AerePQCKeyExchange");
    c = await F.deploy();
    await c.waitForDeployment();
    cAddr = await c.getAddress();
    chainId = (await ethers.provider.getNetwork()).chainId;
    // Bob registers his ML-KEM encapsulation key.
    await (await c.connect(bob).registerKemKey(EK)).wait();
  });

  // Alice initiates a standard handshake to Bob; returns { id, kc, commit }.
  async function initiateStd(overrides = {}) {
    const kc = overrides.kc ?? deriveKc(alice.address, bob.address, EK_HASH, KEM_CT_HASH, K);
    const o = {
      responder: bob.address,
      ekHash: EK_HASH,
      kemCt: KEM_CT,
      commit: overrides.commit ?? commitFor(kc),
      window: WINDOW,
      ...overrides,
    };
    const id = await c.nextHandshakeId();
    await (
      await c.connect(alice).initiateHandshake(o.responder, o.ekHash, o.kemCt, o.commit, o.window)
    ).wait();
    return { id, kc, commit: o.commit };
  }

  // =========================================================================
  //  ML-KEM key registry
  // =========================================================================
  describe("register / rotate / retire", function () {
    it("registers a key, stores its hash, epoch 1, active", async function () {
      const k = await c.kemKeyOf(bob.address);
      expect(k.keyHash).to.equal(EK_HASH);
      expect(k.epoch).to.equal(1n);
      expect(k.active).to.equal(true);
    });

    it("rejects a wrong-length encapsulation key", async function () {
      await expect(c.connect(stranger).registerKemKey(fill(EK_LEN - 1, "short")))
        .to.be.revertedWithCustomError(c, "BadEncapKeyLength");
    });

    it("rejects a second register while a key is already active", async function () {
      await expect(c.connect(bob).registerKemKey(fill(EK_LEN, "bob-ek-1")))
        .to.be.revertedWithCustomError(c, "KeyAlreadyActive");
    });

    it("rotate bumps the epoch and changes the bound hash", async function () {
      const ek2 = fill(EK_LEN, "bob-ek-1");
      await expect(c.connect(bob).rotateKemKey(ek2)).to.emit(c, "KemKeyRotated");
      const k = await c.kemKeyOf(bob.address);
      expect(k.epoch).to.equal(2n);
      expect(k.keyHash).to.equal(ethers.keccak256(ek2));
      // a handshake that binds the OLD hash is now stale
      await expect(initiateStd({ ekHash: EK_HASH })).to.be.revertedWithCustomError(c, "StaleEncapKey");
    });

    it("rotate without an active key reverts", async function () {
      await expect(c.connect(stranger).rotateKemKey(fill(EK_LEN, "x"))).to.be.revertedWithCustomError(c, "NoActiveKey");
    });

    it("retire blocks new handshakes but re-register is allowed afterwards", async function () {
      await (await c.connect(bob).retireKemKey()).wait();
      expect((await c.kemKeyOf(bob.address)).active).to.equal(false);
      await expect(initiateStd()).to.be.revertedWithCustomError(c, "ResponderKeyInactive");
      // retire twice reverts
      await expect(c.connect(bob).retireKemKey()).to.be.revertedWithCustomError(c, "NoActiveKey");
      // re-register with a fresh key works and bumps epoch to 2
      const ek2 = fill(EK_LEN, "bob-ek-2");
      await (await c.connect(bob).registerKemKey(ek2)).wait();
      const k = await c.kemKeyOf(bob.address);
      expect(k.epoch).to.equal(2n);
      expect(k.active).to.equal(true);
      expect(k.keyHash).to.equal(ethers.keccak256(ek2));
    });
  });

  // =========================================================================
  //  initiateHandshake
  // =========================================================================
  describe("initiateHandshake", function () {
    it("stores ct hash + commitment, arms the deadline, emits the full ciphertext", async function () {
      const kc = deriveKc(alice.address, bob.address, EK_HASH, KEM_CT_HASH, K);
      const commit = commitFor(kc);
      const id = await c.nextHandshakeId();
      await expect(c.connect(alice).initiateHandshake(bob.address, EK_HASH, KEM_CT, commit, WINDOW))
        .to.emit(c, "HandshakeInitiated")
        .withArgs(id, alice.address, bob.address, EK_HASH, KEM_CT_HASH, commit, anyValue, KEM_CT);

      const h = await c.getHandshake(id);
      expect(h.initiator).to.equal(alice.address);
      expect(h.responder).to.equal(bob.address);
      expect(h.state).to.equal(1n); // Initiated
      expect(h.ekHash).to.equal(EK_HASH);
      expect(h.kemCtHash).to.equal(KEM_CT_HASH);
      expect(h.confirmCommit).to.equal(commit);
      expect(h.confirmOpening).to.equal(ethers.ZeroHash);
    });

    it("rejects an inactive/unknown responder", async function () {
      await expect(initiateStd({ responder: stranger.address }))
        .to.be.revertedWithCustomError(c, "ResponderKeyInactive");
    });

    it("rejects a stale ekHash", async function () {
      await expect(initiateStd({ ekHash: ethers.ZeroHash })).to.be.revertedWithCustomError(c, "StaleEncapKey");
    });

    it("rejects a wrong-length KEM ciphertext", async function () {
      await expect(initiateStd({ kemCt: fill(CT_LEN - 1, "badct") }))
        .to.be.revertedWithCustomError(c, "BadCiphertextLength");
    });

    it("rejects a zero commitment", async function () {
      await expect(initiateStd({ commit: ethers.ZeroHash })).to.be.revertedWithCustomError(c, "ZeroCommitment");
    });

    it("rejects out-of-range windows", async function () {
      await expect(initiateStd({ window: 0 })).to.be.revertedWithCustomError(c, "HandshakeWindowOutOfRange");
      await expect(initiateStd({ window: 31 * DAY })).to.be.revertedWithCustomError(c, "HandshakeWindowOutOfRange");
    });

    it("confirmCommitFor view matches the off-chain derivation", async function () {
      const kc = deriveKc(alice.address, bob.address, EK_HASH, KEM_CT_HASH, K);
      expect(await c.confirmCommitFor(kc)).to.equal(commitFor(kc));
    });
  });

  // =========================================================================
  //  confirmHandshake
  // =========================================================================
  describe("confirmHandshake", function () {
    it("responder reveals the matching opening; handshake becomes Confirmed and records kc", async function () {
      const { id, kc } = await initiateStd();
      await expect(c.connect(bob).confirmHandshake(id, kc))
        .to.emit(c, "HandshakeConfirmed")
        .withArgs(id, bob.address, kc);
      const h = await c.getHandshake(id);
      expect(h.state).to.equal(2n); // Confirmed
      expect(h.confirmOpening).to.equal(kc);
    });

    it("only the designated responder may confirm", async function () {
      const { id, kc } = await initiateStd();
      await expect(c.connect(stranger).confirmHandshake(id, kc))
        .to.be.revertedWithCustomError(c, "NotResponder");
      await expect(c.connect(alice).confirmHandshake(id, kc))
        .to.be.revertedWithCustomError(c, "NotResponder");
    });

    it("a mismatched opening reverts and leaves state untouched; the correct one then confirms", async function () {
      const { id, kc } = await initiateStd();
      const wrongKc = deriveKc(alice.address, bob.address, EK_HASH, KEM_CT_HASH, ethers.keccak256(ethers.toUtf8Bytes("wrong-K")));
      await expect(c.connect(bob).confirmHandshake(id, wrongKc))
        .to.be.revertedWithCustomError(c, "ConfirmMismatch");
      // still Initiated
      expect((await c.getHandshake(id)).state).to.equal(1n);
      // correct opening within the window still works
      await (await c.connect(bob).confirmHandshake(id, kc)).wait();
      expect((await c.getHandshake(id)).state).to.equal(2n);
    });

    it("cannot confirm after the deadline", async function () {
      const { id, kc } = await initiateStd();
      await time.increase(WINDOW + 60);
      await expect(c.connect(bob).confirmHandshake(id, kc))
        .to.be.revertedWithCustomError(c, "HandshakeExpired");
    });

    it("cannot confirm an unknown handshake", async function () {
      await expect(c.connect(bob).confirmHandshake(999, ethers.ZeroHash))
        .to.be.revertedWithCustomError(c, "UnknownHandshake");
    });

    it("cannot confirm twice (state no longer Initiated)", async function () {
      const { id, kc } = await initiateStd();
      await (await c.connect(bob).confirmHandshake(id, kc)).wait();
      await expect(c.connect(bob).confirmHandshake(id, kc))
        .to.be.revertedWithCustomError(c, "WrongState");
    });
  });

  // =========================================================================
  //  cancel / timeout-cancel
  // =========================================================================
  describe("cancel / timeout-cancel", function () {
    it("the initiator may abort a pending handshake; confirm then impossible", async function () {
      const { id, kc } = await initiateStd();
      await expect(c.connect(alice).cancelHandshake(id))
        .to.emit(c, "HandshakeCancelled")
        .withArgs(id, alice.address, false);
      expect((await c.getHandshake(id)).state).to.equal(3n); // Cancelled
      await expect(c.connect(bob).confirmHandshake(id, kc))
        .to.be.revertedWithCustomError(c, "WrongState");
    });

    it("a stranger cannot cancel before the deadline", async function () {
      const { id } = await initiateStd();
      expect(await c.isTimedOut(id)).to.equal(false);
      await expect(c.connect(stranger).cancelHandshake(id))
        .to.be.revertedWithCustomError(c, "NotCancellable");
    });

    it("anyone may timeout-cancel once the deadline passes; confirm then impossible", async function () {
      const { id, kc } = await initiateStd();
      await time.increase(WINDOW + 60);
      expect(await c.isTimedOut(id)).to.equal(true);
      await expect(c.connect(relayer).cancelHandshake(id))
        .to.emit(c, "HandshakeCancelled")
        .withArgs(id, relayer.address, true);
      expect((await c.getHandshake(id)).state).to.equal(3n); // Cancelled
      expect(await c.isTimedOut(id)).to.equal(false); // no longer Initiated
      await expect(c.connect(bob).confirmHandshake(id, kc))
        .to.be.revertedWithCustomError(c, "WrongState");
    });

    it("cannot cancel a confirmed handshake", async function () {
      const { id, kc } = await initiateStd();
      await (await c.connect(bob).confirmHandshake(id, kc)).wait();
      await expect(c.connect(alice).cancelHandshake(id))
        .to.be.revertedWithCustomError(c, "WrongState");
    });

    it("cannot cancel an unknown handshake", async function () {
      await expect(c.connect(alice).cancelHandshake(12345))
        .to.be.revertedWithCustomError(c, "UnknownHandshake");
    });
  });
});