// AereZKScreen + AereAIProof tests.

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

async function deployZkStack() {
  const [deployer, alice, bob] = await ethers.getSigners();
  const Verifier = await ethers.getContractFactory("MockSp1Verifier");
  const verifier = await Verifier.deploy();
  await verifier.waitForDeployment();
  const ZK = await ethers.getContractFactory("AereZKScreen");
  const zk = await ZK.deploy(await verifier.getAddress());
  await zk.waitForDeployment();
  return { deployer, alice, bob, verifier, zk };
}

// Build publicValues + pre-program the void mock to accept the exact triple.
async function makeProof(verifier, vKey, user, attestedAt, extraDataHash, proof = "0xdeadbeef") {
  const pv = ethers.AbiCoder.defaultAbiCoder().encode(
    ["uint256", "address", "uint256", "bytes32"],
    [(await ethers.provider.getNetwork()).chainId, user, attestedAt, extraDataHash]
  );
  const marker = ethers.keccak256(ethers.AbiCoder.defaultAbiCoder().encode(
    ["bytes32", "bytes", "bytes"], [vKey, pv, proof]
  ));
  await verifier.setValid(marker, true);
  return { pv, proof };
}

describe("AereZKScreen", function () {
  const ROOT = ethers.id("authorized-allowlist-root-v1");

  it("Foundation registers program; duplicate rejected; zero root rejected", async function () {
    const { zk } = await deployZkStack();
    const vKey = ethers.id("OFAC-screen-v1");
    await zk.registerProgram(vKey, 90 * 86400, ROOT, "OFAC SDN screen v1");
    await expect(zk.registerProgram(vKey, 0, ROOT, "")).to.be.revertedWithCustomError(zk, "AlreadyRegistered");
    const vKey2 = ethers.id("zero-root");
    await expect(zk.registerProgram(vKey2, 0, ethers.ZeroHash, "")).to.be.revertedWithCustomError(zk, "ZeroRoot");
  });

  it("submits valid proof (root == authorizedRoot); isCleared true", async function () {
    const { alice, verifier, zk } = await deployZkStack();
    const vKey = ethers.id("OFAC-screen-v1");
    await zk.registerProgram(vKey, 90 * 86400, ROOT, "OFAC");
    const attestedAt = (await ethers.provider.getBlock("latest")).timestamp - 60;
    const { pv, proof } = await makeProof(verifier, vKey, alice.address, attestedAt, ROOT);
    await zk.connect(alice).submitProof(vKey, pv, proof);
    expect(await zk.isCleared(alice.address, vKey, 0)).to.equal(true);
  });

  it("rejects proof whose root != authorizedRoot (self-clear attack)", async function () {
    const { alice, verifier, zk } = await deployZkStack();
    const vKey = ethers.id("OFAC-screen-v1");
    await zk.registerProgram(vKey, 0, ROOT, "OFAC");
    // Attacker proves membership in a root of their OWN making.
    const attackerRoot = ethers.id("attacker-picked-root");
    const attestedAt = (await ethers.provider.getBlock("latest")).timestamp - 60;
    const { pv, proof } = await makeProof(verifier, vKey, alice.address, attestedAt, attackerRoot);
    await expect(zk.connect(alice).submitProof(vKey, pv, proof))
      .to.be.revertedWithCustomError(zk, "UnauthorizedRoot");
    expect(await zk.isCleared(alice.address, vKey, 0)).to.equal(false);
  });

  it("rejects proof for wrong user", async function () {
    const { alice, bob, verifier, zk } = await deployZkStack();
    const vKey = ethers.id("OFAC-screen-v1");
    await zk.registerProgram(vKey, 0, ROOT, "OFAC");
    const attestedAt = (await ethers.provider.getBlock("latest")).timestamp - 60;
    const { pv, proof } = await makeProof(verifier, vKey, bob.address, attestedAt, ROOT);
    await expect(zk.connect(alice).submitProof(vKey, pv, proof))
      .to.be.revertedWithCustomError(zk, "UserMismatch");
  });

  it("rejects invalid proof", async function () {
    const { alice, zk } = await deployZkStack();
    const vKey = ethers.id("v");
    await zk.registerProgram(vKey, 0, ROOT, "x");
    const pv = ethers.AbiCoder.defaultAbiCoder().encode(
      ["uint256", "address", "uint256", "bytes32"],
      [(await ethers.provider.getNetwork()).chainId, alice.address, 100, ROOT]
    );
    await expect(zk.connect(alice).submitProof(vKey, pv, "0x"))
      .to.be.revertedWithCustomError(zk, "InvalidProof");
  });

  it("expiry causes isCleared to flip false", async function () {
    const { alice, verifier, zk } = await deployZkStack();
    const vKey = ethers.id("expiry-test");
    await zk.registerProgram(vKey, 100, ROOT, "short");
    const ts = (await ethers.provider.getBlock("latest")).timestamp - 50;
    const { pv, proof } = await makeProof(verifier, vKey, alice.address, ts, ROOT, "0x01");
    await zk.connect(alice).submitProof(vKey, pv, proof);
    expect(await zk.isCleared(alice.address, vKey, 0)).to.equal(true);
    await ethers.provider.send("evm_increaseTime", [200]);
    await ethers.provider.send("evm_mine", []);
    expect(await zk.isCleared(alice.address, vKey, 0)).to.equal(false);
  });

  it("setAuthorizedRoot rotates the accepted root", async function () {
    const { alice, verifier, zk } = await deployZkStack();
    const vKey = ethers.id("rotate");
    await zk.registerProgram(vKey, 0, ROOT, "rot");
    const newRoot = ethers.id("rotated-root");
    // Proof against the OLD root now fails after rotation.
    await zk.setAuthorizedRoot(vKey, newRoot);
    const attestedAt = (await ethers.provider.getBlock("latest")).timestamp - 30;
    const old = await makeProof(verifier, vKey, alice.address, attestedAt, ROOT);
    await expect(zk.connect(alice).submitProof(vKey, old.pv, old.proof))
      .to.be.revertedWithCustomError(zk, "UnauthorizedRoot");
    // Proof against the NEW root succeeds.
    const fresh = await makeProof(verifier, vKey, alice.address, attestedAt, newRoot);
    await zk.connect(alice).submitProof(vKey, fresh.pv, fresh.proof);
    expect(await zk.isCleared(alice.address, vKey, 0)).to.equal(true);
    await expect(zk.setAuthorizedRoot(vKey, ethers.ZeroHash)).to.be.revertedWithCustomError(zk, "ZeroRoot");
  });

  it("revoke voids a clearance; a fresh newer proof re-clears", async function () {
    const { alice, verifier, zk } = await deployZkStack();
    const vKey = ethers.id("revoke-test");
    await zk.registerProgram(vKey, 0, ROOT, "rev");
    const t1 = (await ethers.provider.getBlock("latest")).timestamp - 100;
    const p1 = await makeProof(verifier, vKey, alice.address, t1, ROOT);
    await zk.connect(alice).submitProof(vKey, p1.pv, p1.proof);
    expect(await zk.isCleared(alice.address, vKey, 0)).to.equal(true);

    await zk.revoke(alice.address, vKey);
    expect(await zk.isCleared(alice.address, vKey, 0)).to.equal(false);

    // A fresh proof with attestedAt strictly after the revocation re-clears.
    await ethers.provider.send("evm_increaseTime", [10]);
    await ethers.provider.send("evm_mine", []);
    const t2 = (await ethers.provider.getBlock("latest")).timestamp;
    const p2 = await makeProof(verifier, vKey, alice.address, t2, ROOT);
    await zk.connect(alice).submitProof(vKey, p2.pv, p2.proof);
    expect(await zk.isCleared(alice.address, vKey, 0)).to.equal(true);
  });

  it("only owner can register, rotate, revoke", async function () {
    const { alice, zk } = await deployZkStack();
    const vKey = ethers.id("owner-test");
    await expect(zk.connect(alice).registerProgram(vKey, 0, ROOT, "x"))
      .to.be.revertedWith("Ownable: caller is not the owner");
    await zk.registerProgram(vKey, 0, ROOT, "x");
    await expect(zk.connect(alice).setAuthorizedRoot(vKey, ethers.id("y")))
      .to.be.revertedWith("Ownable: caller is not the owner");
    await expect(zk.connect(alice).revoke(alice.address, vKey))
      .to.be.revertedWith("Ownable: caller is not the owner");
  });
});

describe("AereAIProof", function () {
  async function deploy() {
    const [deployer, signer, operator, requester] = await ethers.getSigners();
    const AI = await ethers.getContractFactory("AereAIProof");
    const ai = await AI.deploy();
    await ai.waitForDeployment();
    return { deployer, signer, operator, requester, ai };
  }

  async function signAttestation(wallet, ai, attest) {
    const domain = {
      name: "AereAIProof", version: "1",
      chainId: (await ethers.provider.getNetwork()).chainId,
      verifyingContract: await ai.getAddress(),
    };
    const types = {
      Attestation: [
        { name: "modelId", type: "bytes32" },
        { name: "inputHash", type: "bytes32" },
        { name: "outputHash", type: "bytes32" },
        { name: "requesterHash", type: "bytes32" },
        { name: "timestamp", type: "uint256" },
        { name: "nonce", type: "uint256" },
      ],
    };
    return wallet.signTypedData(domain, types, attest);
  }

  it("operator registers model; signer rotation works", async function () {
    const { operator, signer, requester, ai } = await deploy();
    const modelId = ethers.id("claude-opus-4.7");
    await ai.connect(operator).registerModel(modelId, signer.address, "Claude Opus 4.7");
    const m = await ai.models(modelId);
    expect(m.signer).to.equal(signer.address);
    expect(m.operator).to.equal(operator.address);
    await ai.connect(operator).setSigner(modelId, requester.address);
    expect((await ai.models(modelId)).signer).to.equal(requester.address);
  });

  it("anchors a valid attestation; rejects replay + bad-sig + wrong-chain timestamp", async function () {
    const { operator, signer, requester, ai } = await deploy();
    const modelId = ethers.id("claude-opus-4.7");
    await ai.connect(operator).registerModel(modelId, signer.address, "Claude Opus 4.7");

    const inputHash    = ethers.id("input-X");
    const outputHash   = ethers.id("output-Y");
    const requesterHash= ethers.id("user-Z-nonce-1");
    const ts = (await ethers.provider.getBlock("latest")).timestamp - 10;
    const nonce = 1n;
    const attest = { modelId, inputHash, outputHash, requesterHash, timestamp: ts, nonce };

    const sig = await signAttestation(signer, ai, attest);
    await expect(ai.connect(requester).anchor(modelId, inputHash, outputHash, requesterHash, ts, nonce, sig))
      .to.emit(ai, "Anchored");
    expect(await ai.attestationCount(modelId)).to.equal(1n);

    // Replay rejected.
    await expect(ai.connect(requester).anchor(modelId, inputHash, outputHash, requesterHash, ts, nonce, sig))
      .to.be.revertedWithCustomError(ai, "NonceConsumed");

    // Wrong signer rejected.
    const badSig = await signAttestation(operator, ai, { ...attest, nonce: 2n });
    await expect(ai.connect(requester).anchor(modelId, inputHash, outputHash, requesterHash, ts, 2n, badSig))
      .to.be.revertedWithCustomError(ai, "BadSignature");

    // Future timestamp rejected.
    const future = (await ethers.provider.getBlock("latest")).timestamp + 3600;
    const futureSig = await signAttestation(signer, ai, { ...attest, nonce: 3n, timestamp: future });
    await expect(ai.connect(requester).anchor(modelId, inputHash, outputHash, requesterHash, future, 3n, futureSig))
      .to.be.revertedWithCustomError(ai, "AttestationInFuture");
  });
});