pragma solidity 0.8.20;

import {Test} from "forge-std/Test.sol";

import "../../src/FatsoERC20.sol";
import "../../src/StinkyERC20.sol";
import "../../src/GhstERC20.sol";
import "../../src/GhostAuthority.sol";
import "../../src/StakingDistributor.sol";
import "../../src/Treasury.sol";
import "../../src/Staking.sol";
import "../../src/mocks/ERC20Mock.sol";

contract StakingTest is Test {
    address constant initializer       = 0x0000000000000000000000000000000000000001;
    address constant governor          = 0x0000000000000000000000000000000000000003;
    address constant guardian          = 0x0000000000000000000000000000000000000004;
    address constant policy            = 0x0000000000000000000000000000000000000005;
    address constant vault             = 0x0000000000000000000000000000000000000006;
    address constant alice             = 0x0000000000000000000000000000000000000007;
    address constant bob               = 0x0000000000000000000000000000000000000008;

    uint48 public constant EPOCH_LENGTH = 2200;
    uint48 public constant EPOCH_NUMBER = 1;
    uint48 public constant EPOCH_END_TIME = 1337;

    uint256 constant public INITIAL_INDEX = 10819917194513808e56;
    uint256 constant public TOTAL_INITIAL_SUPPLY = 5000000000000000;

    Fatso ftso;
    Stinky stnk;
    Ghost ghst;
    GhostStaking staking;
    GhostTreasury treasury;
    GhostAuthority authority;

    uint256 public constant amount = 69;

    event DistributorSet(address distributor);
    event WarmupSet(uint256 warmup);

    function setUp() public {
        vm.startPrank(initializer);
        authority = new GhostAuthority(
            governor,
            guardian,
            policy,
            vault
        );
        ftso = new Fatso(address(authority));
        stnk = new Stinky(INITIAL_INDEX);
        ghst = new Ghost(address(stnk));
        staking = new GhostStaking(
            address(ftso),
            address(stnk),
            address(ghst),
            EPOCH_LENGTH,
            EPOCH_NUMBER,
            EPOCH_END_TIME,
            address(authority)
        );
        treasury = new GhostTreasury(address(ftso), 69, address(authority));
        stnk.initialize(address(staking), address(treasury), address(ghst));
        ghst.initialize(address(staking));
        vm.stopPrank();
    }

    function test_correctAfterConstruction() public view {
        assertEq(staking.stnk(), address(stnk));
        assertEq(staking.ftso(), address(ftso));
        assertEq(staking.ghst(), address(ghst));
        (uint48 length, uint48 number, uint48 end, uint256 distribute) = staking.epoch();
        assertEq(length, EPOCH_LENGTH);
        assertEq(number, EPOCH_NUMBER);
        assertEq(end, EPOCH_END_TIME);
        assertEq(distribute, 0);
    }

    function test_governorCouldSetDistributor(address maybeDistributor) public {
        vm.assume(maybeDistributor != address(0));
        assertEq(staking.distributor(), address(0));
        vm.prank(governor);
        staking.setDistributor(maybeDistributor);
        assertEq(staking.distributor(), maybeDistributor);
    }

    function test_emitsDistributorSetEvent(address maybeDistributor) public {
        vm.assume(maybeDistributor != address(0));
        vm.expectEmit(true, true, true, false, address(staking));
        emit DistributorSet(maybeDistributor);
        vm.prank(governor);
        staking.setDistributor(maybeDistributor);
    }

    function test_arbitraryAddressCouldNotAddDistributor(address someone, address maybeDistributor) public {
        vm.assume(maybeDistributor != address(0) && someone != governor);
        vm.expectRevert();
        vm.prank(someone);
        staking.setDistributor(maybeDistributor);
    }

    function test_governorCouldSetWarmupPeriod(uint48 maybeWarmupPeriod) public {
        vm.assume(maybeWarmupPeriod > 0);
        assertEq(staking.warmupPeriod(), 0);
        vm.prank(governor);
        staking.setWarmupPeriod(maybeWarmupPeriod);
        assertEq(staking.warmupPeriod(), maybeWarmupPeriod);
    }

    function test_emitsWarmupEvent(uint48 maybeWarmupPeriod) public {
        vm.expectEmit(true, true, true, false, address(staking));
        emit WarmupSet(maybeWarmupPeriod);
        vm.prank(governor);
        staking.setWarmupPeriod(maybeWarmupPeriod);
    }

    function test_arbitraryAddressCouldNotSetWarmupPeriod(address someone, uint48 maybeWarmupPeriod) public {
        vm.assume(maybeWarmupPeriod > 0 && someone != governor);
        vm.expectRevert();
        vm.prank(someone);
        staking.setWarmupPeriod(maybeWarmupPeriod);
    }

    function test_stake_addAmountToTheWarmupWhenClaimIsFalse() public {
        _mintAndApprove(alice, amount);

        vm.prank(alice);
        uint256 rebased = staking.stake(amount, alice, false, false);

        assertEq(ftso.balanceOf(alice), 0);
        assertEq(staking.sharesInWarmup(), stnk.sharesForBalance(rebased));

        (uint256 deposit, uint256 shares, uint48 expiry, bool lock) = staking.warmupInfo(alice);
        assertEq(deposit, amount);
        assertEq(shares, stnk.sharesForBalance(amount));
        assertEq(expiry, 1);
        assertEq(lock, false);
    }

    function test_stake_exchangesFatsoToStinkyWhenClaimIsTrueAndRebasingIsTrue() public {
        _mintAndApprove(alice, amount);

        vm.prank(alice);
        uint256 rebasedAmount = staking.stake(amount, alice, true, true);

        assertEq(ftso.balanceOf(alice), 0);
        assertEq(stnk.balanceOf(alice), rebasedAmount);
    }

    function test_stake_exchangesFatsoForNewlyMintedGhostWhenClaimIsTrueAndRebasingIsFalse() public {
        _mintAndApprove(alice, amount);

        vm.prank(alice);
        uint256 rebasedAmount = staking.stake(amount, alice, false, true);

        assertEq(ftso.balanceOf(alice), 0);
        assertEq(ghst.balanceOf(alice), rebasedAmount);
    }

    function test_stake_addAmountToWarmupWhenClaimIsTrueAndWarmupGtZero() public {
        _mintAndApprove(alice, amount);
        vm.prank(governor);
        staking.setWarmupPeriod(1);

        vm.prank(alice);
        uint256 rebased = staking.stake(amount, alice, false, true);

        assertEq(ftso.balanceOf(alice), 0);
        assertEq(staking.sharesInWarmup(), stnk.sharesForBalance(rebased));

        (uint256 deposit, uint256 shares, uint48 expiry, bool lock) = staking.warmupInfo(alice);
        assertEq(deposit, amount);
        assertEq(shares, stnk.sharesForBalance(amount));
        assertEq(expiry, 2);
        assertEq(lock, false);
    }

    function test_stake_allowsSelfDepositWhenNotLocked() public {
        _mintAndApprove(alice, amount);

        vm.prank(alice);
        staking.stake(amount, bob, false, false);

        assertEq(ftso.balanceOf(alice), 0);
    }

    function test_stake_disablesExternalDepositsWhenLocked() public {
        _mintAndApprove(alice, amount);
        vm.prank(bob);
        staking.toggleLock();

        vm.expectRevert();
        vm.prank(alice);
        staking.stake(amount, bob, false, false);

        assertEq(ftso.balanceOf(alice), amount);
    }

    function test_claim_transferStinkyWhenRebasingIsTrue() public {
        _prepareAndRoll(alice, amount, false, false);

        assertEq(stnk.balanceOf(alice), 0);
        vm.prank(alice);
        uint256 rebased = staking.claim(alice, true);
        assertEq(stnk.balanceOf(alice), rebased);
        assertEq(rebased > 0, true);
    }

    function test_claim_mintsGhostWhenRebasingIsFalse() public {
        _prepareAndRoll(alice, amount, false, false);

        assertEq(ghst.balanceOf(alice), 0);
        vm.prank(alice);
        uint256 rebased = staking.claim(alice, false);
        assertEq(ghst.balanceOf(alice), rebased);
        assertEq(rebased > 0, true);
    }

    function test_claim_preventsExternalClaimsWhenLocked() public {
        _prepareAndRoll(alice, amount, false, false);

        assertEq(ghst.balanceOf(alice), 0);
        vm.prank(alice);
        staking.toggleLock();

        vm.expectRevert();
        vm.prank(bob);
        staking.claim(alice, false);
        assertEq(ghst.balanceOf(alice), 0);
    }

    function test_claim_allowExternalClaimsWhenNotLocked() public {
        _prepareAndRoll(alice, amount, false, false);

        assertEq(ghst.balanceOf(alice), 0);
        vm.prank(bob);
        uint256 rebased = staking.claim(alice, false);
        assertEq(ghst.balanceOf(alice), rebased);
        assertEq(rebased > 0, true);
    }

    function test_claim_doesNothingWhenThereIsNothingToClaim() public {
        assertEq(ghst.balanceOf(alice), 0);
        vm.prank(alice);
        uint256 rebased = staking.claim(alice, false);
        assertEq(ghst.balanceOf(alice), 0);
        assertEq(rebased, 0);
    }

    function test_claim_doesNothingWhenWarmupIsntOver() public {
        _mintAndApprove(alice, amount);
        vm.prank(governor);
        staking.setWarmupPeriod(1337);
        vm.prank(alice);
        staking.stake(amount, alice, false, false);

        assertEq(stnk.balanceOf(alice), 0);
        vm.prank(alice);
        uint256 rebased = staking.claim(alice, true);
        assertEq(stnk.balanceOf(alice), 0);
        assertEq(rebased, 0);
    }

    function test_forefeit_removesStakeFromWarmupAndReturnsFatso() public {
        _prepareAndRoll(alice, amount, false, false);

        assertEq(ftso.balanceOf(alice), 0);
        assertEq(staking.sharesInWarmup(), stnk.sharesForBalance(amount));

        vm.prank(alice);
        uint256 deposit = staking.forfeit();

        assertEq(ftso.balanceOf(alice), deposit);
        assertEq(staking.sharesInWarmup(), 0);

        (uint256 depositInWarmup, uint256 shares, uint48 expiry,) = staking.warmupInfo(alice);
        assertEq(depositInWarmup, 0);
        assertEq(shares, 0);
        assertEq(expiry, 0);
    }

    function test_forefeit_transfersZeroIfThereIsNoBalanceInWarmup() public {
        vm.prank(alice);
        uint256 deposit = staking.forfeit();
        assertEq(deposit, 0);
        assertEq(ftso.balanceOf(alice), 0);
        assertEq(staking.sharesInWarmup(), 0);
    }

    function test_unstake_canRedeemStinkyToFatso() public {
        _prepareAndRoll(alice, amount, true, true);
        uint256 aliceBalance = stnk.balanceOf(alice);
        vm.startPrank(alice);
        stnk.approve(address(staking), aliceBalance);
        staking.unstake(aliceBalance, alice, false, true);
        vm.stopPrank();
        assertEq(stnk.balanceOf(alice), 0);
        assertEq(ftso.balanceOf(alice), amount);
    }

    function test_unstake_canRedeemGhostToFatso() public {
        _prepareAndRoll(alice, amount, false, true);
        uint256 aliceBalance = ghst.balanceOf(alice);
        vm.startPrank(alice);
        ghst.approve(address(staking), aliceBalance);
        vm.roll(block.number + 69);
        staking.unstake(aliceBalance, alice, false, false);
        vm.stopPrank();
        assertEq(ghst.balanceOf(alice), 0);
        assertEq(ftso.balanceOf(alice), ghst.balanceFrom(aliceBalance));
    }

    function test_wrap_convertsStinkyIntoGhost() public {
        _prepareAndRoll(alice, amount, true, true);
        uint256 aliceBalance = stnk.balanceOf(alice);
        assertEq(ghst.balanceOf(alice), 0);

        vm.startPrank(alice);
        stnk.approve(address(staking), aliceBalance);
        uint256 newBalance = staking.wrap(alice, aliceBalance);
        vm.stopPrank();

        assertEq(stnk.balanceOf(alice), 0);
        assertEq(ghst.balanceOf(alice), newBalance);
    }

    function test_wrap_convertsGhostToStinky() public {
        _prepareAndRoll(alice, amount, false, true);
        uint256 aliceBalance = ghst.balanceOf(alice);
        assertEq(stnk.balanceOf(alice), 0);

        vm.roll(block.number + 69);
        vm.startPrank(alice);
        ghst.approve(address(staking), aliceBalance);
        uint256 newBalance = staking.unwrap(alice, aliceBalance);
        vm.stopPrank();

        assertEq(ghst.balanceOf(alice), 0);
        assertEq(stnk.balanceOf(alice), newBalance);
    }

    function test_rebase_doesNothingIfTheBlockIsBeforeTheEpochEndBlock() public {
        (, uint48 numberBefore, uint48 endBefore, uint256 distributeBefore) = staking.epoch();
        assertEq(endBefore > block.timestamp, true);
        staking.rebase();
        (, uint48 numberAfter, uint48 endAfter, uint256 distributeAfter) = staking.epoch();
        assertEq(distributeAfter, distributeBefore);
        assertEq(endAfter, endBefore);
        assertEq(numberAfter, numberBefore);
    }

    function test_rebase_incrementsEpochNumberAndCallRebase() public {
        (uint48 length, uint48 numberBefore, uint48 endBefore,) = staking.epoch();
        assertEq(endBefore > block.timestamp, true);
        vm.warp(endBefore);
        staking.rebase();
        (, uint48 numberAfter, uint48 endAfter,) = staking.epoch();

        assertEq(endAfter, endBefore + length);
        assertEq(numberAfter, numberBefore + 1);
    }

    function test_rebase_whenTheFatsoBalanceEqualsToStinkySupplyZeroDistribution() public {
        (uint48 length, uint48 numberBefore, uint48 endBefore, uint256 distributeBefore) = staking.epoch();
        assertEq(endBefore > block.timestamp, true);

        vm.prank(address(staking));
        stnk.transfer(alice, amount);
        vm.prank(vault);
        ftso.mint(address(staking), amount);

        assertEq(stnk.balanceOf(alice), ftso.balanceOf(address(staking)));

        vm.warp(endBefore);
        staking.rebase();

        (, uint48 numberAfter, uint48 endAfter, uint256 distributeAfter) = staking.epoch();
        assertEq(endAfter, endBefore + length);
        assertEq(numberAfter, numberBefore + 1);
        assertEq(distributeAfter, distributeBefore);
    }

    function test_rebase_distributeDifferenceBetweenStakedAndTotalSupply() public {
        (uint48 length, uint48 numberBefore, uint48 endBefore,) = staking.epoch();
        assertEq(endBefore > block.timestamp, true);

        vm.prank(address(staking));
        stnk.transfer(alice, amount);
        vm.prank(vault);
        ftso.mint(address(staking), amount * 2);

        vm.warp(endBefore);
        staking.rebase();

        (, uint48 numberAfter, uint48 endAfter, uint256 distributeAfter) = staking.epoch();
        assertEq(endAfter, endBefore + length);
        assertEq(numberAfter, numberBefore + 1);
        assertEq(distributeAfter, amount);
    }

    function test_rebase_callDistributorIfSet() public {
        (uint48 length, uint48 numberBefore, uint48 endBefore,) = staking.epoch();
        assertEq(endBefore > block.timestamp, true);
        uint256 extendedAmount = amount * 1e18;

        GhostDistributor distributor = new GhostDistributor(
            address(treasury),
            address(ftso),
            address(staking),
            address(authority),
            1e4 // 1%
        );
        ERC20Mock reserve = new ERC20Mock("Reserve Token", "RET");

        vm.startPrank(governor);
        authority.pushVault(address(treasury));
        staking.setDistributor(address(distributor));
        treasury.enable(ITreasury.STATUS.REWARDMANAGER, address(distributor), address(0));
        treasury.enable(ITreasury.STATUS.RESERVEDEPOSITOR, alice, address(0));
        treasury.enable(ITreasury.STATUS.RESERVEDEPOSITOR, bob, address(0));
        treasury.enable(ITreasury.STATUS.RESERVETOKEN, address(reserve), address(0));
        vm.stopPrank();

        vm.startPrank(bob);
        reserve.mint(bob, extendedAmount);
        reserve.approve(address(treasury), extendedAmount);
        uint256 zeroProfit = treasury.tokenValue(address(reserve), extendedAmount);
        treasury.deposit(address(reserve), extendedAmount, zeroProfit);
        vm.stopPrank();

        vm.startPrank(alice);
        assertEq(ftso.balanceOf(address(staking)), 0);
        reserve.mint(alice, extendedAmount);
        reserve.approve(address(treasury), extendedAmount);
        uint256 send = treasury.deposit(address(reserve), extendedAmount, 0);
        assertEq(ftso.balanceOf(alice), send);
        assertEq(reserve.balanceOf(address(treasury)), extendedAmount * 2);
        ftso.approve(address(staking), send);
        staking.stake(send, alice, true, true);
        assertEq(ftso.balanceOf(address(staking)), send);
        assertEq(stnk.balanceOf(alice), send);
        vm.stopPrank();

        assertEq(treasury.excessReserves(), send);
        assertEq(treasury.totalReserves(), send * 2);

        skip(endBefore);
        uint256 extraToStaking = send * 1e4 / 1e6;
        vm.prank(address(staking));
        staking.rebase();
        assertEq(ftso.balanceOf(address(staking)), send + extraToStaking);

        (, uint48 numberAfter, uint48 endAfter, uint256 distributeAfter) = staking.epoch();
        assertEq(endAfter, endBefore + length);
        assertEq(numberAfter, numberBefore + 1);
        assertEq(distributeAfter, extraToStaking);
    }

    function test_correctStakeAmountDuringRebase() public {
        uint256 extendedAmount = amount * 1e18;
        uint256 bounty = 1;
        GhostDistributor distributor = new GhostDistributor(
            address(treasury),
            address(ftso),
            address(staking),
            address(authority),
            1e6
        );
        ERC20Mock reserve = new ERC20Mock("Reserve Token", "RET");

        vm.startPrank(governor);
        authority.pushVault(address(treasury));
        staking.setDistributor(address(distributor));
        distributor.setBounty(bounty);
        treasury.enable(ITreasury.STATUS.REWARDMANAGER, address(distributor), address(0));
        treasury.enable(ITreasury.STATUS.RESERVEDEPOSITOR, alice, address(0));
        treasury.enable(ITreasury.STATUS.RESERVETOKEN, address(reserve), address(0));
        vm.stopPrank();

        vm.startPrank(alice);
        reserve.mint(alice, extendedAmount);
        reserve.approve(address(treasury), type(uint256).max);

        uint256 profit = treasury.tokenValue(address(reserve), extendedAmount) - amount;
        treasury.deposit(address(reserve), extendedAmount, profit);
        vm.stopPrank();

        (,, uint48 end,) = staking.epoch();
        skip(end);

        vm.startPrank(alice);
        ftso.approve(address(staking), type(uint256).max);
        staking.stake(amount, alice, true, true);
        vm.stopPrank();

        uint256 postBounty = amount + bounty;

        assertEq(stnk.balanceOf(alice), postBounty);
        assertEq(ftso.balanceOf(address(staking)), postBounty);

        skip(end + EPOCH_LENGTH);
        staking.rebase();

        (,,, uint256 distribute) = staking.epoch();
        assertEq(distribute, postBounty);
        assertEq(stnk.balanceOf(alice), postBounty);
        assertEq(ftso.balanceOf(address(staking)), 2 * postBounty + bounty);
    }

    function _mintAndApprove(address who, uint256 value) internal {
        vm.prank(vault);
        ftso.mint(who, value);
        vm.prank(who);
        ftso.approve(address(staking), value);
    }

    function _prepareAndRoll(address who, uint256 value, bool rebase, bool claim) internal returns (uint256) {
        _mintAndApprove(who, value);

        (,, uint48 end,) = staking.epoch();
        skip(end);

        vm.prank(who);
        uint256 rebased = staking.stake(value, who, rebase, claim);

        (,, uint48 expiry,) = staking.warmupInfo(alice);
        vm.roll(expiry);
        return rebased;
    }
}