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

// ---------------------------------------------------------------------------
// AereRecoveryRegistry test.
//
// Hardhat has no Falcon-512 precompile, so we install the repo's MockPQCPrecompile
// (scheme 1) at 0x0AE1 via hardhat_setCode. The mock parses the input at the LIVE
// precompile's spec offsets and accepts iff commitment == keccak256(pk || message).
// So a positive record here proves the contract built the 0x0AE1 input at the correct
// offsets and enforces fail-closed / append-only semantics.
//
// [MEASURE] The REAL precompile path (that the on-chain 0x0AE1 accepts a genuine
// Falcon-512 signature) is exercised separately against LIVE mainnet 2800 via eth_call
// (scripts/pqc/verify-live-precompile.js) and is NOT re-proven inside this unit test.
// The contract stays hard-wired to 0x0AE1 for production; the test only swaps in mock
// bytecode at that address.
// ---------------------------------------------------------------------------

const FALCON512 = ethers.getAddress("0x0000000000000000000000000000000000000ae1");
const PK_LEN = 897;
const PK_HEADER = 0x09;
const ESIG_HEADER = 0x29;

const FAULT_CRASH = 1;
const FAULT_NETSPLIT = 2;
const FAULT_BYZANTINE = 3;

// Deterministic seeded RNG (keccak hash-chain) so the test 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;
  }
  return { bytes };
}

function makePubKey(rng) {
  const raw = rng.bytes(PK_LEN);
  raw[0] = PK_HEADER; // Falcon-512 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]));
}

// A genuine Falcon-512 signature envelope binding to `message` under `pubKey`:
// envelope = nonce(40) || esig ; esig = 0x29 || commitment(32)  => 73 bytes.
function genuine(pubKey, message, rng) {
  const nonce = rng.bytes(40);
  const esig = ethers.concat([new Uint8Array([ESIG_HEADER]), commitmentFor(pubKey, message)]);
  return ethers.hexlify(ethers.concat([nonce, esig]));
}

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

// N=7 fault-boundary context used across the suite (matches the security spec: quorum 5, f=2).
const N7 = { n: 7, q: 5, aliveKill2: 5, aliveKill3: 4 };

describe("AereRecoveryRegistry", function () {
  let reg, mockCode;

  before(async function () {
    const F = await ethers.getContractFactory("MockPQCPrecompile");
    const mock = await F.deploy(1); // scheme 1 = Falcon-512
    await mock.waitForDeployment();
    mockCode = await ethers.provider.getCode(await mock.getAddress());
  });

  beforeEach(async function () {
    await network.provider.send("hardhat_setCode", [FALCON512, mockCode]);
    const R = await ethers.getContractFactory("AereRecoveryRegistry");
    reg = await R.deploy();
    await reg.waitForDeployment();
    // Advance block height so historical halt/resume heights are <= block.number.
    await network.provider.send("hardhat_mine", ["0x400"]); // +1024 blocks
  });

  async function registerOp(pubKey) {
    const tx = await reg.registerOperator(pubKey);
    const rc = await tx.wait();
    const log = rc.logs.map((l) => reg.interface.parseLog(l)).find((p) => p && p.name === "OperatorRegistered");
    return Number(log.args.operatorId);
  }

  // Build a signed submitRecovery call for a genuine record under `operatorId`.
  async function signedArgs(operatorId, pubKey, rng, f) {
    const nonce = await reg.nonceOf(operatorId);
    const recordHash = await reg.recordChallenge(
      operatorId, nonce, f.halt, f.resume, f.kind, f.n, f.q, f.alive
    );
    const sig = genuine(pubKey, recordHash, rng);
    return [operatorId, f.halt, f.resume, f.kind, f.n, f.q, f.alive, sig];
  }

  // -----------------------------------------------------------------------
  // Operator registration
  // -----------------------------------------------------------------------
  describe("registerOperator", function () {
    it("accepts a well-formed Falcon-512 key and assigns sequential ids", async function () {
      const rng = makeRng("reg");
      for (let i = 0; i < 3; i++) {
        const pk = makePubKey(rng);
        const id = await registerOp(pk);
        expect(id).to.equal(i);
        const op = await reg.getOperator(id);
        expect(op.pubKey).to.equal(pk);
        expect(Number(op.nonce)).to.equal(0);
      }
      expect(await reg.operatorCount()).to.equal(3n);
    });

    it("rejects a wrong-length or wrong-header key", async function () {
      await expect(reg.registerOperator("0x" + "09" + "00".repeat(895))).to.be.revertedWithCustomError(reg, "InvalidPubKey"); // 896 bytes
      await expect(reg.registerOperator("0x" + "0a" + "00".repeat(896))).to.be.revertedWithCustomError(reg, "InvalidPubKey"); // wrong header
    });
  });

  // -----------------------------------------------------------------------
  // Happy path: a valid attestation records
  // -----------------------------------------------------------------------
  describe("submitRecovery (valid attestation records)", function () {
    it("records a genuine kill-2 (f=2 boundary) crash recovery and stores every field", async function () {
      const rng = makeRng("happy");
      const pk = makePubKey(rng);
      const operatorId = await registerOp(pk);
      const f = { halt: 1000, resume: 1013, kind: FAULT_CRASH, n: N7.n, q: N7.q, alive: N7.aliveKill2 };

      const before = await reg.recordCount();
      const args = await signedArgs(operatorId, pk, rng, f);
      await expect(reg.submitRecovery(...args)).to.emit(reg, "RecoveryAttested");

      expect(await reg.recordCount()).to.equal(before + 1n);
      expect(Number(await reg.nonceOf(operatorId))).to.equal(1);

      const r = await reg.getRecord(0);
      expect(Number(r.haltHeight)).to.equal(1000);
      expect(Number(r.resumeHeight)).to.equal(1013);
      expect(Number(r.faultKind)).to.equal(FAULT_CRASH);
      expect(Number(r.nValidators)).to.equal(7);
      expect(Number(r.quorum)).to.equal(5);
      expect(Number(r.aliveDuringFault)).to.equal(5);
      expect(Number(r.operatorId)).to.equal(operatorId);
      expect(Number(r.nonce)).to.equal(0);
    });

    it("records a netsplit and a byzantine-proposer recovery under the same operator (nonce advances)", async function () {
      const rng = makeRng("kinds");
      const pk = makePubKey(rng);
      const operatorId = await registerOp(pk);

      await reg.submitRecovery(
        ...(await signedArgs(operatorId, pk, rng, { halt: 500, resume: 540, kind: FAULT_NETSPLIT, n: 7, q: 5, alive: 4 }))
      );
      await reg.submitRecovery(
        ...(await signedArgs(operatorId, pk, rng, { halt: 600, resume: 612, kind: FAULT_BYZANTINE, n: 7, q: 5, alive: 5 }))
      );

      expect(await reg.recordCount()).to.equal(2n);
      expect(Number(await reg.nonceOf(operatorId))).to.equal(2);
      expect(Number((await reg.getRecord(0)).faultKind)).to.equal(FAULT_NETSPLIT);
      expect(Number((await reg.getRecord(1)).faultKind)).to.equal(FAULT_BYZANTINE);
    });
  });

  // -----------------------------------------------------------------------
  // Fail-closed: a tampered attestation is rejected
  // -----------------------------------------------------------------------
  describe("fail-closed rejection (nothing recorded)", function () {
    it("rejects a tampered signature and records nothing", async function () {
      const rng = makeRng("tamper");
      const pk = makePubKey(rng);
      const operatorId = await registerOp(pk);
      const f = { halt: 1000, resume: 1013, kind: FAULT_CRASH, n: 7, q: 5, alive: 5 };
      const args = await signedArgs(operatorId, pk, rng, f);
      args[7] = tamperCommitment(args[7]); // corrupt the signature

      await expect(reg.submitRecovery(...args)).to.be.revertedWithCustomError(reg, "PQCVerificationFailed");
      expect(await reg.recordCount()).to.equal(0n);
      expect(Number(await reg.nonceOf(operatorId))).to.equal(0);
    });

    it("rejects a signature bound to DIFFERENT record fields (heights / fault kind)", async function () {
      const rng = makeRng("bind");
      const pk = makePubKey(rng);
      const operatorId = await registerOp(pk);
      // Sign a record for kind=CRASH at 1000..1013, then try to submit BYZANTINE with the same sig.
      const signed = await signedArgs(operatorId, pk, rng, { halt: 1000, resume: 1013, kind: FAULT_CRASH, n: 7, q: 5, alive: 5 });
      const sig = signed[7];
      await expect(
        reg.submitRecovery(operatorId, 1000, 1013, FAULT_BYZANTINE, 7, 5, 5, sig)
      ).to.be.revertedWithCustomError(reg, "PQCVerificationFailed");
      // ... and a shifted resume height, same sig.
      await expect(
        reg.submitRecovery(operatorId, 1000, 1099, FAULT_CRASH, 7, 5, 5, sig)
      ).to.be.revertedWithCustomError(reg, "PQCVerificationFailed");
      expect(await reg.recordCount()).to.equal(0n);
    });

    it("treats an empty precompile return (pre-fork / stale chain) as invalid", async function () {
      const rng = makeRng("stale");
      const pk = makePubKey(rng);
      const operatorId = await registerOp(pk);
      const args = await signedArgs(operatorId, pk, rng, { halt: 1000, resume: 1013, kind: FAULT_CRASH, n: 7, q: 5, alive: 5 });
      // Wipe the precompile: an empty account returns success + empty returndata.
      await network.provider.send("hardhat_setCode", [FALCON512, "0x"]);
      await expect(reg.submitRecovery(...args)).to.be.revertedWithCustomError(reg, "PQCVerificationFailed");
      expect(await reg.recordCount()).to.equal(0n);
    });

    it("reverts on unknown operator, invalid fault kind, bad heights, and bad validator set", async function () {
      const rng = makeRng("validate");
      const pk = makePubKey(rng);
      const operatorId = await registerOp(pk);
      const dummySig = "0x" + "00".repeat(73);
      const cur = await ethers.provider.getBlockNumber();

      await expect(reg.submitRecovery(99, 1000, 1013, FAULT_CRASH, 7, 5, 5, dummySig)).to.be.revertedWithCustomError(reg, "UnknownOperator");
      await expect(reg.submitRecovery(operatorId, 1000, 1013, 0, 7, 5, 5, dummySig)).to.be.revertedWithCustomError(reg, "InvalidFaultKind");
      await expect(reg.submitRecovery(operatorId, 1000, 1013, 4, 7, 5, 5, dummySig)).to.be.revertedWithCustomError(reg, "InvalidFaultKind");
      await expect(reg.submitRecovery(operatorId, 0, 10, FAULT_CRASH, 7, 5, 5, dummySig)).to.be.revertedWithCustomError(reg, "InvalidHeights");
      await expect(reg.submitRecovery(operatorId, 1013, 1000, FAULT_CRASH, 7, 5, 5, dummySig)).to.be.revertedWithCustomError(reg, "InvalidHeights"); // resume < halt
      await expect(reg.submitRecovery(operatorId, 1000, cur + 1000, FAULT_CRASH, 7, 5, 5, dummySig)).to.be.revertedWithCustomError(reg, "InvalidHeights"); // future resume
      await expect(reg.submitRecovery(operatorId, 1000, 1013, FAULT_CRASH, 5, 7, 5, dummySig)).to.be.revertedWithCustomError(reg, "InvalidValidatorSet"); // quorum > n
      await expect(reg.submitRecovery(operatorId, 1000, 1013, FAULT_CRASH, 7, 5, 9, dummySig)).to.be.revertedWithCustomError(reg, "InvalidValidatorSet"); // alive > n
      expect(await reg.recordCount()).to.equal(0n);
    });
  });

  // -----------------------------------------------------------------------
  // Append-only: cannot overwrite
  // -----------------------------------------------------------------------
  describe("append-only (cannot overwrite)", function () {
    it("appends new records without ever rewriting an earlier record, and rejects replay", async function () {
      const rng = makeRng("append");
      const pk = makePubKey(rng);
      const operatorId = await registerOp(pk);

      // Record 0.
      const args0 = await signedArgs(operatorId, pk, rng, { halt: 1000, resume: 1013, kind: FAULT_CRASH, n: 7, q: 5, alive: 5 });
      await reg.submitRecovery(...args0);
      const r0before = await reg.getRecord(0);

      // Record 1 (nonce has advanced to 1).
      const args1 = await signedArgs(operatorId, pk, rng, { halt: 2000, resume: 2050, kind: FAULT_NETSPLIT, n: 7, q: 5, alive: 4 });
      await reg.submitRecovery(...args1);

      expect(await reg.recordCount()).to.equal(2n);
      // Record 0 is byte-for-byte unchanged after record 1 was appended.
      const r0after = await reg.getRecord(0);
      expect(r0after.recordHash).to.equal(r0before.recordHash);
      expect(Number(r0after.haltHeight)).to.equal(1000);
      expect(Number(r0after.resumeHeight)).to.equal(1013);
      expect(Number(r0after.faultKind)).to.equal(FAULT_CRASH);
      // Record 1 landed at a fresh id.
      expect(Number((await reg.getRecord(1)).haltHeight)).to.equal(2000);

      // Replay of the exact consumed (record 0) signature now fails: the challenge uses nonce 2.
      await expect(reg.submitRecovery(...args0)).to.be.revertedWithCustomError(reg, "PQCVerificationFailed");
      expect(await reg.recordCount()).to.equal(2n);
      // Record 0 still unchanged after the replay attempt.
      expect((await reg.getRecord(0)).recordHash).to.equal(r0before.recordHash);
    });

    it("exposes no function that can overwrite or delete a record", async function () {
      // Structural check: the ABI has no setter/updater/delete entry point for records.
      const names = reg.interface.fragments.filter((f) => f.type === "function").map((f) => f.name.toLowerCase());
      for (const forbidden of ["set", "update", "edit", "delete", "remove", "overwrite", "rewrite", "setrecord", "admin", "owner"]) {
        expect(names.some((n) => n.includes(forbidden))).to.equal(false);
      }
    });
  });
});
