const { expect } = require("chai");
const { ethers } = require("hardhat");
const { time } = require("@nomicfoundation/hardhat-network-helpers");

// AUDIT FINDING 2 (medium) — Season 1 distributor confiscation.
// setExcluded has no guard against excluding an account that has ALREADY
// partially claimed. Owner excludes such an account -> claim() then reverts
// IsExcluded -> clawbackUnclaimed sweeps the ENTIRE remaining balance
// (including that user's vested-but-unclaimed AERE) after expiry. That
// contradicts the advertised invariant that the owner CANNOT redirect a
// valid, vested claim away from its rightful account.
//
// This suite reproduces the confiscation against the flawed contract and
// proves AereRewardDistributorV2 closes it while preserving the legit uses:
//   (a) excluding a never-claimed account still works, and
//   (b) legit clawback of truly-unallocated funds still works.

const abi = ethers.AbiCoder.defaultAbiCoder();

/* ----------------------------- merkle helpers ----------------------------- */
function leafHash(account, total) {
  const inner = ethers.keccak256(abi.encode(["address", "uint256"], [account, total]));
  return ethers.keccak256(inner);
}
function hashPair(a, b) {
  return BigInt(a) < BigInt(b)
    ? ethers.keccak256(ethers.concat([a, b]))
    : ethers.keccak256(ethers.concat([b, a]));
}
function buildLayers(leaves) {
  const layers = [leaves.slice()];
  while (layers[layers.length - 1].length > 1) {
    const prev = layers[layers.length - 1];
    const next = [];
    for (let i = 0; i < prev.length; i += 2) {
      if (i + 1 < prev.length) next.push(hashPair(prev[i], prev[i + 1]));
      else next.push(prev[i]);
    }
    layers.push(next);
  }
  return layers;
}
function getProof(layers, index) {
  const proof = [];
  let idx = index;
  for (let l = 0; l < layers.length - 1; l++) {
    const layer = layers[l];
    const sib = idx ^ 1;
    if (sib < layer.length) proof.push(layer[sib]);
    idx = Math.floor(idx / 2);
  }
  return proof;
}
function makeTree(entries) {
  const leaves = entries.map((e) => leafHash(e.account, e.total));
  const layers = buildLayers(leaves);
  const root = layers[layers.length - 1][0];
  const indexOf = {};
  entries.forEach((e, i) => (indexOf[e.account.toLowerCase()] = i));
  return { root, proofFor: (account) => getProof(layers, indexOf[account.toLowerCase()]) };
}

const E = ethers.parseEther;

// Deploy helper parameterized by contract name so we can point the same
// scenario at the flawed V1 and the corrected V2.
async function deployDist(name, waere, saere, sink) {
  return (await ethers.getContractFactory(name)).deploy(
    await waere.getAddress(),
    await saere.getAddress(),
    await sink.getAddress()
  );
}

describe("AERE Season 1 — Distributor confiscation (FINDING 2)", () => {
  let waere, saere, sink, owner, alice, bob, relayer;

  beforeEach(async () => {
    [owner, alice, bob, relayer] = await ethers.getSigners();
    waere = await (await ethers.getContractFactory("WAERE")).deploy();
    saere = await (await ethers.getContractFactory("MockSAERE")).deploy(await waere.getAddress());
    sink = await (await ethers.getContractFactory("MockSink")).deploy();
  });

  // The core attack: a partial claimer must NOT be excludable, so the owner
  // can never freeze-then-clawback their vested-but-unclaimed AERE.
  it("owner CANNOT exclude an account that has already claimed (confiscation blocked)", async () => {
    const dist = await deployDist("AereRewardDistributorV2", waere, saere, sink);
    const tree = makeTree([
      { account: alice.address, total: E("100") },
      { account: bob.address, total: E("50") },
    ]);
    const now = await time.latest();
    const expiry = now + 10_000_000;
    await dist.publishRoot(tree.root, 0, 1000, expiry, false, E("150"));
    await owner.sendTransaction({ to: await dist.getAddress(), value: E("150") });

    // Alice partially claims at half vest (50 of 100).
    const start = Number(await dist.rootPublishedAt());
    await time.setNextBlockTimestamp(start + 500);
    await dist.connect(relayer).claim(alice.address, E("100"), tree.proofFor(alice.address));
    expect(await dist.claimed(alice.address)).to.equal(E("50"));

    // The confiscation weapon: excluding a partial claimer. Must revert.
    await expect(dist.setExcluded(alice.address, true)).to.be.revertedWithCustomError(
      dist,
      "AlreadyClaimed"
    );

    // Alice can therefore still claim her remaining vested AERE.
    await time.setNextBlockTimestamp(start + 1000 + 10);
    const balBefore = await ethers.provider.getBalance(alice.address);
    await dist.connect(relayer).claim(alice.address, E("100"), tree.proofFor(alice.address));
    const balAfter = await ethers.provider.getBalance(alice.address);
    expect(balAfter - balBefore).to.equal(E("50"));
    expect(await dist.claimed(alice.address)).to.equal(E("100"));
  });

  // Legit use (a): excluding a never-claimed sybil still works (and re-include).
  it("still allows excluding a never-claimed account (sybil revocation)", async () => {
    const dist = await deployDist("AereRewardDistributorV2", waere, saere, sink);
    const tree = makeTree([
      { account: alice.address, total: E("100") },
      { account: bob.address, total: E("50") },
    ]);
    const now = await time.latest();
    await dist.publishRoot(tree.root, 0, 10, now + 10_000_000, false, E("150"));
    await owner.sendTransaction({ to: await dist.getAddress(), value: E("150") });

    // bob never claimed -> exclude allowed.
    await dist.setExcluded(bob.address, true);
    await time.increase(20);
    await expect(
      dist.claim(bob.address, E("50"), tree.proofFor(bob.address))
    ).to.be.revertedWithCustomError(dist, "IsExcluded");

    // re-include is always permitted, and then bob can claim.
    await dist.setExcluded(bob.address, false);
    await dist.claim(bob.address, E("50"), tree.proofFor(bob.address));
    expect(await dist.claimed(bob.address)).to.equal(E("50"));
  });

  // Legit use (b): after expiry, clawback of genuinely-unallocated funds (the
  // balance beyond every leaf's reserved entitlement) still works. Here Alice
  // fully claims her 100 and the distributor was overfunded by 10 of dust; only
  // that 10 is recoverable, never a claimant's owed AERE.
  it("still allows legit post-expiry clawback of unallocated dust", async () => {
    const dist = await deployDist("AereRewardDistributorV2", waere, saere, sink);
    const tree = makeTree([{ account: alice.address, total: E("100") }]);
    const now = await time.latest();
    const expiry = now + 5000;
    await dist.publishRoot(tree.root, 0, 1000, expiry, false, E("100"));
    // Overfund by 10 AERE of dust no leaf is entitled to.
    await owner.sendTransaction({ to: await dist.getAddress(), value: E("110") });

    // Alice claims her full entitlement (past full vest).
    await time.increaseTo(now + 1000 + 10);
    await dist.claim(alice.address, E("100"), tree.proofFor(alice.address));
    expect(await dist.claimed(alice.address)).to.equal(E("100"));

    // After expiry, reserve is 0 (fully claimed); only the 10 dust is clawable.
    await time.increaseTo(expiry + 1);
    const sinkBefore = await ethers.provider.getBalance(await sink.getAddress());
    await dist.clawbackUnclaimed(true);
    const sinkAfter = await ethers.provider.getBalance(await sink.getAddress());
    expect(sinkAfter - sinkBefore).to.equal(E("10"));
    expect(await ethers.provider.getBalance(await dist.getAddress())).to.equal(0);
  });
});