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

describe("AereInferNet", () => {
  let provider, regulator, outsider;
  let net;
  const MODEL_ID = ethers.keccak256(ethers.toUtf8Bytes("claude-4.7-sonnet-20251022"));
  const MODEL_DIGEST = ethers.keccak256(ethers.toUtf8Bytes("model-weights-blob"));
  const ATT_DIGEST = ethers.keccak256(ethers.toUtf8Bytes("sgx-attestation-quote"));

  beforeEach(async () => {
    [provider, regulator, outsider] = await ethers.getSigners();
    const F = await ethers.getContractFactory("AereInferNet");
    net = await F.deploy();
    await net.waitForDeployment();
  });

  function buildLeaf(req, resp, ts, deployer) {
    return ethers.keccak256(ethers.AbiCoder.defaultAbiCoder().encode(
      ["bytes32", "bytes32", "uint64", "bytes32"], [req, resp, ts, deployer]
    ));
  }

  // tiny sort-pair Merkle: builds a 4-leaf tree and returns root + a proof for leaf 0
  function tinyTree(leaves) {
    const hashPair = (a, b) => {
      const [lo, hi] = BigInt(a) <= BigInt(b) ? [a, b] : [b, a];
      return ethers.keccak256(ethers.concat([lo, hi]));
    };
    while (leaves.length > 1) {
      const next = [];
      for (let i = 0; i < leaves.length; i += 2) {
        if (i + 1 < leaves.length) next.push(hashPair(leaves[i], leaves[i + 1]));
        else next.push(leaves[i]);
      }
      leaves = next;
    }
    return leaves[0];
  }

  function buildTreeWithProof(allLeaves, idx) {
    // returns { root, proof } for the leaf at idx in a sorted-pair tree.
    const hashPair = (a, b) => {
      const [lo, hi] = BigInt(a) <= BigInt(b) ? [a, b] : [b, a];
      return ethers.keccak256(ethers.concat([lo, hi]));
    };
    let layer = allLeaves.slice();
    let proof = [];
    let cur = idx;
    while (layer.length > 1) {
      const sibling = cur ^ 1;
      if (sibling < layer.length) proof.push(layer[sibling]);
      const next = [];
      for (let i = 0; i < layer.length; i += 2) {
        if (i + 1 < layer.length) next.push(hashPair(layer[i], layer[i + 1]));
        else next.push(layer[i]);
      }
      layer = next;
      cur = Math.floor(cur / 2);
    }
    return { root: layer[0], proof };
  }

  it("registers a model + reads it back", async () => {
    await expect(
      net.connect(provider).registerModel(
        MODEL_ID, MODEL_DIGEST, ATT_DIGEST, 2, "claude-4.7-sonnet-20251022", "ipfs://attestation"
      )
    ).to.emit(net, "ModelRegistered");
    const m = await net.modelOf(provider.address, MODEL_ID);
    expect(m.exists).to.equal(true);
    expect(m.modelDigest).to.equal(MODEL_DIGEST);
    expect(m.risk).to.equal(2);
  });

  it("rejects duplicate model registration + malformed inputs", async () => {
    await net.connect(provider).registerModel(MODEL_ID, MODEL_DIGEST, ATT_DIGEST, 2, "n", "u");
    await expect(
      net.connect(provider).registerModel(MODEL_ID, MODEL_DIGEST, ATT_DIGEST, 2, "n", "u")
    ).to.be.revertedWithCustomError(net, "ModelExists");
    await expect(
      net.connect(outsider).registerModel(MODEL_ID, ethers.ZeroHash, ATT_DIGEST, 2, "n", "u")
    ).to.be.revertedWithCustomError(net, "ZeroDigest");
    await expect(
      net.connect(outsider).registerModel(MODEL_ID, MODEL_DIGEST, ATT_DIGEST, 2, "", "u")
    ).to.be.revertedWithCustomError(net, "EmptyName");
  });

  it("commits an epoch only when model is registered", async () => {
    const root = ethers.keccak256(ethers.toUtf8Bytes("root-x"));
    const today = BigInt(Math.floor((await ethers.provider.getBlock("latest")).timestamp / 86400));
    await expect(
      net.connect(provider).commitEpoch(MODEL_ID, today, root, 100, "ipfs://meta")
    ).to.be.revertedWithCustomError(net, "ModelUnknown");
    await net.connect(provider).registerModel(MODEL_ID, MODEL_DIGEST, ATT_DIGEST, 2, "n", "u");
    await expect(
      net.connect(provider).commitEpoch(MODEL_ID, today, root, 100, "ipfs://meta")
    ).to.emit(net, "EpochCommitted").withArgs(provider.address, MODEL_ID, today, root, 100, "ipfs://meta");
    const e = await net.epochOf(provider.address, MODEL_ID, today);
    expect(e.root).to.equal(root);
    expect(e.inferenceCount).to.equal(100n);
    expect(await net.modelEpochCount(provider.address, MODEL_ID)).to.equal(1n);
  });

  it("rejects re-committing the same epoch + zero root + out-of-range epoch (audit fix HIGH #26)", async () => {
    const root = ethers.keccak256(ethers.toUtf8Bytes("root"));
    const today = BigInt(Math.floor((await ethers.provider.getBlock("latest")).timestamp / 86400));
    await net.connect(provider).registerModel(MODEL_ID, MODEL_DIGEST, ATT_DIGEST, 2, "n", "u");
    // Old epoch (>1 day in past) rejected.
    await expect(
      net.connect(provider).commitEpoch(MODEL_ID, 1n, root, 1, "")
    ).to.be.revertedWithCustomError(net, "EpochOutOfRange");
    // Future epoch rejected.
    await expect(
      net.connect(provider).commitEpoch(MODEL_ID, today + 5n, root, 1, "")
    ).to.be.revertedWithCustomError(net, "EpochOutOfRange");
    await expect(
      net.connect(provider).commitEpoch(MODEL_ID, today, ethers.ZeroHash, 1, "")
    ).to.be.revertedWithCustomError(net, "ZeroDigest");
    await net.connect(provider).commitEpoch(MODEL_ID, today, root, 1, "");
    await expect(
      net.connect(provider).commitEpoch(MODEL_ID, today, root, 1, "")
    ).to.be.revertedWithCustomError(net, "EpochExists");
  });

  it("verifyInference proves inclusion against committed root", async () => {
    await net.connect(provider).registerModel(MODEL_ID, MODEL_DIGEST, ATT_DIGEST, 2, "n", "u");

    const leaves = [];
    const targets = [];
    for (let i = 0; i < 4; i++) {
      const req = ethers.keccak256(ethers.toUtf8Bytes(`req-${i}`));
      const resp = ethers.keccak256(ethers.toUtf8Bytes(`resp-${i}`));
      const ts = BigInt(1717880000000 + i * 1000);
      const dep = ethers.keccak256(ethers.toUtf8Bytes(`deployer-${i}`));
      leaves.push(buildLeaf(req, resp, ts, dep));
      targets.push({ req, resp, ts, dep });
    }

    const idx = 2;
    const { root, proof } = buildTreeWithProof(leaves, idx);
    const today = BigInt(Math.floor((await ethers.provider.getBlock("latest")).timestamp / 86400));
    await net.connect(provider).commitEpoch(MODEL_ID, today, root, 4, "");

    const t = targets[idx];
    expect(await net.verifyInference(provider.address, MODEL_ID, today, t.req, t.resp, t.ts, t.dep, proof))
      .to.equal(true);

    // Wrong leaf → false.
    expect(await net.verifyInference(
      provider.address, MODEL_ID, today,
      ethers.keccak256(ethers.toUtf8Bytes("not-req")), t.resp, t.ts, t.dep, proof
    )).to.equal(false);

    // Unknown epoch → false (no revert).
    expect(await net.verifyInference(provider.address, MODEL_ID, 999999n, t.req, t.resp, t.ts, t.dep, proof))
      .to.equal(false);
  });
});