Tests: InfiniteGame2 & InfiniteTreasuryManager

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.27;

import "forge-std/Test.sol";
import "../src/InfiniteGame.sol";
import "../src/InfiniteGame2.sol";
import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import "@openzeppelin/contracts/access/Ownable.sol";

import {InfiniteToken} from "../src/InfiniteToken.sol";

uint256 constant TOTAL_SUPPLY = 1_000_000_000 ether;

// Contract for testing reentrancy protection
contract MaliciousToken is ERC20 {
    InfiniteGame2 private game2;

    constructor(address _game2) ERC20("Evil Token", "EVIL") {
        game2 = InfiniteGame2(_game2);
        _mint(address(this), TOTAL_SUPPLY);
    }

    function attack() external {
        approve(address(game2), type(uint256).max);
        game2.play();
    }

    function transferFrom(
        address sender,
        address recipient,
        uint256 amount
    ) public override returns (bool) {
        if (recipient == address(game2)) {
            game2.play(); // Attempt reentrancy
        }
        return super.transferFrom(sender, recipient, amount);
    }
}

contract InfiniteGame2Test is Test {
    InfiniteGame public game1;
    InfiniteGame2 public game2;
    InfiniteToken public gameToken;
    address public owner;
    address public player1;
    address public player2;

    function setUp() public {
        owner = address(0x42069);
        player1 = address(0x456);
        player2 = address(0x789);

        // Deploy mock token and game contract
        vm.startPrank(owner);

        gameToken = deployToken();

        game1 = new InfiniteGame(
            address(gameToken),
            block.timestamp,
            address(0)
        );
        game2 = new InfiniteGame2(
            address(gameToken),
            address(game1),
            block.timestamp
        );

        // Give tokens to test players
        gameToken.transfer(player1, 100_000_000 ether);
        gameToken.transfer(player2, 100_000_000 ether);
        vm.stopPrank();
    }

    function deployToken() public returns (InfiniteToken) {
        return new InfiniteToken(owner);
    }

    function test_InitialState() public view {
        assertEq(address(game2.gameToken()), address(gameToken));
        assertEq(address(game2.game1()), address(game1));
        assertEq(game2.startTime(), block.timestamp);
        assertEq(game2.roundStartTime(), block.timestamp);
        assertEq(game2.roundDeadline(), 0);
        assertEq(game2.lastPlayedAt(), 0);
        assertEq(game2.currentPrize(), 0);
        assertEq(game2.roundPlayCount(), 0);
        assertEq(game2.totalPlayCount(), 0);
        assertEq(game2.finalizedRounds(), 0);
        assertEq(game2.grandTotalFees(), 0);
        assertEq(game2.playerTotalFees(player1), 0);
        assertEq(game2.playerTotalFees(player2), 0);
        assertEq(game2.latestPlayer(), address(0));
    }

    function test_PlayerPlayFees() public {
        uint256 fee1 = game2.playFee();

        vm.startPrank(player1);
        gameToken.approve(address(game2), type(uint256).max);

        // First play
        game2.play();

        uint256 playerPlayFees = game2.playerTotalFees(player1);

        assertEq(playerPlayFees, fee1);
        assertEq(game2.grandTotalFees(), fee1);

        uint256 fee2 = game2.playFee();

        // Second play
        game2.play();

        playerPlayFees = game2.playerTotalFees(player1);
        assertEq(playerPlayFees, fee1 + fee2);
        assertEq(game2.grandTotalFees(), fee1 + fee2);
        vm.stopPrank();
    }

    function test_FeeIncreasesAccurately() public {
        vm.prank(owner);

        uint256 testReferenceFee = game2.playFee();

        for (uint i = 0; i < 200; i++) {
            vm.startPrank(player1);
            gameToken.approve(address(game2), game2.playFee());
            game2.play();
            vm.stopPrank();

            uint256 currentFee = game2.playFee();

            testReferenceFee = (testReferenceFee * 105) / 100;

            // Fee should increase 5% each play
            assertEq(currentFee, testReferenceFee);
        }
    }

    function test_CompleteRound() public {
        // Fund game with arbitrary amount
        vm.prank(owner);
        gameToken.transfer(address(game2), 1000 ether);

        // Player 1 plays
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        uint256 initContractBalance = gameToken.balanceOf(address(game2));
        uint256 initPlayer1Balance = gameToken.balanceOf(player1);

        // Fast forward past deadline
        vm.warp(block.timestamp + game2.countdown());

        // Finalize round
        game2.completeRound();

        // Check prize distribution
        uint256 expectedWinnerPrize = (initContractBalance *
            game2.winnerShare()) / 100;
        uint256 expectedRetained = initContractBalance - expectedWinnerPrize;

        // Verify distributions
        assertEq(
            gameToken.balanceOf(player1),
            initPlayer1Balance + expectedWinnerPrize,
            "Winner prize incorrect"
        );
        assertEq(
            gameToken.balanceOf(address(game2)),
            expectedRetained,
            "Retained amount incorrect"
        );

        assertEq(game2.roundPlayCount(), 0, "Play count not reset");
        assertEq(game2.totalPlayCount(), 1, "Total play count not incremented");
    }

    function test_RoundCompletedEventEmission() public {
        // Fund the contract
        vm.prank(owner);
        gameToken.transfer(address(game2), 1000 ether);

        // Player plays
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        uint256 gameBalance = gameToken.balanceOf(address(game2));

        // Calculate expected amounts
        uint256 expectedWinnerPrize = (gameBalance * game2.winnerShare()) / 100;
        uint256 expectedRetained = gameBalance - expectedWinnerPrize;

        // Move to round end
        vm.warp(block.timestamp + game2.countdown());

        // Expect RoundComplete event with correct parameters
        vm.expectEmit(true, true, true, true, address(game2));
        emit IInfiniteGame2.RoundCompleted(
            1,
            player1,
            expectedWinnerPrize,
            expectedRetained
        );
        game2.completeRound();
    }

    function test_FinalizedRoundsInitializedUsingGame1() public {
        // play a round on game 1
        vm.startPrank(player1);
        gameToken.approve(address(game1), game1.playFee());
        game1.play();
        vm.stopPrank();

        // Fast forward past deadline
        vm.warp(block.timestamp + game1.countDown());
        game1.completeRound();

        // Verify finalized rounds is 1
        assertEq(game1.finalizedRounds(), 1);

        // Play and complete a round on game 2
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();

        uint256 gameBalance = gameToken.balanceOf(address(game2));
        uint256 expectedWinnerPrize = (gameBalance * game2.winnerShare()) / 100;
        uint256 expectedRetained = gameBalance - expectedWinnerPrize;

        vm.warp(block.timestamp + game2.countdown());

        vm.expectEmit(true, true, true, true, address(game2));
        emit IInfiniteGame2.RoundCompleted(
            2,
            player1,
            expectedWinnerPrize,
            expectedRetained
        );

        game2.completeRound();

        // Verify finalized rounds is 2
        assertEq(game2.finalizedRounds(), 2);
    }

    function test_CompleteRound_RevertIfRoundNotOver() public {
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        vm.warp(block.timestamp + game2.countdown() - 1);
        vm.expectRevert(IInfiniteGame2.RoundStillActive.selector);
        game2.completeRound();
    }

    function test_MultipleRounds() public {
        vm.startPrank(owner);

        // Fund game with substantial amount for prizes
        gameToken.transfer(address(game2), 1000 ether);
        vm.stopPrank();

        uint256 player1initBalance = gameToken.balanceOf(player1);
        uint256 playFee1 = game2.playFee();

        // Round 1
        vm.startPrank(player1);
        gameToken.approve(address(game2), 1000 ether);
        game2.play();
        vm.stopPrank();

        vm.warp(block.timestamp + game2.countdown());
        game2.completeRound();

        // Verify Round 1 completion
        uint256 player1NewBalance = gameToken.balanceOf(player1);
        assertTrue(
            player1NewBalance > player1initBalance - playFee1,
            "Player 1 should profit from winning"
        );
        assertEq(game2.grandTotalFees(), playFee1);
        assertEq(game2.playerTotalFees(player1), playFee1);

        uint256 player2initBalance = gameToken.balanceOf(player2);

        uint256 playFee2 = game2.playFee();

        vm.startPrank(player2);
        gameToken.approve(address(game2), 1000 ether);
        game2.play();
        vm.stopPrank();

        vm.warp(block.timestamp + game2.countdown());
        game2.completeRound();

        uint256 player2NewBalance = gameToken.balanceOf(player2);

        // Player 2's final balance should be greater than their init balance minus play fee
        assertTrue(
            player2NewBalance > player2initBalance - playFee2,
            "Player 2 should profit from winning"
        );
        assertEq(game2.grandTotalFees(), playFee1 + playFee2);
        assertEq(game2.totalPlayCount(), 2);
    }

    function test_Play_RevertIfRoundNotCompleted() public {
        // Round 1
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        // Player wins game
        vm.warp(block.timestamp + game2.countdown());

        // Player 2 attempts play
        vm.startPrank(player2);
        gameToken.approve(address(game2), game2.playFee());
        vm.expectRevert(IInfiniteGame2.NextRoundPending.selector);
        game2.play();
        vm.stopPrank();
    }

    function test_ArbitraryTimeBeforeCompleting() public {
        // Round 1
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        // Player wins game
        vm.warp(block.timestamp + 42069 days);

        // Simulate player 2 playing
        vm.startPrank(player2);
        gameToken.approve(address(game2), game2.playFee());
        // Should revert
        vm.expectRevert(IInfiniteGame2.NextRoundPending.selector);
        game2.play();
        vm.stopPrank();

        uint256 player1Balance = gameToken.balanceOf(player1);

        uint256 gameBalanceBeforePayout = gameToken.balanceOf(address(game2));

        game2.completeRound();

        // Verify player 1 received prize
        assertEq(
            gameToken.balanceOf(player1),
            player1Balance +
                (gameBalanceBeforePayout * game2.winnerShare()) /
                100,
            "Player 1 should receive prize"
        );
    }

    function test_Play_RevertIfNoBalance() public {
        // Player 2 attempts play
        vm.startPrank(address(2397092));
        gameToken.approve(address(game2), game2.playFee());
        vm.expectRevert();
        game2.play();
        vm.stopPrank();
    }

    function test_Play_RevertIfBeforeStartTime() public {
        InfiniteGame2 newGame2 = new InfiniteGame2(
            address(gameToken),
            address(game1),
            block.timestamp + 10
        );

        // Give player some tokens
        vm.prank(owner);
        gameToken.transfer(player1, 1000);

        // Player 1 attempts play
        vm.startPrank(player1);
        gameToken.approve(address(newGame2), newGame2.playFee());
        vm.expectRevert(IInfiniteGame2.GameNotStarted.selector);
        newGame2.play();
        vm.stopPrank();
    }

    function test_WinnerWithdrawERC20_RevertIfNotOwner() public {
        InfiniteToken extraToken = deployToken();

        vm.prank(owner);
        extraToken.transfer(address(game2), 1000);

        vm.prank(player2);
        vm.expectRevert();
        game2.withdrawERC20(address(extraToken));
    }

    function test_WithdrawERC20_RevertIfGameToken() public {
        vm.startPrank(owner);

        // Try to withdraw game token
        vm.expectRevert(IInfiniteGame2.CannotWithdrawGameToken.selector);
        game2.withdrawERC20(address(gameToken));
        vm.stopPrank();
    }

    function test_WithdrawERC20_RevertIfNoBalance() public {
        InfiniteToken extraToken = deployToken();

        vm.prank(owner);
        vm.expectRevert(IInfiniteGame2.NoTokenBalance.selector);
        game2.withdrawERC20(address(extraToken));
    }

    function test_ReentrancyProtection() public {
        // Deploy malicious token that attempts reentrancy
        MaliciousToken malToken = new MaliciousToken(address(game2));

        vm.prank(owner);

        // Try to reenter during play
        vm.expectRevert();
        malToken.attack();
    }

    function test_PrizeCalculationRounding() public {
        // Send an amount that could cause rounding issues
        vm.prank(owner);
        gameToken.transfer(address(game2), 100);

        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        vm.warp(block.timestamp + game2.countdown());

        uint256 balanceBefore = gameToken.balanceOf(address(game2));
        game2.completeRound();
        uint256 balanceAfter = gameToken.balanceOf(address(game2));

        assertTrue(
            balanceAfter <= balanceBefore,
            "Prize distribution should not create tokens"
        );
    }

    function test_ConsecutiveRoundsWithSamePlayer() public {
        // First round
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        vm.warp(block.timestamp + game2.countdown());
        game2.completeRound();

        // Same player tries next round
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        assertTrue(
            game2.latestPlayer() == player1,
            "Same player should be able to play consecutive rounds"
        );
    }

    function test_EmergencyTokenRecovery() public {
        vm.startPrank(owner);
        InfiniteToken randomToken = deployToken();

        // Test recovery of accidentally sent tokens
        uint256 ownerBalanceBefore = randomToken.balanceOf(owner);

        randomToken.transfer(address(game2), 1000);

        // Owner should be able to recover random tokens
        game2.withdrawERC20(address(randomToken));

        vm.stopPrank();

        assertEq(
            randomToken.balanceOf(owner),
            ownerBalanceBefore,
            "Owner should receive recovered tokens"
        );
    }

    function test_PlayedEventEmission() public {
        vm.prank(owner);

        uint256 playFee = game2.playFee();

        vm.startPrank(player1);
        gameToken.approve(address(game2), playFee);

        vm.expectEmit(true, false, false, true, address(game2));
        emit IInfiniteGame2.Played(player1, playFee, 1);
        game2.play();
        vm.stopPrank();
    }

    // Tests the game won't brick if nobody plays at the start of a round
    function test_CompleteRound_NoPlayers() public {
        // Assert this is the first round
        assertEq(game2.finalizedRounds(), 0);

        // Warp to end of round + arbitrary time to simulate being able to continue the game
        // no matter now long nobody has played
        vm.warp(game2.roundDeadline() + 1000 weeks);

        assertEq(game2.latestPlayer(), address(0));

        // Can't complete the round until someone plays
        vm.expectRevert(IInfiniteGame2.RoundStillActive.selector);
        game2.completeRound();

        // Player can still play even though countdown has expired
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        uint256 gameBalanceBefore = gameToken.balanceOf(address(game2));
        uint256 prizeBeforeRoundEnd = game2.currentPrize();
        uint256 player1BalanceBefore = gameToken.balanceOf(player1);

        // Still need to wait for new countdown to expire
        vm.warp(game2.roundDeadline());

        // Now round can end
        game2.completeRound();

        // Player received round prize
        assertEq(
            gameToken.balanceOf(address(game2)),
            gameBalanceBefore - prizeBeforeRoundEnd
        );
        assertApproxEqAbs(
            gameToken.balanceOf(player1),
            player1BalanceBefore + prizeBeforeRoundEnd,
            0
        );

        // New round - player should be able to play
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        // Verify round count
        assertEq(game2.finalizedRounds(), 1);
    }

    function test_CurrentPrize_Calculation() public {
        // Test with zero balance
        assertEq(game2.currentPrize(), 0);

        // Test with specific balance
        vm.prank(owner);
        gameToken.transfer(address(game2), 1000 ether);

        // Safely calculate expected prize: (balance * winnerShare) / 100
        uint256 balance = gameToken.balanceOf(address(game2));
        uint256 expectedPrize = (balance / 100) * game2.winnerShare();

        assertEq(
            game2.currentPrize(),
            expectedPrize,
            "Prize calculation incorrect"
        );
    }

    function test_Play_MultiplePlayers() public {
        // First player
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        assertEq(game2.latestPlayer(), player1);

        // Second player
        vm.startPrank(player2);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        assertEq(game2.latestPlayer(), player2);
        assertEq(game2.roundPlayCount(), 2);
        assertEq(game2.totalPlayCount(), 2);
    }

    function test_Play_ApprovalRequired() public {
        vm.startPrank(player1);

        // Try to play without approving tokens
        vm.expectRevert();
        game2.play();

        // Approve less than required
        gameToken.approve(address(game2), game2.playFee() - 1);
        vm.expectRevert();
        game2.play();

        vm.stopPrank();
    }

    function test_SetCountdown_DuringActiveRound() public {
        // Start a round
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        // Should not update until round completes
        vm.prank(owner);
        game2.setCountdown(2 hours);
        assertEq(game2.countdown(), 1 hours);

        // Complete round and verify update
        vm.warp(block.timestamp + game2.countdown());

        vm.expectEmit(true, true, true, true, address(game2));
        emit IInfiniteGame2.CountdownUpdated(2 hours);

        game2.completeRound();
        assertEq(game2.countdown(), 2 hours);
    }

    function test_SetCountdown_RevertIfInvalid() public {
        vm.startPrank(owner);
        vm.expectRevert(IInfiniteGame2.InvalidValue.selector);
        game2.setCountdown(59); // Must be at least 60 seconds

        vm.expectRevert(IInfiniteGame2.InvalidValue.selector);
        game2.setCountdown(86401); // 1 day is the max
        vm.stopPrank();
    }

    function test_SetInterestRate_DuringActiveRound() public {
        // Start a round
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        // Should not update until round completes
        vm.prank(owner);
        game2.setInterestRate(1100);
        assertEq(game2.interestRate(), 1050);

        // Complete round and verify update
        vm.warp(block.timestamp + game2.countdown());

        vm.expectEmit(true, true, true, true, address(game2));
        emit IInfiniteGame2.InterestRateUpdated(1100);

        game2.completeRound();
        assertEq(game2.interestRate(), 1100);
    }

    function test_SetInterestRate_RevertIfInvalid() public {
        vm.startPrank(owner);
        vm.expectRevert(IInfiniteGame2.InvalidValue.selector);
        game2.setInterestRate(0);

        vm.expectRevert(IInfiniteGame2.InvalidValue.selector);
        game2.setInterestRate(1501);

        vm.stopPrank();
    }

    function test_SetWinnerShare_DuringActiveRound() public {
        // Start a round
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        // Should not update until round completes
        vm.prank(owner);
        game2.setWinnerShare(60);
        assertEq(game2.winnerShare(), 50);

        // Complete round and verify update
        vm.warp(block.timestamp + game2.countdown());

        vm.expectEmit(true, true, true, true, address(game2));
        emit IInfiniteGame2.WinnerShareUpdated(60);
        game2.completeRound();
        assertEq(game2.winnerShare(), 60);
    }

    function test_SetWinnerShare_RevertIfInvalid() public {
        vm.startPrank(owner);
        vm.expectRevert(IInfiniteGame2.InvalidValue.selector);
        game2.setWinnerShare(0);
        vm.expectRevert(IInfiniteGame2.InvalidValue.selector);
        game2.setWinnerShare(100);
        vm.stopPrank();
    }

    function test_SetInitialPlayFee_DuringActiveRound() public {
        // Start a round
        vm.startPrank(player1);
        gameToken.approve(address(game2), game2.playFee());
        game2.play();
        vm.stopPrank();

        // Should not update until round completes
        vm.prank(owner);
        game2.setInitialPlayFee(2000);
        assertEq(game2.initialPlayFee(), 100 ether);

        // Complete round and verify update
        vm.warp(block.timestamp + game2.countdown());

        vm.expectEmit(true, true, true, true, address(game2));
        emit IInfiniteGame2.InitialPlayFeeUpdated(2000);

        game2.completeRound();

        assertEq(game2.initialPlayFee(), 2000);
    }

    function test_SetInitialPlayFee_RevertIfInvalid() public {
        vm.startPrank(owner);
        vm.expectRevert(IInfiniteGame2.InvalidValue.selector);
        game2.setInitialPlayFee(0);
        vm.expectRevert(IInfiniteGame2.InvalidValue.selector);
        game2.setInitialPlayFee(100_001);
        vm.stopPrank();
    }

    receive() external payable {}
}

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "@uniswap/v3-periphery/contracts/interfaces/INonfungiblePositionManager.sol";
import "@uniswap/v3-core/contracts/interfaces/IUniswapV3Pool.sol";

import "./MockERC20.sol";

contract MockNonfungiblePositionManager {
    mapping(uint256 => address) public ownerOf;
    mapping(uint256 => Position) private _positions;

    struct Position {
        address token0;
        address token1;
        uint256 fees0;
        uint256 fees1;
    }

    function setOwner(uint256 tokenId, address owner) external {
        ownerOf[tokenId] = owner;
    }

    // This matches Uniswap's full interface
    function positions(
        uint256 tokenId
    )
        external
        view
        returns (
            uint96 nonce,
            address operator,
            address token0,
            address token1,
            uint24 fee,
            int24 tickLower,
            int24 tickUpper,
            uint128 liquidity,
            uint256 feeGrowthInside0LastX128,
            uint256 feeGrowthInside1LastX128,
            uint128 tokensOwed0,
            uint128 tokensOwed1
        )
    {
        Position memory pos = _positions[tokenId];
        return (
            0, // nonce
            address(0), // operator
            pos.token0,
            pos.token1,
            3000, // fee
            0, // tickLower
            0, // tickUpper
            0, // liquidity
            0, // feeGrowthInside0LastX128
            0, // feeGrowthInside1LastX128
            uint128(pos.fees0), // tokensOwed0
            uint128(pos.fees1) // tokensOwed1
        );
    }

    function setPosition(
        uint256 tokenId,
        address token0,
        address token1,
        uint256 fees0,
        uint256 fees1
    ) external {
        _positions[tokenId] = Position(token0, token1, fees0, fees1);
    }

    function collect(
        INonfungiblePositionManager.CollectParams calldata params
    ) external returns (uint256 amount0, uint256 amount1) {
        Position storage pos = _positions[params.tokenId];
        require(
            pos.token0 != address(0) && pos.token1 != address(0),
            "Invalid position"
        );

        amount0 = pos.fees0;
        amount1 = pos.fees1;

        if (amount0 > 0) {
            MockERC20(pos.token0).transfer(params.recipient, amount0);
        }
        if (amount1 > 0) {
            MockERC20(pos.token1).transfer(params.recipient, amount1);
        }

        // Reset fees after collection
        pos.fees0 = 0;
        pos.fees1 = 0;
    }
}

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "forge-std/Test.sol";
import "../src/InfiniteTreasuryManager.sol";
import "./mocks/MockNonfungiblePositionManager.sol";

import "@openzeppelin/contracts/token/ERC20/ERC20.sol";

contract MockInfiniteGameV1 {
    mapping(address => uint256) private _playerPlays;
    uint256 private _totalPlays;

    function setPlayerCount(address player, uint256 plays) external {
        _playerPlays[player] = plays;
    }

    function setTotalPlayCount(uint256 plays) external {
        _totalPlays = plays;
    }

    function playerPlayCount(address player) external view returns (uint256) {
        return _playerPlays[player];
    }

    function totalPlayCount() external view returns (uint256) {
        return _totalPlays;
    }
}

contract MockInfiniteGame2 {
    mapping(address => uint256) private _playerFees;
    uint256 private _totalFees;
    address private _owner;

    constructor() {
        _owner = msg.sender;
    }

    function setPlayerFees(address player, uint256 fees) external {
        _playerFees[player] = fees;
    }

    function setTotalFees(uint256 fees) external {
        _totalFees = fees;
    }

    function playerTotalFees(address player) external view returns (uint256) {
        return _playerFees[player];
    }

    function grandTotalFees() external view returns (uint256) {
        return _totalFees;
    }

    function owner() external view returns (address) {
        return _owner;
    }

    function setOwner(address newOwner) external {
        _owner = newOwner;
    }
}

contract MockTreasury {
    // For the sake of simplicity, we're not including the donation
    // that happens here in the real contract
    function withdraw(address token, address to, uint256 amount) external {
        MockERC20(token).transfer(to, amount);
    }

    function collectFees(uint256 tokenId) external {}
}

contract MockERC721 {
    mapping(uint256 => address) private _owners;

    function mint(address to, uint256 tokenId) external {
        _owners[tokenId] = to;
    }

    function ownerOf(uint256 tokenId) external view returns (address) {
        return _owners[tokenId];
    }

    function transferFrom(address from, address to, uint256 tokenId) external {
        require(_owners[tokenId] == from, "Not owner");
        _owners[tokenId] = to;
    }
}

contract InfiniteTreasuryManagerTest is Test {
    MockERC20 public usdc;
    MockERC20 public dai;
    MockERC20 public weth;
    MockERC20 public infiniteToken;
    MockInfiniteGameV1 public gameV1;
    MockInfiniteGame2 public gameV2;
    MockTreasury public treasury;
    InfiniteTreasuryManager public manager;
    MockNonfungiblePositionManager public positionManager;

    address public constant ZERO_ADDRESS = address(0);
    uint256 public constant POSITION_ID = 1;

    mapping(address => uint256) private playerPlayCounts;

    function setUp() public {
        usdc = new MockERC20();
        dai = new MockERC20();
        weth = new MockERC20();
        infiniteToken = new MockERC20();
        gameV1 = new MockInfiniteGameV1();
        gameV2 = new MockInfiniteGame2();
        treasury = new MockTreasury();
        positionManager = new MockNonfungiblePositionManager();

        uint256 infiniteBalance = 5313 ether;
        uint256 wethBalance = 459 ether;

        // Create position
        positionManager.setPosition(
            POSITION_ID,
            address(infiniteToken),
            address(weth),
            infiniteBalance,
            wethBalance
        );

        positionManager.setOwner(POSITION_ID, address(treasury));

        manager = new InfiniteTreasuryManager(
            address(treasury),
            address(positionManager),
            address(gameV1),
            address(gameV2),
            address(infiniteToken)
        );
    }

    function test_Constructor_InitialState() public view {
        assertEq(manager.infiniteTreasury(), address(treasury));
        assertEq(manager.infiniteGame1(), address(gameV1));
        assertEq(manager.infiniteGame2(), address(gameV2));
        assertEq(manager.infiniteToken(), address(infiniteToken));
    }

    function test_Constructor_RevertZeroAddress() public {
        vm.expectRevert(
            abi.encodeWithSelector(InfiniteTreasuryManager.ZeroAddress.selector)
        );
        new InfiniteTreasuryManager(
            ZERO_ADDRESS,
            address(positionManager),
            address(gameV1),
            address(gameV2),
            address(infiniteToken)
        );

        vm.expectRevert(
            abi.encodeWithSelector(InfiniteTreasuryManager.ZeroAddress.selector)
        );
        new InfiniteTreasuryManager(
            address(treasury),
            ZERO_ADDRESS,
            address(gameV1),
            address(gameV2),
            address(infiniteToken)
        );

        vm.expectRevert(
            abi.encodeWithSelector(InfiniteTreasuryManager.ZeroAddress.selector)
        );
        new InfiniteTreasuryManager(
            address(treasury),
            address(positionManager),
            address(gameV1),
            address(gameV2),
            ZERO_ADDRESS
        );
    }

    function test_playerWithdrawals_MultipleCollections() public {
        // Test with 3 players
        address player1 = address(1);
        address player2 = address(2);
        address player3 = address(3);

        // Set up V1 play counts
        uint256 totalPlays = 100;
        gameV1.setTotalPlayCount(totalPlays);
        gameV1.setPlayerCount(player1, 20); // 20% of plays
        gameV1.setPlayerCount(player2, 30); // 30% of plays
        gameV1.setPlayerCount(player3, 50); // 50% of plays

        // Set up V2 fees
        uint256 totalFees = 1000 ether;
        gameV2.setTotalFees(totalFees);
        gameV2.setPlayerFees(player1, 500 ether); // 50% of fees
        gameV2.setPlayerFees(player2, 300 ether); // 30% of fees
        gameV2.setPlayerFees(player3, 200 ether); // 20% of fees

        // First collection - should go to V1
        uint256 firstCollection = 1000 ether;
        weth.transfer(address(treasury), firstCollection);
        manager.collectAndDistribute(POSITION_ID);

        // Second collection - should go to V2
        uint256 secondCollection = 2000 ether;
        weth.transfer(address(treasury), secondCollection);
        manager.collectAndDistribute(POSITION_ID);

        // Test player1 withdrawable
        // Expected: Should be (20% of 1000 ETH) + (50% of 2000 ETH) = 200 + 1000 = 1200 ETH
        // Current test expectation: (20% of 1000 ETH + 50% of 2000 ETH) / 2 = 600 ETH
        uint256 player1Withdrawable = manager.withdrawable(
            player1,
            address(weth)
        );

        // The actual value from the contract should be 1200 ETH
        assertEq(
            player1Withdrawable,
            1200 ether,
            "Player 1 withdrawable incorrect"
        );

        // Test player2 withdrawable
        // Expected: Should be (30% of 1000 ETH) + (30% of 2000 ETH) = 300 + 600 = 900 ETH
        // Current test expectation: (30% of 1000 ETH + 30% of 2000 ETH) / 2 = 450 ETH
        uint256 player2Withdrawable = manager.withdrawable(
            player2,
            address(weth)
        );

        // The actual value from the contract should be 900 ETH
        assertEq(
            player2Withdrawable,
            900 ether,
            "Player 2 withdrawable incorrect"
        );

        // Test player3 withdrawable
        // Expected: Should be (50% of 1000 ETH) + (20% of 2000 ETH) = 500 + 400 = 900 ETH
        // Current test expectation: (50% of 1000 ETH + 20% of 2000 ETH) / 2 = 450 ETH
        uint256 player3Withdrawable = manager.withdrawable(
            player3,
            address(weth)
        );

        // The actual value from the contract should be 900 ETH
        assertEq(
            player3Withdrawable,
            900 ether,
            "Player 3 withdrawable incorrect"
        );
    }

    function test_Withdraw_RevertZeroAddress() public {
        vm.expectRevert(
            abi.encodeWithSelector(InfiniteTreasuryManager.ZeroAddress.selector)
        );
        manager.withdraw(ZERO_ADDRESS, address(weth));
    }

    function test_Withdraw_RevertInsufficientBalance() public {
        // No deposits made
        vm.expectRevert(abi.encodeWithSignature("InsufficientBalance()"));
        manager.withdraw(address(this), address(weth));
    }

    function test_Withdrawable_NoDeposits() public {
        gameV1.setTotalPlayCount(100);
        gameV1.setPlayerCount(address(this), 50);
        gameV2.setTotalFees(100 ether);
        gameV2.setPlayerFees(address(this), 50 ether);

        assertEq(manager.withdrawable(address(this), address(weth)), 0);
    }
    function test_Withdrawable_AlreadyWithdrawn() public {
        // Set up initial state
        uint256 wethAmount = 100 ether;
        weth.transfer(address(treasury), wethAmount);

        // Set equal proportions in V1 and V2
        gameV1.setTotalPlayCount(100);
        gameV1.setPlayerCount(address(this), 50);
        gameV2.setTotalFees(100 ether);
        gameV2.setPlayerFees(address(this), 50 ether);

        manager.collectAndDistribute(POSITION_ID);

        uint256 withdrawable = manager.withdrawable(
            address(this),
            address(weth)
        );
        assertEq(withdrawable, 50 ether); // Should be 50% in both V1 and V2, averaged to 50%

        manager.withdraw(address(this), address(weth));

        // Should be 0 after withdrawal
        assertEq(manager.withdrawable(address(this), address(weth)), 0);
    }

    function test_Withdrawable_NoV2Activity() public {
        uint256 wethAmount = 100 ether;
        weth.transfer(address(treasury), wethAmount);

        // Only V1 activity
        gameV1.setTotalPlayCount(100);
        gameV1.setPlayerCount(address(this), 50);
        gameV2.setTotalFees(0); // No V2 activity

        manager.collectAndDistribute(POSITION_ID);

        // Should use only V1 share (50%) since V2 has no activity
        assertEq(manager.withdrawable(address(this), address(weth)), 50 ether);
    }

    function test_Withdrawable_NoV1Activity() public {
        uint256 wethAmount = 100 ether;
        weth.transfer(address(treasury), wethAmount);

        // Only V2 activity
        gameV1.setTotalPlayCount(0);
        gameV2.setTotalFees(1000 ether);
        gameV2.setPlayerFees(address(this), 500 ether);

        manager.collectAndDistribute(POSITION_ID);

        uint256 withdrawable = manager.withdrawable(
            address(this),
            address(weth)
        );

        console.log("Withdrawable: %d", withdrawable);

        // Should be zero because player didn't play V1
        assertEq(withdrawable, 0);
    }

    function test_Withdrawable_MixedActivity() public {
        uint256 v1UniswapFees = 1000 ether;
        weth.transfer(address(treasury), v1UniswapFees);

        address player1 = address(1);
        address player2 = address(2);

        // Different shares in V1 and V2
        gameV1.setTotalPlayCount(100);
        gameV1.setPlayerCount(player1, 25); // 25% in V1
        gameV1.setPlayerCount(player2, 75); // 75% in V1

        manager.collectAndDistribute(POSITION_ID);

        // First collection should only be V1 share (25%)
        assertEq(manager.withdrawable(player1, address(weth)), 250 ether);
        assertEq(manager.withdrawable(player2, address(weth)), 750 ether);

        // Player 1 withdraws
        manager.withdraw(player1, address(weth));
        assertEq(weth.balanceOf(player1), 250 ether);

        // V2 starts after V1 fees have been collected
        uint256 v2UniswapFees = 1000 ether;
        weth.transfer(address(treasury), v2UniswapFees);

        // Set V2 playing fees
        gameV2.setTotalFees(100 ether);
        gameV2.setPlayerFees(player1, 60 ether);
        gameV2.setPlayerFees(player2, 40 ether);

        // Distribute v2 uniswap fees
        manager.collectAndDistribute(POSITION_ID);

        // Player 1 has withdrawn V1 uniswap fees, so should only get 60% of V2
        assertEq(manager.withdrawable(player1, address(weth)), 600 ether);

        // Player 2 hasn't withdrawn v2 uniswap yet, so should be 75% of V1 + 40% of V2
        assertEq(
            manager.withdrawable(player2, address(weth)),
            750 ether + 400 ether
        );
    }

    function test_Withdrawable_SmallAmounts() public {
        // Test with 1 wei to verify rounding behavior
        weth.transfer(address(treasury), 1);

        gameV1.setTotalPlayCount(3);
        gameV1.setPlayerCount(address(this), 1); // 33.33...% in V1
        gameV2.setTotalFees(3);
        gameV2.setPlayerFees(address(this), 1); // 33.33...% in V2

        manager.collectAndDistribute(POSITION_ID);

        // With floor rounding, should get 0
        assertEq(manager.withdrawable(address(this), address(weth)), 0);
    }

    // Add these tests to InfiniteTreasuryManagerTest

    function test_CollectAndDistribute_InfiniteTokenFees() public {
        // Test collecting Infinite tokens
        uint256 infiniteTokenAmount = 1000 ether;
        infiniteToken.transfer(address(treasury), infiniteTokenAmount);

        // Should transfer to V2 game
        manager.collectAndDistribute(POSITION_ID);

        assertEq(infiniteToken.balanceOf(address(gameV2)), infiniteTokenAmount);
    }

    function test_CollectAndDistribute_EmitsEvent() public {
        uint256 wethAmount = 1000 ether;
        uint256 infiniteTokenAmount = 500 ether;

        weth.transfer(address(treasury), wethAmount);
        infiniteToken.transfer(address(treasury), infiniteTokenAmount);

        vm.expectEmit(true, true, true, true);
        emit InfiniteTreasuryManager.FeesCollected(
            POSITION_ID,
            address(infiniteToken),
            address(weth),
            infiniteTokenAmount,
            wethAmount
        );

        manager.collectAndDistribute(POSITION_ID);
    }

    function test_CollectAndDistribute_ZeroFees() public {
        // Should not revert when collecting zero fees
        manager.collectAndDistribute(POSITION_ID);

        assertEq(weth.balanceOf(address(manager)), 0);
        assertEq(infiniteToken.balanceOf(address(gameV2)), 0);
    }

    function test_GetterFunctions() public view {
        assertEq(manager.infiniteTreasury(), address(treasury));
        assertEq(manager.infiniteGame1(), address(gameV1));
        assertEq(manager.infiniteToken(), address(infiniteToken));
    }

    function test_MultipleWithdrawals() public {
        uint256 wethAmount = 1000 ether;
        weth.transfer(address(treasury), wethAmount);

        address player = address(12345);

        gameV1.setTotalPlayCount(100);
        gameV1.setPlayerCount(player, 50);

        manager.collectAndDistribute(POSITION_ID);

        // First withdrawal
        manager.withdraw(player, address(weth));
        assertEq(weth.balanceOf(player), 500 ether);

        // Second collection
        weth.transfer(address(treasury), wethAmount);
        manager.collectAndDistribute(POSITION_ID);

        // Second withdrawal - should revert since player has no activity on V2 game
        vm.expectRevert(
            abi.encodeWithSelector(
                InfiniteTreasuryManager.InsufficientBalance.selector
            )
        );
        manager.withdraw(player, address(weth));

        // Play on V2
        gameV2.setTotalFees(1000 ether);
        gameV2.setPlayerFees(player, 500 ether);

        // Collect and distribute V2 fees
        manager.collectAndDistribute(POSITION_ID);

        // Second withdrawal should now succeed
        manager.withdraw(player, address(weth));

        assertEq(weth.balanceOf(player), 1000 ether);
    }

    function test_WithdrawableWithMaxPlayers() public {
        uint256 numPlayers = 1000; // Or higher, test gas limits
        uint256 wethAmount = 1000 ether;
        uint256 playShare = 1;

        weth.transfer(address(treasury), wethAmount);

        // Set up many players with small shares
        for (uint256 i = 1; i <= numPlayers; i++) {
            address player = address(uint160(i));
            gameV1.setPlayerCount(player, playShare);
        }
        gameV1.setTotalPlayCount(numPlayers * playShare);

        manager.collectAndDistribute(POSITION_ID);

        // Test a few random players can withdraw
        for (uint256 i = 1; i <= 10; i++) {
            address player = address(uint160(i));
            uint256 expectedAmount = wethAmount / numPlayers;
            assertEq(
                manager.withdrawable(player, address(weth)),
                expectedAmount
            );
        }
    }

    function test_WithdrawableWithPrimeAmounts() public {
        // Test with prime numbers to force recurring decimals
        uint256 wethAmount = 1000 ether;
        uint256 totalPlays = 7; // Prime number

        weth.transfer(address(treasury), wethAmount);

        gameV1.setTotalPlayCount(totalPlays);
        gameV1.setPlayerCount(address(1), 2);
        gameV1.setPlayerCount(address(2), 3);
        gameV1.setPlayerCount(address(3), 2);

        manager.collectAndDistribute(POSITION_ID);

        // Verify rounding behavior
        uint256 expected1 = (wethAmount * 2) / totalPlays;
        uint256 expected2 = (wethAmount * 3) / totalPlays;

        assertEq(manager.withdrawable(address(1), address(weth)), expected1);
        assertEq(manager.withdrawable(address(2), address(weth)), expected2);
    }

    function test_MultiplePartialWithdrawals() public {
        uint256 wethAmount = 1000 ether;
        address player = address(1);

        // Setup initial state
        weth.transfer(address(treasury), wethAmount);
        gameV1.setTotalPlayCount(100);
        gameV1.setPlayerCount(player, 50);

        manager.collectAndDistribute(POSITION_ID);

        // First partial withdrawal
        manager.withdraw(player, address(weth));

        // Second collection
        weth.transfer(address(treasury), wethAmount);
        manager.collectAndDistribute(POSITION_ID);

        // Verify correct accounting after multiple operations
        uint256 withdrawable2 = manager.withdrawable(player, address(weth));
        assertEq(withdrawable2, 0); // Should be 0 since no V2 activity
    }

    function test_CollectAndDistribute_ArbitraryTokenPair() public {
        // Set up a new position with USDC/DAI pair
        uint256 newPositionId = 2;
        uint256 usdcAmount = 1000000 * 10 ** 6; // 1M USDC
        uint256 daiAmount = 1000000 * 10 ** 18; // 1M DAI

        positionManager.setPosition(
            newPositionId,
            address(usdc),
            address(dai),
            usdcAmount,
            daiAmount
        );
        positionManager.setOwner(newPositionId, address(treasury));

        // Transfer tokens to treasury
        usdc.transfer(address(treasury), usdcAmount);
        dai.transfer(address(treasury), daiAmount);

        // Set up V1 and V2 activity
        address player1 = address(1);
        address player2 = address(2);

        // V1: player1 = 30%, player2 = 70%
        gameV1.setTotalPlayCount(100);
        gameV1.setPlayerCount(player1, 30);
        gameV1.setPlayerCount(player2, 70);

        // Collect and distribute first batch (goes to V1)
        vm.expectEmit(true, true, true, true);
        emit InfiniteTreasuryManager.FeesCollected(
            newPositionId,
            address(usdc),
            address(dai),
            usdcAmount,
            daiAmount
        );
        manager.collectAndDistribute(newPositionId);

        // Verify V1 distribution
        assertEq(
            manager.withdrawable(player1, address(usdc)),
            (usdcAmount * 30) / 100
        );
        assertEq(
            manager.withdrawable(player1, address(dai)),
            (daiAmount * 30) / 100
        );
        assertEq(
            manager.withdrawable(player2, address(usdc)),
            (usdcAmount * 70) / 100
        );
        assertEq(
            manager.withdrawable(player2, address(dai)),
            (daiAmount * 70) / 100
        );

        // Second batch of fees
        usdc.transfer(address(treasury), usdcAmount);
        dai.transfer(address(treasury), daiAmount);

        // Set up V2 activity before second distribution
        // V2: player1 = 60%, player2 = 40%
        gameV2.setTotalFees(1000 ether);
        gameV2.setPlayerFees(player1, 600 ether);
        gameV2.setPlayerFees(player2, 400 ether);

        // Collect and distribute second batch (goes to V2)
        manager.collectAndDistribute(newPositionId);

        // Verify total withdrawable amounts
        assertEq(
            manager.withdrawable(player1, address(usdc)),
            ((usdcAmount * 30) / 100) + ((usdcAmount * 60) / 100)
        );
        assertEq(
            manager.withdrawable(player2, address(usdc)),
            ((usdcAmount * 70) / 100) + ((usdcAmount * 40) / 100)
        );
    }

    function test_EmergencyWithdrawERC20() public {
        // Transfer some tokens directly to the manager contract
        uint256 amount = 100 ether;
        address recipient = address(12345);

        // Only game owner can withdraw
        address nonOwner = address(123);
        vm.prank(nonOwner);
        vm.expectRevert(InfiniteTreasuryManager.NotGameOwner.selector);
        manager.emergencyWithdrawERC20(address(usdc), recipient);

        // Can't withdraw infinite token
        vm.prank(gameV2.owner());
        vm.expectRevert(
            InfiniteTreasuryManager.CannotWithdrawInfiniteToken.selector
        );
        manager.emergencyWithdrawERC20(address(infiniteToken), recipient);

        // Can't withdraw if no balance
        vm.prank(gameV2.owner());
        vm.expectRevert(InfiniteTreasuryManager.NoTokenBalance.selector);
        manager.emergencyWithdrawERC20(address(usdc), recipient);

        // Add balance
        usdc.transfer(address(manager), amount);

        // Successful withdrawal
        vm.prank(gameV2.owner());
        manager.emergencyWithdrawERC20(address(usdc), recipient);

        assertEq(usdc.balanceOf(address(manager)), 0);
        assertEq(usdc.balanceOf(recipient), amount);
    }

    function test_EmergencyWithdrawERC721() public {
        // Create a mock ERC721 token
        MockERC721 mockNFT = new MockERC721();
        uint256 tokenId = 123;
        address recipient = address(12345);

        // Mint NFT to manager contract
        mockNFT.mint(address(manager), tokenId);

        // Only game owner can withdraw
        address nonOwner = address(123);
        vm.prank(nonOwner);
        vm.expectRevert(InfiniteTreasuryManager.NotGameOwner.selector);
        manager.emergencyWithdrawERC721(address(mockNFT), tokenId, recipient);

        // Can't withdraw non-owned token
        uint256 nonOwnedTokenId = 456;
        vm.prank(Ownable(address(gameV2)).owner());
        vm.expectRevert(InfiniteTreasuryManager.NoTokenBalance.selector);
        manager.emergencyWithdrawERC721(
            address(mockNFT),
            nonOwnedTokenId,
            recipient
        );

        // Successful withdrawal
        vm.prank(Ownable(address(gameV2)).owner());
        manager.emergencyWithdrawERC721(address(mockNFT), tokenId, recipient);
        assertEq(mockNFT.ownerOf(tokenId), recipient);
    }
}

// SPDX-License-Identifier: MIT
pragma solidity ^0.8.20;

import "forge-std/Test.sol";
import "../src/InfiniteTreasuryManager.sol";
import "../src/InfiniteTreasury.sol";
import "../src/InfiniteGame.sol";
import "../src/InfiniteGame2.sol";
import "./mocks/MockNonfungiblePositionManager.sol";
import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import "@uniswap/v3-periphery/contracts/interfaces/INonfungiblePositionManager.sol";

contract MockWETH is ERC20 {
    constructor() ERC20("Wrapped Ether", "WETH") {
        _mint(msg.sender, 1e9 ether);
    }
}

contract MockInfiniteToken is ERC20 {
    constructor() ERC20("Infinite Token", "INF") {
        _mint(msg.sender, 1e9 ether);
    }
}

contract InfiniteTreasuryManagerIntegration is Test {
    InfiniteTreasuryManager public manager;
    InfiniteTreasury public treasury;
    InfiniteGame public gameV1;
    InfiniteGame2 public gameV2;
    MockWETH public weth;
    MockInfiniteToken public infiniteToken;
    MockNonfungiblePositionManager public positionManager;

    address public admin;
    address public player1;
    address public player2;

    uint256 public infiniteBalance = 5313 ether;
    uint256 public wethBalance = 459 ether;

    uint256 public constant POSITION_ID = 1;

    event FeesCollected(uint256 wethAmount, uint256 infiniteTokenAmount);

    function setUp() public {
        admin = address(this);
        player1 = address(0x1);
        player2 = address(0x2);

        // Deploy mock tokens
        weth = new MockWETH();
        infiniteToken = new MockInfiniteToken();
        positionManager = new MockNonfungiblePositionManager();

        // Deploy core contracts
        treasury = new InfiniteTreasury(address(positionManager));

        // Setup test position with accumulated fees
        positionManager.setPosition(
            POSITION_ID,
            address(infiniteToken),
            address(weth),
            infiniteBalance,
            wethBalance
        );

        infiniteToken.transfer(address(positionManager), infiniteBalance);
        weth.transfer(address(positionManager), wethBalance);

        positionManager.setOwner(POSITION_ID, address(treasury));

        gameV1 = new InfiniteGame(
            address(infiniteToken),
            block.timestamp,
            admin
        );

        gameV2 = new InfiniteGame2(
            address(infiniteToken),
            address(1),
            block.timestamp
        );

        // Deploy manager
        manager = new InfiniteTreasuryManager(
            address(treasury),
            address(positionManager),
            address(gameV1),
            address(gameV2),
            address(infiniteToken)
        );

        // Transfer treasury ownership to manager
        treasury.transferOwnership(address(manager));

        // Setup initial balances
        vm.startPrank(admin);
        infiniteToken.transfer(player1, 1_000_000 ether);
        infiniteToken.transfer(player2, 1_000_000 ether);
        weth.transfer(player1, 1_000_000 ether);
        weth.transfer(player2, 1_000_000 ether);
        vm.stopPrank();
    }

    // Helper function to simulate V1 gameplay
    function playGameV1(address player, uint256 times) internal {
        vm.startPrank(player);
        infiniteToken.approve(address(gameV1), type(uint256).max);

        for (uint256 i = 0; i < times; i++) {
            gameV1.play();
            vm.warp(block.timestamp + 1 seconds);
        }
        vm.stopPrank();
    }

    // Helper function to simulate V2 gameplay with fees
    function playGameV2(address player, uint256 times) internal {
        vm.startPrank(player);
        infiniteToken.approve(address(gameV2), type(uint256).max);
        weth.approve(address(gameV2), type(uint256).max);

        for (uint256 i = 0; i < times; i++) {
            gameV2.play();
            vm.warp(block.timestamp + 1 seconds);
        }
        vm.stopPrank();
    }

    function test_WithdrawalDistribution() public {
        // Setup plays in V1 and fees in V2
        playGameV1(player1, 6);
        playGameV1(player2, 4);
        playGameV2(player2, 4);

        uint256 player1BalanceBefore = weth.balanceOf(player1);
        uint256 player2BalanceBefore = weth.balanceOf(player2);

        vm.warp(block.timestamp + treasury.TIMELOCK_DURATION());

        manager.collectAndDistribute(POSITION_ID);

        // Verify treasury manager has funds
        uint256 donatedAmount = (wethBalance * treasury.donationShare()) / 100;
        uint256 distributedWeth = wethBalance - donatedAmount;

        assertEq(weth.balanceOf(address(manager)), distributedWeth);

        // First collection should go to V1 players
        uint256 player1Game1Share = manager.withdrawable(
            player1,
            address(weth)
        );
        uint256 player2Game1Share = manager.withdrawable(
            player2,
            address(weth)
        );

        assertEq(player1Game1Share + player2Game1Share, distributedWeth);

        manager.withdraw(player1, address(weth));

        assertEq(
            weth.balanceOf(player1),
            player1BalanceBefore + player1Game1Share
        );

        // 60% of remaining amount after donation
        assertEq(player1Game1Share, (distributedWeth * 60) / 100);
        // 40% of remaining amount after donation
        assertEq(player2Game1Share, (distributedWeth * 40) / 100);

        uint256 newWethAmount = 5 ether;

        // Add more WETH to the position
        positionManager.setPosition(
            POSITION_ID,
            address(infiniteToken),
            address(weth),
            0,
            newWethAmount
        );

        weth.transfer(address(positionManager), newWethAmount);

        // This should all go to the V2 game (only player 2)
        manager.collectAndDistribute(POSITION_ID);

        uint256 newWethDistributedAmount = newWethAmount -
            (newWethAmount * treasury.donationShare()) /
            100;

        // Verify V2 player can now withdraw
        uint256 player2Withdrawable = manager.withdrawable(
            player2,
            address(weth)
        );

        assertEq(
            player2Withdrawable,
            player2Game1Share + newWethDistributedAmount
        );

        manager.withdraw(player2, address(weth));

        assertEq(
            weth.balanceOf(player2),
            player2BalanceBefore + player2Withdrawable
        );
    }

    function test_RevertOnInsufficientBalance() public {
        vm.expectRevert(InfiniteTreasuryManager.InsufficientBalance.selector);
        manager.withdraw(player1, address(weth));
    }

    function test_RevertOnZeroAddress() public {
        vm.expectRevert(InfiniteTreasuryManager.ZeroAddress.selector);
        manager.withdraw(address(0), address(weth));
    }

    function test_WithdrawableCalculation() public {
        // Setup different play counts and fee amounts
        playGameV1(player1, 10); // 10 plays in V1
        playGameV1(player2, 20); // 20 plays in V1

        vm.warp(block.timestamp + treasury.TIMELOCK_DURATION());

        manager.collectAndDistribute(POSITION_ID);

        // Player2 should have ~2x the withdrawable amount of player1
        uint256 player1Share = manager.withdrawable(player1, address(weth));
        uint256 player2Share = manager.withdrawable(player2, address(weth));

        assertApproxEqRel(player2Share, player1Share * 2, 0.01e18); // Within 1% due to rounding
    }

    receive() external payable {}
}

Tests: InfiniteGame2 & InfiniteTreasuryManager

Tests: InfiniteGame2 & InfiniteTreasuryManager

The Infinite Game

Tests: InfiniteGame2 & InfiniteTreasuryManager

Tests: InfiniteGame2 & InfiniteTreasuryManager

Tests: InfiniteGame2 & InfiniteTreasuryManager

The Infinite Game

Tests: InfiniteGame2 & InfiniteTreasuryManager