// SPDX-License-Identifier: GPL-3.0-or-later
pragma solidity ^0.8.19;

import "forge-std/Test.sol";
import "forge-std/console.sol";
import {TwabController} from "pt-v5-twab-controller/TwabController.sol";
import "../src/Harberg.sol";
import {TooMuchSnatch, Stake} from "../src/Stake.sol";

contract StakeTest is Test {
    TwabController tc;
    Harberg harberg;
    Stake stakingPool;
    address liquidityPool;
    address liquidityManager;
    address taxPool;

    event PositionCreated(uint256 indexed positionId, address indexed owner, uint256 harbergDeposit, uint256 share, uint32 taxRate);
    event PositionRemoved(uint256 indexed positionId, address indexed owner, uint256 harbergPayout);


    function setUp() public {
        tc = new TwabController(60 * 60, uint32(block.timestamp));
        harberg = new Harberg("HARB", "HARB", tc);
        taxPool = harberg.TAX_POOL();
        stakingPool = new Stake(address(harberg));
        harberg.setStakingPool(address(stakingPool));
        liquidityPool = makeAddr("liquidityPool");
        harberg.setLiquidityPool(liquidityPool);
        liquidityManager = makeAddr("liquidityManager");
        harberg.setLiquidityManager(liquidityManager);
    }


    function testBasicStaking() public {
        // Setup
        uint256 stakeAmount = 1 ether;
        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        harberg.mint(stakeAmount * 5);
        harberg.transfer(staker, stakeAmount);
        vm.stopPrank();

        vm.startPrank(staker);

        // Approve and stake
        harberg.approve(address(stakingPool), stakeAmount);
        uint256[] memory empty;
        uint256 sharesExpected = stakingPool.assetsToShares(stakeAmount);
        vm.expectEmit(address(stakingPool));
        emit PositionCreated(654321, staker, stakeAmount, sharesExpected, 1);
        uint256 positionId = stakingPool.snatch(stakeAmount, staker, 1, empty);

        // Check results
        assertEq(stakingPool.outstandingStake(), stakingPool.assetsToShares(stakeAmount), "Outstanding stake did not update correctly");
        (uint256 share, address owner, uint32 creationTime, , uint32 taxRate) = stakingPool.positions(positionId);
        assertEq(stakingPool.sharesToAssets(share), stakeAmount, "Stake amount in position is incorrect");
        assertEq(owner, staker, "Stake owner is incorrect");
        assertEq(creationTime, uint32(block.timestamp), "Creation time is incorrect");
        assertEq(taxRate, 1, "Tax rate should be initialized to 1");

        vm.stopPrank();
    }

    function testUnstaking() public {
        // Setup: Create a staking position first
        uint256 stakeAmount = 1 ether;
        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        harberg.mint(stakeAmount * 5);  // Ensuring the staker has enough balance
        harberg.transfer(staker, stakeAmount);
        vm.stopPrank();

        // Staker stakes tokens
        vm.startPrank(staker);
        harberg.approve(address(stakingPool), stakeAmount);
        uint256[] memory empty;
        uint256 positionId = stakingPool.snatch(stakeAmount, staker, 1, empty);

        // Simulate time passage to accumulate tax liability
        vm.warp(block.timestamp + 3 days);

        // Calculate expected tax due at the moment of unstaking
        uint256 taxAmount = stakingPool.taxDue(positionId, 0);
        uint256 assetsAfterTax = stakeAmount - taxAmount;

        // Expect the PositionRemoved event with the expected parameters
        vm.expectEmit(true, true, true, true);
        emit PositionRemoved(positionId, staker, assetsAfterTax);

        // Perform unstaking
        stakingPool.exitPosition(positionId);

        // Check results after unstaking
        assertEq(harberg.balanceOf(staker), assetsAfterTax, "Assets after tax not returned correctly");
        assertEq(stakingPool.outstandingStake(), 0, "Outstanding stake not updated correctly");

        // Ensure the position is cleared
        (, address owner, uint32 time, , ) = stakingPool.positions(positionId);
        assertEq(time, 0, "Position time not cleared");
        assertEq(owner, address(0), "Position owner not cleared");

        vm.stopPrank();
    }

    function testSnatchFunction() public {
        uint256 initialStake1 = 1 ether;
        uint256 initialStake2 = 2 ether;
        address firstStaker = makeAddr("firstStaker");
        address secondStaker = makeAddr("secondStaker");
        address newStaker = makeAddr("newStaker");
        uint256 snatchAmount = 1.5 ether;

        // Mint and distribute tokens
        vm.startPrank(liquidityManager);
        harberg.mint((initialStake1 + initialStake2) * 5);
        harberg.transfer(firstStaker, initialStake1);
        harberg.transfer(secondStaker, initialStake2);
        harberg.transfer(newStaker, snatchAmount);
        vm.stopPrank();

        // Setup initial stakers
        uint256 positionId1 = setupStaker(firstStaker, initialStake1, 1);
        uint256 positionId2 = setupStaker(secondStaker, initialStake2, 5);

        // Snatch setup
        vm.startPrank(newStaker);
        harberg.approve(address(stakingPool), snatchAmount);
        uint256 snatchShares = stakingPool.assetsToShares(snatchAmount);
        uint256[] memory targetPositions = new uint256[](2);
        targetPositions[0] = positionId1;
        targetPositions[1] = positionId2;

        // Perform snatch
        uint256 newPositionId = stakingPool.snatch(snatchAmount, newStaker, 10, targetPositions);

        // Verify new position
        assertPosition(newPositionId, snatchShares, 10);

        // Check old positions
        verifyPositionShrunkOrRemoved(1, initialStake1);
        verifyPositionShrunkOrRemoved(2, initialStake2);

        vm.stopPrank();
    }

    function setupStaker(address staker, uint256 amount, uint32 taxRate) private returns (uint256 positionId) {
        vm.startPrank(staker);
        harberg.approve(address(stakingPool), amount);
        uint256[] memory empty;
        positionId = stakingPool.snatch(amount, staker, taxRate, empty);
        vm.stopPrank();
    }

    function assertPosition(uint256 positionId, uint256 expectedShares, uint32 expectedTaxRate) private view {
        (uint256 shares, , , , uint32 taxRate) = stakingPool.positions(positionId);
        assertEq(shares, expectedShares, "Incorrect share amount for new position");
        assertEq(taxRate, expectedTaxRate, "Incorrect tax rate for new position");
    }

    function verifyPositionShrunkOrRemoved(uint256 positionId, uint256 initialStake) private view {
        (uint256 remainingShare, , , , ) = stakingPool.positions(positionId);
        uint256 expectedInitialShares = stakingPool.assetsToShares(initialStake);
        bool positionRemoved = remainingShare == 0;
        bool positionShrunk = remainingShare < expectedInitialShares;

        assertTrue(positionRemoved || positionShrunk, "Position was not correctly shrunk or removed");
    }

    function testRevert_SharesTooLow() public {
        address staker = makeAddr("staker");
        vm.startPrank(liquidityManager);
        harberg.mint(10 ether);
        uint256 tooSmallStake = harberg.previousTotalSupply() / 4000;  // Less than minStake calculation
        harberg.transfer(staker, tooSmallStake);
        vm.stopPrank();

        vm.startPrank(staker);
        harberg.approve(address(stakingPool), tooSmallStake);

        uint256[] memory empty;
        vm.expectRevert(abi.encodeWithSelector(Stake.StakeTooLow.selector, staker, tooSmallStake, harberg.previousTotalSupply() / 3000));
        stakingPool.snatch(tooSmallStake, staker, 1, empty);
        vm.stopPrank();
    }

    function testRevert_TaxTooLow() public {
        address existingStaker = makeAddr("existingStaker");
        address newStaker = makeAddr("newStaker");
        vm.startPrank(liquidityManager);
        harberg.mint(10 ether);
        harberg.transfer(existingStaker, 1 ether);
        harberg.transfer(newStaker, 1 ether);
        vm.stopPrank();

        uint256 positionId = setupStaker(existingStaker, 1 ether, 5); // Existing staker with tax rate 5

        vm.startPrank(newStaker);
        harberg.transfer(newStaker, 1 ether);
        harberg.approve(address(stakingPool), 1 ether);

        uint256[] memory positions = new uint256[](1);
        positions[0] = positionId;  // Assuming position ID 1 has tax rate 5

        vm.expectRevert(abi.encodeWithSelector(Stake.TaxTooLow.selector, newStaker, 5, 5, positionId));
        stakingPool.snatch(1 ether, newStaker, 5, positions);  // Same tax rate should fail
        vm.stopPrank();
    }

    function testRevert_TooMuchSnatch() public {
        address staker = makeAddr("staker");
        address ambitiousStaker = makeAddr("ambitiousStaker");

        vm.startPrank(liquidityManager);
        harberg.mint(20 ether);
        harberg.transfer(staker, 2 ether);
        harberg.transfer(ambitiousStaker, 1 ether);
        vm.stopPrank();

        uint256 positionId = setupStaker(staker, 2 ether, 10);

        vm.startPrank(ambitiousStaker);
        harberg.approve(address(stakingPool), 1 ether);

        uint256[] memory positions = new uint256[](1);
        positions[0] = positionId;
        vm.expectRevert(abi.encodeWithSelector(TooMuchSnatch.selector, ambitiousStaker, 500000 ether, 1000000 ether, 1000000 ether));
        stakingPool.snatch(1 ether, ambitiousStaker, 20, positions);
        vm.stopPrank();
    }

    function testRevert_PositionNotFound() public {
        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        harberg.mint(10 ether);
        harberg.transfer(staker, 1 ether);
        vm.stopPrank();

        vm.startPrank(staker);
        harberg.approve(address(stakingPool), 1 ether);

        uint256[] memory nonExistentPositions = new uint256[](1);
        nonExistentPositions[0] = 999; // Assumed non-existent position ID

        vm.expectRevert(abi.encodeWithSelector(Stake.PositionNotFound.selector, 999, staker));
        stakingPool.snatch(1 ether, staker, 15, nonExistentPositions);
        vm.stopPrank();
    }

    function testSuccessfulChangeTaxAndReverts() public {
        uint32 newTaxRate = 3; // New valid tax rate

        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        harberg.mint(10 ether);
        harberg.transfer(staker, 1 ether);
        vm.stopPrank();

        vm.startPrank(staker);
        harberg.approve(address(stakingPool), 1 ether);
        uint256[] memory empty;
        uint256 positionId = stakingPool.snatch(1 ether, staker, 1, empty);

        uint32 invalidTaxRate = 9999; // Assuming this is out of bounds
        vm.expectRevert(bytes("tax rate out of bounds"));
        stakingPool.changeTax(positionId, invalidTaxRate);

        stakingPool.changeTax(positionId, newTaxRate);
        vm.stopPrank();

        // Verify the change
        (, , , , uint32 taxRate) = stakingPool.positions(positionId);
        assertEq(taxRate, newTaxRate, "Tax rate did not update correctly");

        // notOwner tries to change tax rate
        address notOwner = makeAddr("notOwner");
        vm.prank(notOwner);
        vm.expectRevert(abi.encodeWithSelector(Stake.NoPermission.selector, notOwner, staker));
        stakingPool.changeTax(positionId, newTaxRate * 2);
    }

    function testNormalTaxPayment() public {
        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        harberg.mint(10 ether);
        harberg.transfer(staker, 1 ether);
        vm.stopPrank();

        vm.startPrank(staker);
        harberg.approve(address(stakingPool), 1 ether);
        uint256[] memory empty;
        uint256 positionId = stakingPool.snatch(1 ether, staker, 5, empty);  // Using tax rate index 5, which is 18% per year
        (uint256 shareBefore, , , , ) = stakingPool.positions(positionId);


        // Immediately after staking, no tax due
        stakingPool.payTax(positionId);
        // Verify no change in position
        (uint256 share, , , , ) = stakingPool.positions(positionId);
        assertEq(share, shareBefore, "Share should not change when no tax is due");

        // Move time forward 30 days
        vm.warp(block.timestamp + 30 days);
        stakingPool.payTax(positionId);
        vm.stopPrank();

        // Check that the tax was paid and position updated
        (share, , , , ) = stakingPool.positions(positionId);
        uint256 daysElapsed = 30;
        uint256 taxRate = 18; // Corresponding to 18% annually
        uint256 daysInYear = 365;
        uint256 taxBase = 100;

        uint256 taxFractionForTime = taxRate * daysElapsed * 1 ether / daysInYear / taxBase;
        uint256 expectedShareAfterTax = (1 ether - taxFractionForTime) * 1000000;

        assertTrue(share < shareBefore, "Share should decrease correctly after tax payment 1");
        assertEq(share, expectedShareAfterTax, "Share should decrease correctly after tax payment 2");
    }

    function testLiquidation() public {
        address staker = makeAddr("staker");

        vm.startPrank(liquidityManager);
        harberg.mint(10 ether);
        harberg.transfer(staker, 1 ether);
        vm.stopPrank();

        vm.startPrank(staker);
        harberg.approve(address(stakingPool), 1 ether);
        uint256[] memory empty;
        uint256 positionId = stakingPool.snatch(1 ether, staker, 12, empty);  // Using tax rate index 5, which is 100% per year
        vm.warp(block.timestamp + 365 days); // Move time forward to ensure maximum tax due
        stakingPool.payTax(positionId);
        vm.stopPrank();

        // Verify position is liquidated
        (uint256 share, , , , ) = stakingPool.positions(positionId);
        assertEq(share, 0, "Share should be zero after liquidation");
    }

}