import { expect } from "chai";
import { ethers } from "hardhat";
import { time } from "@nomicfoundation/hardhat-network-helpers";
import { KeyVault, MockIdentityRegistry, ReentrancyAttacker } from "../typechain-types";
import { SignerWithAddress } from "@nomicfoundation/hardhat-ethers/signers";

describe("KeyVault", function () {
  let identity: MockIdentityRegistry;
  let vault: KeyVault;
  let owner: SignerWithAddress;
  let attacker: SignerWithAddress;
  let recipient: SignerWithAddress;
  let random: SignerWithAddress;

  const MAX_PER_TX = ethers.parseEther("1");      // 1 ETH per tx
  const MAX_HOURLY = ethers.parseEther("5");       // 5 ETH per hour
  const MAX_DAILY = ethers.parseEther("10");       // 10 ETH per day
  const EMERGENCY_THRESHOLD = ethers.parseEther("20"); // freeze at 20 ETH total

  beforeEach(async function () {
    [owner, attacker, recipient, random] = await ethers.getSigners();

    // Deploy mock identity registry and register the owner
    const IdentityFactory = await ethers.getContractFactory("MockIdentityRegistry");
    identity = (await IdentityFactory.deploy()) as MockIdentityRegistry;
    await identity.register(owner.address);

    // Deploy KeyVault
    const VaultFactory = await ethers.getContractFactory("KeyVault");
    vault = (await VaultFactory.deploy(
      owner.address,
      await identity.getAddress(),
      MAX_PER_TX,
      MAX_HOURLY,
      MAX_DAILY,
      EMERGENCY_THRESHOLD
    )) as KeyVault;

    // Fund the vault with 50 ETH
    await owner.sendTransaction({
      to: await vault.getAddress(),
      value: ethers.parseEther("50"),
    });
  });

  // ═══════════════════════════════════════════════════
  // DEPLOYMENT
  // ═══════════════════════════════════════════════════

  describe("Deployment", function () {
    it("Should set the correct owner", async function () {
      expect(await vault.owner()).to.equal(owner.address);
    });

    it("Should set correct policy values", async function () {
      const policy = await vault.policy();
      expect(policy.maxPerTx).to.equal(MAX_PER_TX);
      expect(policy.maxHourly).to.equal(MAX_HOURLY);
      expect(policy.maxDaily).to.equal(MAX_DAILY);
      expect(policy.emergencyThreshold).to.equal(EMERGENCY_THRESHOLD);
      expect(policy.frozen).to.equal(false);
    });

    it("Should reject unregistered owner", async function () {
      const VaultFactory = await ethers.getContractFactory("KeyVault");
      await expect(
        VaultFactory.deploy(
          attacker.address, // not registered
          await identity.getAddress(),
          MAX_PER_TX,
          MAX_HOURLY,
          MAX_DAILY,
          EMERGENCY_THRESHOLD
        )
      ).to.be.revertedWithCustomError(vault, "NotRegistered");
    });

    it("Should reject zero address for owner", async function () {
      const VaultFactory = await ethers.getContractFactory("KeyVault");
      await expect(
        VaultFactory.deploy(
          ethers.ZeroAddress,
          await identity.getAddress(),
          MAX_PER_TX,
          MAX_HOURLY,
          MAX_DAILY,
          EMERGENCY_THRESHOLD
        )
      ).to.be.revertedWithCustomError(vault, "ZeroAddress");
    });

    it("Should reject zero address for identity registry", async function () {
      const VaultFactory = await ethers.getContractFactory("KeyVault");
      await expect(
        VaultFactory.deploy(
          owner.address,
          ethers.ZeroAddress,
          MAX_PER_TX,
          MAX_HOURLY,
          MAX_DAILY,
          EMERGENCY_THRESHOLD
        )
      ).to.be.revertedWithCustomError(vault, "ZeroAddress");
    });

    it("Should accept deposits via receive()", async function () {
      expect(await vault.balance()).to.equal(ethers.parseEther("50"));
    });

    it("Should accept deposits via deposit()", async function () {
      await vault.deposit({ value: ethers.parseEther("5") });
      expect(await vault.balance()).to.equal(ethers.parseEther("55"));
    });
  });

  // ═══════════════════════════════════════════════════
  // SPENDING LIMITS
  // ═══════════════════════════════════════════════════

  describe("Spending Limits", function () {
    it("Should allow transaction within per-tx limit", async function () {
      await expect(vault.execute(recipient.address, ethers.parseEther("0.5"), "0x"))
        .to.emit(vault, "TransactionExecuted");
    });

    it("Should reject transaction exceeding per-tx limit", async function () {
      await expect(
        vault.execute(recipient.address, ethers.parseEther("1.5"), "0x")
      ).to.be.revertedWithCustomError(vault, "ExceedsPerTxLimit");
    });

    it("Should enforce hourly spending limit across multiple txs", async function () {
      // Spend 1 ETH x 5 = 5 ETH (hits hourly max)
      for (let i = 0; i < 5; i++) {
        await vault.execute(recipient.address, MAX_PER_TX, "0x");
      }
      // 6th should fail — hourly limit reached
      await expect(
        vault.execute(recipient.address, ethers.parseEther("0.01"), "0x")
      ).to.be.revertedWithCustomError(vault, "ExceedsHourlyLimit");
    });

    it("Should reset hourly limit after 1 hour", async function () {
      // Spend to hourly max
      for (let i = 0; i < 5; i++) {
        await vault.execute(recipient.address, MAX_PER_TX, "0x");
      }
      // Advance time 1 hour + 1 second
      await time.increase(3601);

      // Should work again
      await expect(vault.execute(recipient.address, MAX_PER_TX, "0x"))
        .to.emit(vault, "TransactionExecuted");
    });

    it("Should enforce daily spending limit", async function () {
      // Spend 1 ETH x 5, wait an hour, repeat = 10 ETH (daily max)
      for (let i = 0; i < 5; i++) {
        await vault.execute(recipient.address, MAX_PER_TX, "0x");
      }
      await time.increase(3601);
      for (let i = 0; i < 5; i++) {
        await vault.execute(recipient.address, MAX_PER_TX, "0x");
      }

      // Advance another hour, but still same day
      await time.increase(3601);

      // Should fail — daily limit reached
      await expect(
        vault.execute(recipient.address, ethers.parseEther("0.01"), "0x")
      ).to.be.revertedWithCustomError(vault, "ExceedsDailyLimit");
    });

    it("Should reset daily limit after 24 hours", async function () {
      // Spend to daily max (10 ETH) across multiple hours
      for (let i = 0; i < 5; i++) {
        await vault.execute(recipient.address, MAX_PER_TX, "0x");
      }
      await time.increase(3601);
      for (let i = 0; i < 5; i++) {
        await vault.execute(recipient.address, MAX_PER_TX, "0x");
      }

      // Advance 24 hours
      await time.increase(86401);

      // Should work again
      await expect(vault.execute(recipient.address, MAX_PER_TX, "0x"))
        .to.emit(vault, "TransactionExecuted");
    });

    it("Should track remaining hourly allowance correctly", async function () {
      expect(await vault.remainingHourly()).to.equal(MAX_HOURLY);
      await vault.execute(recipient.address, ethers.parseEther("1"), "0x");
      expect(await vault.remainingHourly()).to.equal(ethers.parseEther("4"));
    });

    it("Should track remaining daily allowance correctly", async function () {
      expect(await vault.remainingDaily()).to.equal(MAX_DAILY);
      await vault.execute(recipient.address, ethers.parseEther("1"), "0x");
      expect(await vault.remainingDaily()).to.equal(ethers.parseEther("9"));
    });

    it("Should only allow owner to execute", async function () {
      await expect(
        vault.connect(attacker).execute(recipient.address, ethers.parseEther("0.1"), "0x")
      ).to.be.revertedWithCustomError(vault, "NotOwner");
    });

    it("Should reject execution to zero address", async function () {
      await expect(
        vault.execute(ethers.ZeroAddress, ethers.parseEther("0.1"), "0x")
      ).to.be.revertedWithCustomError(vault, "ZeroAddress");
    });
  });

  // ═══════════════════════════════════════════════════
  // CONTRACT ALLOWLIST
  // ═══════════════════════════════════════════════════

  describe("Contract Allowlist", function () {
    it("Should allow any contract when allowlist disabled", async function () {
      expect(await vault.allowlistEnabled()).to.equal(false);
      await expect(vault.execute(recipient.address, ethers.parseEther("0.1"), "0x"))
        .to.emit(vault, "TransactionExecuted");
    });

    it("Should restrict to allowlisted contracts when enabled", async function () {
      await vault.allowContract(recipient.address);
      expect(await vault.allowlistEnabled()).to.equal(true);

      // Allowed target should work
      await expect(vault.execute(recipient.address, ethers.parseEther("0.1"), "0x"))
        .to.emit(vault, "TransactionExecuted");

      // Non-allowed target should fail
      await expect(
        vault.execute(random.address, ethers.parseEther("0.1"), "0x")
      ).to.be.revertedWithCustomError(vault, "ContractNotAllowed");
    });

    it("Should auto-disable allowlist when last contract removed", async function () {
      await vault.allowContract(recipient.address);
      expect(await vault.allowlistEnabled()).to.equal(true);

      await vault.removeContract(recipient.address);
      expect(await vault.allowlistEnabled()).to.equal(false);

      // Any target should work again
      await expect(vault.execute(random.address, ethers.parseEther("0.1"), "0x"))
        .to.emit(vault, "TransactionExecuted");
    });

    it("Should reject allowContract from non-owner", async function () {
      await expect(
        vault.connect(attacker).allowContract(recipient.address)
      ).to.be.revertedWithCustomError(vault, "NotOwner");
    });

    it("Should reject zero address in allowContract", async function () {
      await expect(
        vault.allowContract(ethers.ZeroAddress)
      ).to.be.revertedWithCustomError(vault, "ZeroAddress");
    });
  });

  // ═══════════════════════════════════════════════════
  // EMERGENCY CONTROLS
  // ═══════════════════════════════════════════════════

  describe("Emergency Controls", function () {
    it("Should freeze vault when threshold exceeded", async function () {
      // Set a lower threshold so we can test without fighting daily limits
      await vault.updatePolicy(
        ethers.parseEther("1"),  // maxPerTx
        ethers.parseEther("5"),  // maxHourly
        ethers.parseEther("10"), // maxDaily
        ethers.parseEther("3")   // emergencyThreshold: freeze after 3 ETH total
      );

      // Spend 1 ETH x 3 = 3 ETH. On the 3rd, totalSpent (3) > threshold (3)? No.
      // We need totalSpent > threshold, so spend 4 ETH total
      await vault.execute(recipient.address, ethers.parseEther("1"), "0x");
      await vault.execute(recipient.address, ethers.parseEther("1"), "0x");
      await vault.execute(recipient.address, ethers.parseEther("1"), "0x");
      
      // totalSpent is now 3. Not > 3 yet.
      expect(await vault.isFrozen()).to.equal(false);

      // 4th ETH pushes totalSpent to 4 > 3 threshold
      await vault.execute(recipient.address, ethers.parseEther("1"), "0x");
      expect(await vault.isFrozen()).to.equal(true);

      // Further transactions should fail
      await expect(
        vault.execute(recipient.address, ethers.parseEther("0.01"), "0x")
      ).to.be.revertedWithCustomError(vault, "VaultFrozen");
    });

    it("Should allow manual freeze", async function () {
      await vault.freeze();
      expect(await vault.isFrozen()).to.equal(true);
      await expect(
        vault.execute(recipient.address, ethers.parseEther("0.01"), "0x")
      ).to.be.revertedWithCustomError(vault, "VaultFrozen");
    });

    it("Should allow unfreeze and reset total", async function () {
      await vault.freeze();
      await vault.unfreeze();
      expect(await vault.isFrozen()).to.equal(false);

      // Should work again
      await expect(vault.execute(recipient.address, ethers.parseEther("0.1"), "0x"))
        .to.emit(vault, "TransactionExecuted");
    });

    it("Should allow emergency withdraw (only owner)", async function () {
      const balanceBefore = await ethers.provider.getBalance(owner.address);
      await vault.emergencyWithdraw();
      const balanceAfter = await ethers.provider.getBalance(owner.address);

      // Owner should have gained ~50 ETH (minus gas)
      expect(balanceAfter - balanceBefore).to.be.greaterThan(ethers.parseEther("49"));
      expect(await vault.balance()).to.equal(0);
    });

    it("Should reject emergency withdraw from non-owner", async function () {
      await expect(
        vault.connect(attacker).emergencyWithdraw()
      ).to.be.revertedWithCustomError(vault, "NotOwner");
    });

    it("Should reject freeze from non-owner", async function () {
      await expect(
        vault.connect(attacker).freeze()
      ).to.be.revertedWithCustomError(vault, "NotOwner");
    });

    it("Should reject unfreeze from non-owner", async function () {
      await vault.freeze();
      await expect(
        vault.connect(attacker).unfreeze()
      ).to.be.revertedWithCustomError(vault, "NotOwner");
    });
  });

  // ═══════════════════════════════════════════════════
  // POLICY UPDATES
  // ═══════════════════════════════════════════════════

  describe("Policy Updates", function () {
    it("Should allow owner to update policy", async function () {
      const newMaxPerTx = ethers.parseEther("2");
      const newMaxHourly = ethers.parseEther("10");
      const newMaxDaily = ethers.parseEther("20");
      const newThreshold = ethers.parseEther("50");

      await expect(vault.updatePolicy(newMaxPerTx, newMaxHourly, newMaxDaily, newThreshold))
        .to.emit(vault, "PolicyUpdated");

      const policy = await vault.policy();
      expect(policy.maxPerTx).to.equal(newMaxPerTx);
      expect(policy.maxHourly).to.equal(newMaxHourly);
      expect(policy.maxDaily).to.equal(newMaxDaily);
      expect(policy.emergencyThreshold).to.equal(newThreshold);
    });

    it("Should reject policy update from non-owner", async function () {
      await expect(
        vault.connect(attacker).updatePolicy(0, 0, 0, 0)
      ).to.be.revertedWithCustomError(vault, "NotOwner");
    });
  });

  // ═══════════════════════════════════════════════════
  // REENTRANCY ATTACK
  // ═══════════════════════════════════════════════════

  describe("Reentrancy Protection", function () {
    it("Should block reentrancy attack on execute()", async function () {
      // Deploy attacker contract
      const AttackerFactory = await ethers.getContractFactory("ReentrancyAttacker");
      const attackerContract = (await AttackerFactory.deploy(
        await vault.getAddress()
      )) as ReentrancyAttacker;

      // Start the attack mode
      await attackerContract.startAttack();

      // Execute a legitimate send to the attacker contract
      // The attacker's receive() will try to re-enter execute()
      await vault.execute(
        await attackerContract.getAddress(),
        ethers.parseEther("0.1"),
        "0x"
      );

      // The attacker should have failed to re-enter (count stays 0 or reverted)
      // The vault should still have nearly all its funds
      expect(await vault.balance()).to.be.greaterThan(ethers.parseEther("49"));
    });
  });
});
