// Proof suite for the paymaster V2 bug-fixes (roadmap #8 follow-up).
//
// Proves — against the ACTUAL AereStaking + AereLockedStaking sources (byte-
// identical to the live chain-2800 deployments, verified on-chain: real
// AereStaking has NO balanceOf, real AereLockedStaking.getLock returns amount
// first) — that:
//
//   AereStakeQuotaPaymasterV2
//     * reads a REAL AereLockedStaking lock correctly where V1 hard-REVERTS
//     * reads a REAL AereStaking delegation where V1 sees 0 (balanceOf missing)
//     * bounds its lock scan at MAX_LOCKS_SCANNED (F3)
//     * keeps V1's quota math / daily reset / per-op cap
//
//   AereTokenPaymasterV2
//     * withdrawToken(transfer) SUCCEEDS where V1's transferFrom(self) reverts
//
// Runs LOCAL (hardhat). Nothing here is deployed to mainnet.

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

const ONE = 10n ** 18n;
const DAY = 86400;
const ZERO32 = "0x" + "00".repeat(32);

function UO(o = {}) {
  return [
    o.sender ?? ethers.ZeroAddress,
    o.nonce ?? 0n,
    o.initCode ?? "0x",
    o.callData ?? "0x",
    o.accountGasLimits ?? ZERO32,
    o.preVerificationGas ?? 0n,
    o.gasFees ?? ZERO32,
    o.paymasterAndData ?? "0x",
    o.signature ?? "0x",
  ];
}

function tokenPmData(pm, token, maxTokenCost) {
  return ethers.concat([
    pm,
    ZERO32,
    token,
    ethers.zeroPadValue(ethers.toBeHex(maxTokenCost), 32),
  ]);
}

async function impersonate(addr) {
  await network.provider.send("hardhat_setBalance", [addr, "0x3635C9ADC5DEA00000"]); // 1000 ETH
  return ethers.getImpersonatedSigner(addr);
}

async function deployEP() {
  const EP = await (await ethers.getContractFactory("AereEntryPoint")).deploy();
  await EP.waitForDeployment();
  return EP;
}

// ═════════════════════════════════════════════════════════════════════════════
describe("AereStakeQuotaPaymasterV2 — fixes vs REAL AereStaking + AereLockedStaking", () => {
  let ep, epS, realStaking, realLocked, v1, v2, owner, user, val, del;

  beforeEach(async () => {
    [owner, user, val, del] = await ethers.getSigners();
    ep = await deployEP();
    epS = await impersonate(await ep.getAddress());
    realStaking = await (await ethers.getContractFactory("AereStaking")).deploy();
    await realStaking.waitForDeployment();
    realLocked = await (await ethers.getContractFactory("AereLockedStaking")).deploy();
    await realLocked.waitForDeployment();
    // V1 (buggy) and V2 (fixed) wired to the SAME real staking addresses.
    v1 = await (await ethers.getContractFactory("AereStakeQuotaPaymaster")).deploy(
      await ep.getAddress(), await realStaking.getAddress(), await realLocked.getAddress()
    );
    await v1.waitForDeployment();
    v2 = await (await ethers.getContractFactory("AereStakeQuotaPaymasterV2")).deploy(
      await ep.getAddress(), await realStaking.getAddress(), await realLocked.getAddress()
    );
    await v2.waitForDeployment();
  });

  it("REAL lock present: V1.stakedOf REVERTS, V2.stakedOf returns the real amount", async () => {
    await realLocked.connect(user).stake(0, { value: 1000n * ONE }); // 30-day, 1,000 AERE
    const lock = await realLocked.getLock(user.address, 0);
    expect(lock.amount).to.equal(1000n * ONE);

    // V1: mis-decodes the real struct → hard revert (the live bug).
    await expect(v1.stakedOf(user.address)).to.be.reverted;
    // V2: matched field order → reads the lock correctly.
    expect(await v2.stakedOf(user.address)).to.equal(1000n * ONE);
    expect(await v2.dailyQuotaOf(user.address)).to.equal(20n);
  });

  it("REAL lock present: V1 validate REVERTS (DoS), V2 validate SPONSORS", async () => {
    await realLocked.connect(user).stake(0, { value: 1000n * ONE });
    await expect(
      v1.connect(epS).validatePaymasterUserOp(UO({ sender: user.address }), ZERO32, 0)
    ).to.be.reverted;
    await expect(
      v2.connect(epS).validatePaymasterUserOp(UO({ sender: user.address }), ZERO32, 0)
    ).to.emit(v2, "Sponsored");
  });

  it("REAL delegation: V1 sees 0 (balanceOf missing), V2 counts the real delegation", async () => {
    await realStaking.connect(val).registerValidator(1000, { value: 500n * ONE }); // 500 AERE self-stake
    await realStaking.connect(del).delegate(val.address, { value: 250n * ONE });   // delegate 250 AERE

    // V1: balanceOf() reverts (caught → 0); delegator has no locks → 0. Delegation invisible.
    expect(await v1.stakedOf(del.address)).to.equal(0n);
    // V2: sums delegations across validatorList → sees the 250 AERE.
    expect(await v2.stakedOf(del.address)).to.equal(250n * ONE);
    // V2 also counts a validator's own self-stake.
    expect(await v2.stakedOf(val.address)).to.equal(500n * ONE);
  });

  it("REAL delegation + REAL lock combine in V2.stakedOf", async () => {
    await realStaking.connect(val).registerValidator(1000, { value: 500n * ONE });
    await realStaking.connect(del).delegate(val.address, { value: 250n * ONE });
    await realLocked.connect(del).stake(0, { value: 100n * ONE });
    expect(await v2.stakedOf(del.address)).to.equal(350n * ONE); // 250 delegated + 100 locked
  });

  it("V2 ignores a withdrawn / expired lock (same rule as V1 intended)", async () => {
    await realLocked.connect(user).stake(0, { value: 1000n * ONE });
    expect(await v2.stakedOf(user.address)).to.equal(1000n * ONE);
    // early-exit the lock → withdrawn=true → no longer counted
    await realLocked.connect(user).earlyExit(0);
    expect(await v2.stakedOf(user.address)).to.equal(0n);
  });

  it("F3: MAX_LOCKS_SCANNED bounds the lock loop (150 real-shaped locks → counts 100)", async () => {
    // MockLockedStaking has the REAL amount-first struct + a batch seed helper.
    const mock = await (await ethers.getContractFactory("MockLockedStaking")).deploy();
    await mock.waitForDeployment();
    const pm = await (await ethers.getContractFactory("AereStakeQuotaPaymasterV2")).deploy(
      await ep.getAddress(), await realStaking.getAddress(), await mock.getAddress()
    );
    await pm.waitForDeployment();
    const far = BigInt(Math.floor(Date.now() / 1000) + 3650 * DAY);
    // 150 active locks of 1 AERE each.
    await mock.seed(user.address, 60, ONE, far, { gasLimit: 15_000_000 });
    await mock.seed(user.address, 60, ONE, far, { gasLimit: 15_000_000 });
    await mock.seed(user.address, 30, ONE, far, { gasLimit: 15_000_000 });
    expect(await mock.lockCount(user.address)).to.equal(150n);
    const cap = await pm.MAX_LOCKS_SCANNED();
    expect(cap).to.equal(100n);
    // Only the first 100 are scanned → 100 AERE counted (not 150).
    expect(await pm.stakedOf(user.address)).to.equal(100n * ONE);
  });

  // ── quota logic preserved (driven via a REAL lock of known size) ──
  it("quota = staked*20/1000; full quota consumable; (quota+1)-th rejected; resets next day", async () => {
    await realLocked.connect(user).stake(0, { value: 100n * ONE }); // quota = 2/day
    expect(await v2.dailyQuotaOf(user.address)).to.equal(2n);
    await v2.connect(epS).validatePaymasterUserOp(UO({ sender: user.address }), ZERO32, 0);
    await v2.connect(epS).validatePaymasterUserOp(UO({ sender: user.address }), ZERO32, 0);
    await expect(
      v2.connect(epS).validatePaymasterUserOp(UO({ sender: user.address }), ZERO32, 0)
    ).to.be.revertedWith("StakePM: daily quota exhausted");
    await network.provider.send("evm_increaseTime", [DAY + 1]);
    await network.provider.send("evm_mine");
    await v2.connect(epS).validatePaymasterUserOp(UO({ sender: user.address }), ZERO32, 0); // ok again
  });

  it("zero-stake sender rejected; per-op gas cap enforced", async () => {
    await expect(
      v2.connect(epS).validatePaymasterUserOp(UO({ sender: user.address }), ZERO32, 0)
    ).to.be.revertedWith("StakePM: stake more AERE");

    await realLocked.connect(user).stake(0, { value: 1000n * ONE });
    const capWei = await v2.maxSponsoredGasWei();
    await expect(
      v2.connect(epS).validatePaymasterUserOp(UO({ sender: user.address }), ZERO32, capWei + 1n)
    ).to.be.revertedWith("StakePM: tx too expensive");
    await v2.connect(epS).validatePaymasterUserOp(UO({ sender: user.address }), ZERO32, capWei); // boundary ok
  });

  it("stakedOf tolerates a reverting locked-staking source (try/catch)", async () => {
    const mock = await (await ethers.getContractFactory("MockLockedStaking")).deploy();
    await mock.waitForDeployment();
    const pm = await (await ethers.getContractFactory("AereStakeQuotaPaymasterV2")).deploy(
      await ep.getAddress(), await realStaking.getAddress(), await mock.getAddress()
    );
    await pm.waitForDeployment();
    const far = BigInt(Math.floor(Date.now() / 1000) + 3650 * DAY);
    await mock.seed(user.address, 1, 1000n * ONE, far);
    expect(await pm.stakedOf(user.address)).to.equal(1000n * ONE);
    await mock.setBoom(true); // lockCount now reverts
    expect(await pm.stakedOf(user.address)).to.equal(0n); // degrades to 0, no revert
  });
});

// ═════════════════════════════════════════════════════════════════════════════
describe("AereTokenPaymasterV2 — withdrawToken(transfer) fix", () => {
  const AERE_USD = ethers.id("AERE/USD");
  const USDT_USD = ethers.id("USDT/USD");
  let ep, epS, oracle, token, owner, user;

  async function accumulate(pm) {
    // Drive one validate so the paymaster ends up holding `cost` tokens.
    await oracle.set(AERE_USD, 10_000_000); // $0.10
    await oracle.set(USDT_USD, 100_000_000); // $1.00
    const maxCost = ONE / 50n;
    const cost = await pm.quoteTokenAmount(await token.getAddress(), maxCost);
    await token.mint(user.address, cost * 10n);
    await token.connect(user).approve(await pm.getAddress(), ethers.MaxUint256);
    const pmAddr = await pm.getAddress();
    await pm.connect(epS).validatePaymasterUserOp(
      UO({ sender: user.address, paymasterAndData: tokenPmData(pmAddr, await token.getAddress(), cost) }),
      ZERO32, maxCost
    );
    return cost;
  }

  beforeEach(async () => {
    [owner, user] = await ethers.getSigners();
    ep = await deployEP();
    epS = await impersonate(await ep.getAddress());
    oracle = await (await ethers.getContractFactory("MockAereOracle")).deploy();
    await oracle.waitForDeployment();
    token = await (await ethers.getContractFactory("MockERC20Lending")).deploy("USDT.e", "USDT.e", 6);
    await token.waitForDeployment();
  });

  it("V1 CANNOT withdraw (transferFrom self reverts) — reproduces the finding", async () => {
    const pm = await (await ethers.getContractFactory("AereTokenPaymaster")).deploy(
      await ep.getAddress(), await oracle.getAddress()
    );
    await pm.waitForDeployment();
    await pm.setToken(await token.getAddress(), USDT_USD, 6, true);
    const cost = await accumulate(pm);
    expect(await token.balanceOf(await pm.getAddress())).to.equal(cost);
    await expect(pm.withdrawToken(await token.getAddress(), owner.address, cost))
      .to.be.revertedWith("ERC20: insufficient allowance");
  });

  it("V2 CAN withdraw accumulated tokens via transfer()", async () => {
    const pm = await (await ethers.getContractFactory("AereTokenPaymasterV2")).deploy(
      await ep.getAddress(), await oracle.getAddress()
    );
    await pm.waitForDeployment();
    await pm.setToken(await token.getAddress(), USDT_USD, 6, true);
    const cost = await accumulate(pm);
    const pmAddr = await pm.getAddress();
    expect(await token.balanceOf(pmAddr)).to.equal(cost);

    const before = await token.balanceOf(owner.address);
    await expect(pm.withdrawToken(await token.getAddress(), owner.address, cost))
      .to.emit(pm, "Withdrawn").withArgs(await token.getAddress(), owner.address, cost);
    expect(await token.balanceOf(owner.address)).to.equal(before + cost);
    expect(await token.balanceOf(pmAddr)).to.equal(0n);
  });

  it("V2 withdrawToken is onlyOwner", async () => {
    const pm = await (await ethers.getContractFactory("AereTokenPaymasterV2")).deploy(
      await ep.getAddress(), await oracle.getAddress()
    );
    await pm.waitForDeployment();
    await expect(pm.connect(user).withdrawToken(await token.getAddress(), user.address, 1))
      .to.be.revertedWith("Ownable: caller is not the owner");
  });

  it("V2 quote math unchanged from V1 (6-dec, $0.10 AERE, 5% markup → 2100)", async () => {
    const pm = await (await ethers.getContractFactory("AereTokenPaymasterV2")).deploy(
      await ep.getAddress(), await oracle.getAddress()
    );
    await pm.waitForDeployment();
    await pm.setToken(await token.getAddress(), USDT_USD, 6, true);
    await oracle.set(AERE_USD, 10_000_000);
    await oracle.set(USDT_USD, 100_000_000);
    const got = await pm.quoteTokenAmount(await token.getAddress(), ONE / 50n);
    expect(got).to.equal(2100n);
  });
});