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

// AUDIT FINDING 2 — RESIDUAL variant (pause-then-clawback).
//
// Iteration 1 closed the setExcluded-after-claim confiscation, but the review
// found a second path to the SAME confiscation: claim() is whenNotPaused while
// clawbackUnclaimed is NOT, and clawback swept address(this).balance with no
// accounting of outstanding entitlements. So the owner could pause(), wait past
// claimExpiry, and sweep an in-progress claimant's vested-but-unclaimed AERE.
//
// This suite reproduces the pause-then-clawback confiscation and proves the
// refined AereRewardDistributorV2 closes it two ways at once:
//   1. clawbackUnclaimed is now whenNotPaused, so the owner cannot clawback
//      while the campaign is paused; and
//   2. clawbackUnclaimed is capped to genuinely-unallocated funds only
//      (balance - (totalAllocated - totalClaimed)), so even unpaused it can
//      never touch a claimant's still-owed entitlement.
// It also proves the legit flows still work: dust above the allocation is still
// recoverable after expiry, and an emergency pause still halts claims.

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;

async function deployV2(waere, saere, sink) {
  return (await ethers.getContractFactory("AereRewardDistributorV2")).deploy(
    await waere.getAddress(),
    await saere.getAddress(),
    await sink.getAddress()
  );
}

describe("AERE Season 1 — Distributor pause-then-clawback (FINDING 2 residual)", () => {
  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 RESIDUAL ATTACK: owner pauses, waits past expiry, then clawbacks the
  // in-progress claimant's vested-but-unclaimed AERE. Must be fully blocked.
  it("owner CANNOT pause-then-clawback to confiscate an in-progress claim", async () => {
    const dist = await deployV2(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;
    // totalAllocated = 150 (Alice 100 + Bob 50).
    await dist.publishRoot(tree.root, 0, 1000, expiry, false, E("150"));
    await owner.sendTransaction({ to: await dist.getAddress(), value: E("150") });

    // Alice claims half (50 of 100) at mid-vest; she is a claim IN PROGRESS.
    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: pause, wait past expiry, sweep the balance.
    await dist.pause();
    await time.setNextBlockTimestamp(expiry + 1);
    // (1) Blocked while paused.
    await expect(dist.clawbackUnclaimed(false)).to.be.revertedWith("Pausable: paused");

    await dist.unpause();
    // (2) Even unpaused, everything left is reserved for Alice(50)+Bob(50);
    //     nothing is genuinely unallocated, so clawback recovers nothing.
    await expect(dist.clawbackUnclaimed(false)).to.be.revertedWithCustomError(
      dist,
      "NothingToClawback"
    );

    // Alice's remaining vested AERE is untouched and still claimable.
    await time.setNextBlockTimestamp(expiry + 100);
    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 (a): dust ABOVE the allocation is still recoverable after expiry,
  // even while a claimant still has an outstanding reserved entitlement.
  it("still allows post-expiry clawback of genuinely-unallocated dust (reserve protected)", async () => {
    const dist = await deployV2(waere, saere, sink);
    const tree = makeTree([{ account: alice.address, total: E("100") }]);
    const now = await time.latest();
    const expiry = now + 5000;
    // totalAllocated = 100, but fund 110 -> 10 AERE of overfunded dust.
    await dist.publishRoot(tree.root, 0, 1000, expiry, false, E("100"));
    await owner.sendTransaction({ to: await dist.getAddress(), value: E("110") });

    // Alice claims half (50); reserved remainder = 100 - 50 = 50.
    await time.increaseTo(now + 500);
    await dist.claim(alice.address, E("100"), tree.proofFor(alice.address));
    expect(await dist.claimed(alice.address)).to.equal(E("50"));

    // After expiry, only the 10 dust is clawable; Alice's 50 stays put.
    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"));
    // Contract still holds Alice's reserved 50.
    expect(await ethers.provider.getBalance(await dist.getAddress())).to.equal(E("50"));

    // A second clawback recovers nothing (all remaining is reserved).
    await expect(dist.clawbackUnclaimed(true)).to.be.revertedWithCustomError(
      dist,
      "NothingToClawback"
    );

    // Alice can still finish her claim after expiry.
    await time.increaseTo(expiry + 2000);
    const aBefore = await ethers.provider.getBalance(alice.address);
    await dist.connect(relayer).claim(alice.address, E("100"), tree.proofFor(alice.address));
    const aAfter = await ethers.provider.getBalance(alice.address);
    expect(aAfter - aBefore).to.equal(E("50"));
  });

  // LEGIT (b): emergency pause still halts claims, and unpausing restores them.
  it("still allows emergency pause to halt and resume claims", async () => {
    const dist = await deployV2(waere, saere, sink);
    const tree = makeTree([{ account: alice.address, total: E("100") }]);
    const now = await time.latest();
    await dist.publishRoot(tree.root, 0, 10, now + 10_000_000, false, E("100"));
    await owner.sendTransaction({ to: await dist.getAddress(), value: E("100") });

    await dist.pause();
    await time.increase(20);
    await expect(
      dist.claim(alice.address, E("100"), tree.proofFor(alice.address))
    ).to.be.revertedWith("Pausable: paused");

    await dist.unpause();
    await dist.claim(alice.address, E("100"), tree.proofFor(alice.address));
    expect(await dist.claimed(alice.address)).to.equal(E("100"));
  });
});