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

// ---------------------------------------------------------------------------
// AereAgentDID — Falcon-root DID for AI agents issuing short-lived secp256k1
// session keys with scope + spend-policy + on-chain revocation.
//
// The ROOT authority is a Falcon key held in AerePQCKeyRegistry (proof-of-possession
// enforced there). We install MockPQCPrecompile at 0x0AE1..0x0AE4 (hardhat_setCode) so
// the registry's Falcon verification runs; the real precompile encoding is proven
// separately against the LIVE mainnet precompile via eth_call. Session ISSUANCE is
// gated by a Falcon PoP over an issuance challenge; session ACTIONS are cheap secp256k1
// (ecrecover). A genuine Falcon envelope has commitment == keccak256(pk || challenge).
// ---------------------------------------------------------------------------

const PRECOMPILE = {
  1: ethers.getAddress("0x0000000000000000000000000000000000000ae1"),
  2: ethers.getAddress("0x0000000000000000000000000000000000000ae2"),
  3: ethers.getAddress("0x0000000000000000000000000000000000000ae3"),
  4: ethers.getAddress("0x0000000000000000000000000000000000000ae4"),
};

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 },
};

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;
  return ethers.hexlify(raw);
}

function commitmentFor(pubKey, message) {
  return ethers.keccak256(ethers.concat([pubKey, message]));
}

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]);
    return ethers.hexlify(ethers.concat([nonce, esig]));
  } else {
    const sig = new Uint8Array(m.sigLen);
    sig.set(ethers.getBytes(commitment), 0);
    return ethers.hexlify(sig);
  }
}

function genuine(scheme, pubKey, message, rng) {
  return envelope(scheme, commitmentFor(pubKey, message), rng);
}

function tamperCommitment(scheme, sigHex) {
  const b = ethers.getBytes(sigHex);
  const idx = META[scheme].falcon ? 41 : 0;
  b[idx] ^= 0x01;
  return ethers.hexlify(b);
}

// secp256k1 signature (r||s||v, low-s, v in {27,28}) by `wallet` over the 32-byte digest.
function ecdsaSign(wallet, digest) {
  return wallet.signingKey.sign(digest).serialized;
}

const AGENT_BOND = ethers.getAddress("0x32E0015F622a8719d1C380A87CE1a09bcd0cB86A"); // AereAgentBond (mainnet)
const AI_REPUTATION = ethers.getAddress("0x781ef746c08760aa854cDa4621d54db6734bfeBF"); // AereAIReputation

describe("AereAgentDID (Falcon-root DID + secp256k1 sessions + seeded fuzz)", function () {
  let reg, did, owner, other, relayer, mockCode;

  before(async function () {
    [owner, other, relayer] = await ethers.getSigners();
    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 R = await ethers.getContractFactory("AerePQCKeyRegistry");
    reg = await R.deploy();
    await reg.waitForDeployment();
    const D = await ethers.getContractFactory("AereAgentDID");
    did = await D.deploy(await reg.getAddress(), AGENT_BOND, AI_REPUTATION);
    await did.waitForDeployment();
  });

  // Register a Falcon root key under `signer`, return { rootKeyId, rootPk, scheme }.
  async function registerFalconRoot(scheme, signer = owner, rng = makeRng("root-" + scheme)) {
    const pk = makePubKey(scheme, rng);
    const nonce = await reg.identityNonce(signer.address);
    const ch = await reg.popChallenge(signer.address, scheme, pk, nonce);
    const tx = await reg.connect(signer).registerKey(scheme, pk, genuine(scheme, pk, ch, rng));
    const rc = await tx.wait();
    const log = rc.logs.map((l) => reg.interface.parseLog(l)).find((p) => p && p.name === "KeyRegistered");
    return { rootKeyId: Number(log.args.keyId), rootPk: pk, scheme };
  }

  // Issue a session with a genuine Falcon PoP (relayed by `caller`).
  async function issueSession(
    root,
    { sessionAddr, scopeHash, spendCap, expiry, rng, caller = relayer, signer = owner }
  ) {
    const nonce = (await did.getAgent(root.rootKeyId)).sessionNonce;
    const challenge = await did.sessionChallenge(root.rootKeyId, nonce, sessionAddr, scopeHash, spendCap, expiry);
    const pop = genuine(root.scheme, root.rootPk, challenge, rng);
    const tx = await did.connect(caller).issueSession(root.rootKeyId, sessionAddr, scopeHash, spendCap, expiry, pop);
    const rc = await tx.wait();
    const log = rc.logs.map((l) => did.interface.parseLog(l)).find((p) => p && p.name === "SessionIssued");
    return Number(log.args.sessionId);
  }

  // Base expiry on the ON-CHAIN latest block timestamp (evm_increaseTime in earlier
  // tests advances the node clock past wall time, so Date.now() would be stale).
  async function futureTs(secs = 3600) {
    const b = await ethers.provider.getBlock("latest");
    return b.timestamp + secs;
  }

  // -----------------------------------------------------------------------
  // Deployment & cross-links
  // -----------------------------------------------------------------------
  describe("deployment", function () {
    it("wires the registry and exposes the read-only cross-links", async function () {
      expect(await did.keyRegistry()).to.equal(await reg.getAddress());
      expect(await did.agentBond()).to.equal(AGENT_BOND);
      expect(await did.aiReputation()).to.equal(AI_REPUTATION);
    });

    it("reverts on a zero registry", async function () {
      const D = await ethers.getContractFactory("AereAgentDID");
      await expect(D.deploy(ethers.ZeroAddress, AGENT_BOND, AI_REPUTATION)).to.be.revertedWithCustomError(
        D,
        "ZeroRegistry"
      );
    });
  });

  // -----------------------------------------------------------------------
  // createAgent
  // -----------------------------------------------------------------------
  describe("createAgent", function () {
    it("creates a DID for an ACTIVE Falcon root, callable only by the root controller", async function () {
      const rng = makeRng("create");
      const root = await registerFalconRoot(1, owner, rng);
      // non-controller cannot create.
      await expect(did.connect(other).createAgent(root.rootKeyId)).to.be.revertedWithCustomError(
        did,
        "NotRootController"
      );
      await expect(did.createAgent(root.rootKeyId)).to.emit(did, "AgentCreated");
      const a = await did.getAgent(root.rootKeyId);
      expect(a.controller).to.equal(owner.address);
      expect(Number(a.sessionNonce)).to.equal(0);
      expect(await did.agentExists(root.rootKeyId)).to.equal(true);
    });

    it("rejects a non-Falcon root key", async function () {
      const rng = makeRng("nonfalcon");
      // register an ML-DSA-44 key and try to use it as a root.
      const pk = makePubKey(3, rng);
      const nonce = await reg.identityNonce(owner.address);
      const ch = await reg.popChallenge(owner.address, 3, pk, nonce);
      const tx = await reg.registerKey(3, pk, genuine(3, pk, ch, rng));
      const rc = await tx.wait();
      const id = Number(
        rc.logs.map((l) => reg.interface.parseLog(l)).find((p) => p && p.name === "KeyRegistered").args.keyId
      );
      await expect(did.createAgent(id)).to.be.revertedWithCustomError(did, "RootKeyNotFalcon");
    });

    it("rejects a rotated-away (non-ACTIVE) root and a duplicate DID", async function () {
      const rng = makeRng("dup");
      const root = await registerFalconRoot(2, owner, rng);
      await did.createAgent(root.rootKeyId);
      await expect(did.createAgent(root.rootKeyId)).to.be.revertedWithCustomError(did, "AgentExists");

      // A separate root that we rotate away is no longer ACTIVE -> cannot create.
      const root2 = await registerFalconRoot(1, owner, rng);
      const np = makePubKey(1, rng);
      const n = await reg.identityNonce(owner.address);
      const ch = await reg.popChallenge(owner.address, 1, np, n);
      await reg.rotateKey(root2.rootKeyId, 1, np, genuine(1, np, ch, rng));
      await expect(did.createAgent(root2.rootKeyId)).to.be.revertedWithCustomError(did, "RootKeyNotActive");
    });
  });

  // -----------------------------------------------------------------------
  // issueSession (Falcon-PoP gated, relayable)
  // -----------------------------------------------------------------------
  describe("issueSession", function () {
    it("issues a session with a genuine Falcon PoP, relayable by anyone", async function () {
      const rng = makeRng("issue");
      const root = await registerFalconRoot(1, owner, rng);
      await did.createAgent(root.rootKeyId);
      const session = ethers.Wallet.createRandom();
      const scopeHash = ethers.keccak256(ethers.toUtf8Bytes("swap:read"));
      const expiry = await futureTs();

      const sid = await issueSession(root, {
        sessionAddr: session.address,
        scopeHash,
        spendCap: 1000n,
        expiry,
        rng,
        caller: relayer, // relayed, not the controller
      });
      expect(sid).to.equal(0);
      const s = await did.getSession(sid);
      expect(s.sessionAddr).to.equal(session.address);
      expect(s.scopeHash).to.equal(scopeHash);
      expect(s.spendCap).to.equal(1000n);
      expect(s.revoked).to.equal(false);
      expect(await did.isSessionValid(sid)).to.equal(true);
      expect(Number((await did.getAgent(root.rootKeyId)).sessionNonce)).to.equal(1);
    });

    it("rejects a forged PoP and a stale/replayed PoP", async function () {
      const rng = makeRng("issue-bad");
      const root = await registerFalconRoot(2, owner, rng);
      await did.createAgent(root.rootKeyId);
      const session = ethers.Wallet.createRandom();
      const scopeHash = ethers.ZeroHash;
      const expiry = await futureTs();

      const nonce = (await did.getAgent(root.rootKeyId)).sessionNonce;
      const challenge = await did.sessionChallenge(root.rootKeyId, nonce, session.address, scopeHash, 500n, expiry);
      const forged = tamperCommitment(2, genuine(2, root.rootPk, challenge, rng));
      await expect(
        did.issueSession(root.rootKeyId, session.address, scopeHash, 500n, expiry, forged)
      ).to.be.revertedWithCustomError(did, "SessionAuthFailed");

      // A genuine PoP for the CURRENT nonce works once...
      const good = genuine(2, root.rootPk, challenge, rng);
      await did.issueSession(root.rootKeyId, session.address, scopeHash, 500n, expiry, good);
      // ...and replaying it fails (nonce advanced -> different challenge).
      await expect(
        did.issueSession(root.rootKeyId, session.address, scopeHash, 500n, expiry, good)
      ).to.be.revertedWithCustomError(did, "SessionAuthFailed");
    });

    it("rejects a past expiry, a zero session address, and an unknown agent", async function () {
      const rng = makeRng("issue-params");
      const root = await registerFalconRoot(1, owner, rng);
      await did.createAgent(root.rootKeyId);
      const session = ethers.Wallet.createRandom();
      const pastExpiry = 1; // in the past
      const nonce = (await did.getAgent(root.rootKeyId)).sessionNonce;
      const ch = await did.sessionChallenge(root.rootKeyId, nonce, session.address, ethers.ZeroHash, 1n, pastExpiry);
      const pop = genuine(1, root.rootPk, ch, rng);
      await expect(
        did.issueSession(root.rootKeyId, session.address, ethers.ZeroHash, 1n, pastExpiry, pop)
      ).to.be.revertedWithCustomError(did, "InvalidSessionParams");
      // unknown agent.
      await expect(
        did.issueSession(999, session.address, ethers.ZeroHash, 1n, await futureTs(), "0x")
      ).to.be.revertedWithCustomError(did, "UnknownAgent");
    });
  });

  // -----------------------------------------------------------------------
  // authorize (secp256k1 hot path: scope + spend cap + anti-replay)
  // -----------------------------------------------------------------------
  describe("authorize", function () {
    let root, session, sid, scopeHash;
    beforeEach(async function () {
      const rng = makeRng("auth-setup");
      root = await registerFalconRoot(1, owner, rng);
      await did.createAgent(root.rootKeyId);
      session = ethers.Wallet.createRandom();
      scopeHash = ethers.keccak256(ethers.toUtf8Bytes("trade"));
      sid = await issueSession(root, {
        sessionAddr: session.address,
        scopeHash,
        spendCap: 100n,
        expiry: await futureTs(),
        rng,
      });
    });

    async function authorize(amount, actionLabel, opts = {}) {
      const scope = opts.scope ?? scopeHash;
      const actionHash = ethers.keccak256(ethers.toUtf8Bytes(actionLabel));
      const nonce = opts.nonce ?? (await did.getSession(sid)).actionNonce;
      const digest = await did.actionDigest(sid, scope, actionHash, amount, nonce);
      const wallet = opts.wallet ?? session;
      const sig = ecdsaSign(wallet, digest);
      return did.connect(relayer).authorize(sid, scope, actionHash, amount, sig);
    }

    it("authorizes a scoped action signed by the session key and records spend", async function () {
      await expect(authorize(30n, "a1")).to.emit(did, "ActionAuthorized");
      expect((await did.getSession(sid)).spent).to.equal(30n);
      expect(Number((await did.getSession(sid)).actionNonce)).to.equal(1);
      expect(await did.remainingSpend(sid)).to.equal(70n);
      await authorize(70n, "a2");
      expect((await did.getSession(sid)).spent).to.equal(100n);
      expect(await did.remainingSpend(sid)).to.equal(0n);
    });

    it("enforces the spend cap (cumulative)", async function () {
      await authorize(60n, "a1");
      await expect(authorize(41n, "a2")).to.be.revertedWithCustomError(did, "SessionSpendExceeded");
      // exactly hitting the cap is allowed.
      await authorize(40n, "a3");
      expect((await did.getSession(sid)).spent).to.equal(100n);
    });

    it("enforces scope: a non-matching scope tag reverts", async function () {
      const wrongScope = ethers.keccak256(ethers.toUtf8Bytes("withdraw"));
      await expect(authorize(1n, "x", { scope: wrongScope })).to.be.revertedWithCustomError(
        did,
        "SessionScopeViolation"
      );
    });

    it("rejects a signature by a non-session key", async function () {
      const stranger = ethers.Wallet.createRandom();
      await expect(authorize(1n, "x", { wallet: stranger })).to.be.revertedWithCustomError(did, "SessionSigInvalid");
    });

    it("is anti-replay: reusing a consumed action nonce reverts", async function () {
      // Sign at nonce 0, submit -> nonce becomes 1. Re-submitting the SAME signature
      // (still nonce 0) now recovers a wrong digest -> SessionSigInvalid.
      const actionHash = ethers.keccak256(ethers.toUtf8Bytes("replay"));
      const digest0 = await did.actionDigest(sid, scopeHash, actionHash, 10n, 0);
      const sig0 = ecdsaSign(session, digest0);
      await did.authorize(sid, scopeHash, actionHash, 10n, sig0);
      expect(Number((await did.getSession(sid)).actionNonce)).to.equal(1);
      await expect(did.authorize(sid, scopeHash, actionHash, 10n, sig0)).to.be.revertedWithCustomError(
        did,
        "SessionSigInvalid"
      );
    });

    it("rejects a malformed (wrong-length) signature", async function () {
      const actionHash = ethers.keccak256(ethers.toUtf8Bytes("m"));
      await expect(did.authorize(sid, scopeHash, actionHash, 1n, "0x1234")).to.be.revertedWithCustomError(
        did,
        "SessionSigInvalid"
      );
    });

    it("reverts on an unknown sessionId", async function () {
      await expect(did.authorize(999, scopeHash, ethers.ZeroHash, 1n, "0x" + "00".repeat(65)))
        .to.be.revertedWithCustomError(did, "UnknownSession");
    });
  });

  // -----------------------------------------------------------------------
  // revocation & expiry & root-lifecycle cross-check
  // -----------------------------------------------------------------------
  describe("session validity: revoke / expiry / root lifecycle", function () {
    let root, session, sid, scopeHash;
    beforeEach(async function () {
      const rng = makeRng("valid-setup");
      root = await registerFalconRoot(1, owner, rng);
      await did.createAgent(root.rootKeyId);
      session = ethers.Wallet.createRandom();
      scopeHash = ethers.keccak256(ethers.toUtf8Bytes("s"));
      sid = await issueSession(root, {
        sessionAddr: session.address,
        scopeHash,
        spendCap: 100n,
        expiry: await futureTs(1000),
        rng,
      });
      this.rng = rng;
    });

    async function authorizeOnce(amount) {
      const actionHash = ethers.keccak256(ethers.toUtf8Bytes("act"));
      const nonce = (await did.getSession(sid)).actionNonce;
      const digest = await did.actionDigest(sid, scopeHash, actionHash, amount, nonce);
      return did.authorize(sid, scopeHash, actionHash, amount, ecdsaSign(session, digest));
    }

    it("controller can fast-revoke (cheap ECDSA); a stranger cannot", async function () {
      await expect(did.connect(other).revokeSession(sid)).to.be.revertedWithCustomError(did, "NotSessionRevoker");
      await expect(did.connect(owner).revokeSession(sid)).to.emit(did, "SessionRevoked");
      expect(await did.isSessionValid(sid)).to.equal(false);
      await expect(authorizeOnce(1n)).to.be.revertedWithCustomError(did, "SessionExpiredOrInvalid");
      // revoke is idempotent (no-op success on a second call).
      await did.connect(owner).revokeSession(sid);
    });

    it("a session past its expiry is invalid", async function () {
      expect(await did.isSessionValid(sid)).to.equal(true);
      await network.provider.send("evm_increaseTime", [2000]);
      await network.provider.send("evm_mine", []);
      expect(await did.isSessionValid(sid)).to.equal(false);
      await expect(authorizeOnce(1n)).to.be.revertedWithCustomError(did, "SessionExpiredOrInvalid");
    });

    it("REVOKING the Falcon root invalidates all its sessions", async function () {
      expect(await did.isSessionValid(sid)).to.equal(true);
      await reg.connect(owner).revokeKey(root.rootKeyId);
      expect(await did.isSessionValid(sid)).to.equal(false);
      await expect(authorizeOnce(1n)).to.be.revertedWithCustomError(did, "SessionExpiredOrInvalid");
    });

    it("ROTATING the Falcon root (root no longer ACTIVE) invalidates its sessions", async function () {
      const rng = this.rng;
      const np = makePubKey(1, rng);
      const n = await reg.identityNonce(owner.address);
      const ch = await reg.popChallenge(owner.address, 1, np, n);
      await reg.rotateKey(root.rootKeyId, 1, np, genuine(1, np, ch, rng));
      expect(await did.isSessionValid(sid)).to.equal(false);
    });
  });

  // -----------------------------------------------------------------------
  // Seeded invariant fuzz
  // -----------------------------------------------------------------------
  describe("seeded invariant fuzz", function () {
    it("scope + spend-cap + anti-replay + revocation hold across random rounds", async function () {
      const rng = makeRng("fuzz-did-v1");
      const scheme = 2;
      const root = await registerFalconRoot(scheme, owner, rng);
      await did.createAgent(root.rootKeyId);

      const ROUNDS = 24;
      for (let r = 0; r < ROUNDS; r++) {
        const session = ethers.Wallet.createRandom();
        const scopeHash = ethers.hexlify(rng.bytes(32));
        const cap = BigInt(50 + rng.int(200));
        const sid = await issueSession(root, {
          sessionAddr: session.address,
          scopeHash,
          spendCap: cap,
          expiry: await futureTs(5000),
          rng,
        });

        let spent = 0n;
        // A few genuine scoped actions within cap succeed and accumulate.
        const k = 1 + rng.int(3);
        for (let i = 0; i < k; i++) {
          const remaining = cap - spent;
          if (remaining === 0n) break;
          const amount = BigInt(1 + rng.int(Number(remaining > 40n ? 40n : remaining)));
          const actionHash = ethers.hexlify(rng.bytes(32));
          const nonce = (await did.getSession(sid)).actionNonce;
          const digest = await did.actionDigest(sid, scopeHash, actionHash, amount, nonce);
          await did.authorize(sid, scopeHash, actionHash, amount, ecdsaSign(session, digest));
          spent += amount;
          expect((await did.getSession(sid)).spent).to.equal(spent);
        }

        // (a) wrong scope always reverts.
        {
          const badScope = ethers.hexlify(rng.bytes(32));
          const actionHash = ethers.hexlify(rng.bytes(32));
          const nonce = (await did.getSession(sid)).actionNonce;
          const digest = await did.actionDigest(sid, badScope, actionHash, 1n, nonce);
          await expect(
            did.authorize(sid, badScope, actionHash, 1n, ecdsaSign(session, digest))
          ).to.be.revertedWithCustomError(did, "SessionScopeViolation");
        }
        // (b) exceeding the cap always reverts.
        {
          const over = cap - spent + 1n;
          const actionHash = ethers.hexlify(rng.bytes(32));
          const nonce = (await did.getSession(sid)).actionNonce;
          const digest = await did.actionDigest(sid, scopeHash, actionHash, over, nonce);
          await expect(
            did.authorize(sid, scopeHash, actionHash, over, ecdsaSign(session, digest))
          ).to.be.revertedWithCustomError(did, "SessionSpendExceeded");
        }
        // (c) a stranger key never authorizes. Use amount 0 so the spend check (which
        // runs before the signature check) always passes, isolating the SIG failure —
        // the session's cap may already be exhausted this round.
        {
          const stranger = ethers.Wallet.createRandom();
          const actionHash = ethers.hexlify(rng.bytes(32));
          const nonce = (await did.getSession(sid)).actionNonce;
          const digest = await did.actionDigest(sid, scopeHash, actionHash, 0n, nonce);
          await expect(
            did.authorize(sid, scopeHash, actionHash, 0n, ecdsaSign(stranger, digest))
          ).to.be.revertedWithCustomError(did, "SessionSigInvalid");
        }

        // Randomly revoke ~1/3 of sessions; a revoked session never authorizes again.
        if (rng.int(3) === 0) {
          await did.connect(owner).revokeSession(sid);
          expect(await did.isSessionValid(sid)).to.equal(false);
          const actionHash = ethers.hexlify(rng.bytes(32));
          const nonce = (await did.getSession(sid)).actionNonce;
          const digest = await did.actionDigest(sid, scopeHash, actionHash, 1n, nonce);
          await expect(
            did.authorize(sid, scopeHash, actionHash, 1n, ecdsaSign(session, digest))
          ).to.be.revertedWithCustomError(did, "SessionExpiredOrInvalid");
        }
      }

      // sessionsOf reflects everything issued.
      const list = await did.sessionsOf(root.rootKeyId);
      expect(list.length).to.equal(ROUNDS);
      expect(Number(await did.sessionCount())).to.equal(ROUNDS);
    });
  });
});