// SPDX-License-Identifier: UNLICENSED
/* solhint-disable */
pragma solidity ^0.8.30;

import {ABDKMath64x64} from "../lib/abdk-libraries-solidity/ABDKMath64x64.sol";
import {Test} from "../lib/forge-std/src/Test.sol";
import {IERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/IERC20.sol";
import {ERC20} from "../lib/openzeppelin-contracts/contracts/token/ERC20/ERC20.sol";
import {Funding} from "../src/Funding.sol";
import {LMSRStabilized} from "../src/LMSRStabilized.sol";
import {PartyInfo} from "../src/PartyInfo.sol";
import {PartyPool} from "../src/PartyPool.sol";
import {Deploy} from "./Deploy.sol";

/// @notice Minimal ERC20 token for tests with an external mint function.
contract TestERC20 is ERC20 {
    constructor(string memory name_, string memory symbol_, uint256 initialSupply) ERC20(name_, symbol_) {
        if (initialSupply > 0) {
            _mint(msg.sender, initialSupply);
        }
    }

    function mint(address to, uint256 amount) external {
        _mint(to, amount);
    }

    function approveMax(address spender) external {
        _approve(msg.sender, spender, type(uint256).max);
    }
}

/// @notice Test contract that provides token funds when called by PartyPool.swap via a selector.
/// The pool will call the payer with (token, amount) and expects the payer to transfer the input token
/// into the pool. This contract implements that provider function.
contract SwapCallbackContract {
    address public pool;
    address public tokenSource;
    bool public shouldFail;

    constructor(address _pool) {
        pool = _pool;
    }

    function setTokenSource(address _tokenSource) external {
        tokenSource = _tokenSource;
    }

    function setShouldFail(bool _shouldFail) external {
        shouldFail = _shouldFail;
    }

    /// @notice Called by PartyPool.swap on the payer. Signature must be:
    ///         provideFunding(address token, uint256 amount)
    /// @dev The pool will call this function to request the input token; this function
    ///      pulls funds from tokenSource (via ERC20.transferFrom) into the pool.
    function provideFunding(address token, uint256 amount) external {
        require(msg.sender == pool, "Callback not called by pool");
        if (shouldFail) revert("callback failed");
        require(tokenSource != address(0), "no token source");

        // Pull the required tokens from tokenSource into the pool
        TestERC20(token).transferFrom(tokenSource, pool, amount);
    }
}

/// @notice Tests for PartyPool swap functionality using alternative funding mechanisms:
/// pre-funding and callback method. Validates that input/output amounts match swap results
/// view calls and that pool balances are correct.
contract FundingTest is Test {
    using ABDKMath64x64 for int128;

    TestERC20 token0;
    TestERC20 token1;
    TestERC20 token2;
    PartyPool pool;
    PartyPool poolZeroFee;
    PartyInfo info;
    SwapCallbackContract callbackContract;

    address alice;
    address bob;

    // Common parameters
    int128 tradeFrac;
    int128 targetSlippage;

    uint256 constant INIT_BAL = 1_000_000; // initial token units for each token

    // Callback funding selector - the pool will call payer.provideFunding(address token, uint256 amount)
    bytes4 constant CALLBACK = SwapCallbackContract.provideFunding.selector;

    function setUp() public {
        alice = address(0xA11ce);
        bob = address(0xB0b);

        // Deploy three ERC20 test tokens
        token0 = new TestERC20("T0", "T0", 0);
        token1 = new TestERC20("T1", "T1", 0);
        token2 = new TestERC20("T2", "T2", 0);

        // Mint initial balances to the test contract
        token0.mint(address(this), INIT_BAL);
        token1.mint(address(this), INIT_BAL);
        token2.mint(address(this), INIT_BAL);

        // Configure LMSR parameters
        tradeFrac = ABDKMath64x64.divu(100, 10_000); // 0.01
        targetSlippage = ABDKMath64x64.divu(10, 10_000); // 0.001

        // Build arrays for pool constructor
        IERC20[] memory tokens = new IERC20[](3);
        tokens[0] = IERC20(address(token0));
        tokens[1] = IERC20(address(token1));
        tokens[2] = IERC20(address(token2));

        // Deploy pool with a small fee (0.1%)
        uint256 feePpm = 1000;
        int128 kappa = LMSRStabilized.computeKappaFromSlippage(tokens.length, tradeFrac, targetSlippage);
        pool = Deploy.newPartyPool(address(this), "LP", "LP", tokens, kappa, feePpm, feePpm, false);

        // Transfer initial deposit amounts into pool
        token0.transfer(address(pool), INIT_BAL);
        token1.transfer(address(pool), INIT_BAL);
        token2.transfer(address(pool), INIT_BAL);

        // Perform initial mint
        pool.initialMint(address(this), INIT_BAL * tokens.length * 10**18);

        // Deploy pool with zero fees for exact balance matching
        poolZeroFee = Deploy.newPartyPool(address(this), "LP_ZERO", "LP_ZERO", tokens, kappa, 0, 0, false);

        // Mint additional tokens for zero-fee pool
        token0.mint(address(this), INIT_BAL);
        token1.mint(address(this), INIT_BAL);
        token2.mint(address(this), INIT_BAL);

        // Transfer to zero-fee pool
        token0.transfer(address(poolZeroFee), INIT_BAL);
        token1.transfer(address(poolZeroFee), INIT_BAL);
        token2.transfer(address(poolZeroFee), INIT_BAL);

        // Initialize zero-fee pool
        poolZeroFee.initialMint(address(this), INIT_BAL * tokens.length * 10**18);

        // Mint tokens to alice and bob for testing
        token0.mint(alice, INIT_BAL);
        token1.mint(alice, INIT_BAL);
        token2.mint(alice, INIT_BAL);

        token0.mint(bob, INIT_BAL);
        token1.mint(bob, INIT_BAL);
        token2.mint(bob, INIT_BAL);

        // Deploy callback contract
        callbackContract = new SwapCallbackContract(address(pool));

        info = Deploy.newInfo();
    }

    /* ----------------------
       Pre-funding Tests
       ---------------------- */

    /// @notice Test swap using pre-funding mechanism with regular fee pool
    function testSwapWithPreFunding() public {
        uint256 maxIn = 10_000;

        // Pre-fund the pool by transferring tokens before the swap
        vm.startPrank(alice);
        token0.transfer(address(pool), maxIn);

        uint256 poolToken0Before = token0.balanceOf(address(pool));
        uint256 poolToken1Before = token1.balanceOf(address(pool));
        uint256 bobToken1Before = token1.balanceOf(bob);

        // Execute swap using Funding.USE_PREFUNDING for pre-funded: token0 -> token1
        (uint256 amountIn, uint256 amountOut, uint256 fee) = pool.swap(
            alice,    // payer (not used with pre-funded)
            Funding.USE_PREFUNDING,
            bob,      // receiver
            0,        // inputTokenIndex (token0)
            1,        // outputTokenIndex (token1)
            maxIn,    // maxAmountIn
            0,        // limitPrice
            0,        // deadline
            false     // unwrap
        );

        // Verify amounts
        assertTrue(amountIn > 0, "expected some input used");
        assertTrue(amountOut > 0, "expected some output returned");
        assertTrue(amountIn <= maxIn, "used input must not exceed max");
        assertTrue(fee <= amountIn, "fee must not exceed total input");

        // Bob received the output
        assertEq(token1.balanceOf(bob), bobToken1Before + amountOut, "Bob should receive output");

        // Pool balances changed as expected
        // Input token increased (tokens were already transferred)
        assertEq(token0.balanceOf(address(pool)), poolToken0Before, "Pool token0 should remain at pre-funded level");
        
        // Output token decreased
        assertEq(token1.balanceOf(address(pool)), poolToken1Before - amountOut, "Pool token1 should decrease");

        // If any unused tokens, they remain in the pool
        uint256 unusedTokens = maxIn - amountIn;
        if (unusedTokens > 0) {
            // The pre-funded amount stays in the pool
            assertTrue(token0.balanceOf(address(pool)) >= poolToken0Before - amountIn, "Unused tokens remain");
        }

        vm.stopPrank();
    }

    /// @notice Test swap using pre-funding with zero-fee pool to verify exact pool balance matching
    function testSwapWithPreFundingZeroFeeExactBalances() public {
        uint256 maxIn = 10_000;

        // Pre-fund the pool
        vm.startPrank(alice);
        token0.transfer(address(poolZeroFee), maxIn);

        uint256 poolToken0Before = token0.balanceOf(address(poolZeroFee));
        uint256 poolToken1Before = token1.balanceOf(address(poolZeroFee));

        // Execute swap
        (uint256 amountIn, uint256 amountOut, uint256 fee) = poolZeroFee.swap(
            alice,
            Funding.USE_PREFUNDING,
            bob,
            0,        // token0 -> token1
            1,
            maxIn,
            0,
            0,
            false
        );

        // With zero fees, fee should be 0
        assertEq(fee, 0, "Fee should be zero in zero-fee pool");

        // Pool balances should match exactly (no rounding errors with zero fees)
        uint256 poolToken0After = token0.balanceOf(address(poolZeroFee));
        uint256 poolToken1After = token1.balanceOf(address(poolZeroFee));

        // Net change: input increased by amountIn (already pre-funded), output decreased by amountOut
        assertEq(poolToken0After, poolToken0Before, "Pool token0 balance exact");
        assertEq(poolToken1After, poolToken1Before - amountOut, "Pool token1 balance exact");

        vm.stopPrank();
    }

    /// @notice Test that pre-funding with insufficient tokens reverts appropriately
    function testSwapWithPreFundingInsufficientTokensReverts() public {
        uint256 maxIn = 10_000;
        uint256 insufficientAmount = maxIn / 2; // Only half of what's needed

        vm.startPrank(alice);
        token0.transfer(address(pool), insufficientAmount);

        // This should revert because the pool doesn't have enough pre-funded tokens
        vm.expectRevert();
        pool.swap(
            alice,
            Funding.USE_PREFUNDING,
            bob,
            0,
            1,
            maxIn,
            0,
            0,
            false
        );

        vm.stopPrank();
    }

    /* ----------------------
       Callback Method Tests
       ---------------------- */

    /// @notice Test swap using callback mechanism
    function testSwapWithCallback() public {
        uint256 maxIn = 10_000;

        // Setup callback contract to use alice's tokens
        callbackContract.setTokenSource(alice);
        callbackContract.setShouldFail(false);

        vm.startPrank(alice);
        // Alice approves callback contract to transfer tokens
        token0.approve(address(callbackContract), type(uint256).max);

        uint256 aliceToken0Before = token0.balanceOf(alice);
        uint256 poolToken0Before = token0.balanceOf(address(pool));
        uint256 poolToken1Before = token1.balanceOf(address(pool));
        uint256 bobToken1Before = token1.balanceOf(bob);

        // Execute swap using callback: token0 -> token1
        // The payer address (callbackContract) will receive the callback
        vm.stopPrank();
        
        (uint256 amountIn, uint256 amountOut, uint256 fee) = pool.swap(
            address(callbackContract),  // payer (receives callback)
            CALLBACK,
            bob,                        // receiver
            0,                          // inputTokenIndex (token0)
            1,                          // outputTokenIndex (token1)
            maxIn,                      // maxAmountIn
            0,                          // limitPrice
            0,                          // deadline
            false                       // unwrap
        );

        // Verify amounts
        assertTrue(amountIn > 0, "expected some input used");
        assertTrue(amountOut > 0, "expected some output returned");
        assertTrue(amountIn <= maxIn, "used input must not exceed max");
        assertTrue(fee <= amountIn, "fee must not exceed total input");

        // Alice's tokens were used (via callback)
        assertEq(token0.balanceOf(alice), aliceToken0Before - amountIn, "Alice tokens should decrease");

        // Bob received the output
        assertEq(token1.balanceOf(bob), bobToken1Before + amountOut, "Bob should receive output");

        // Pool balances changed as expected
        assertEq(token0.balanceOf(address(pool)), poolToken0Before + amountIn, "Pool token0 should increase");
        assertEq(token1.balanceOf(address(pool)), poolToken1Before - amountOut, "Pool token1 should decrease");
    }

    /// @notice Test swap callback with zero-fee pool for exact balance matching
    function testSwapWithCallbackZeroFeeExactBalances() public {
        uint256 maxIn = 10_000;

        // Setup callback for zero-fee pool
        SwapCallbackContract zeroFeeCallback = new SwapCallbackContract(address(poolZeroFee));
        zeroFeeCallback.setTokenSource(alice);
        zeroFeeCallback.setShouldFail(false);

        vm.startPrank(alice);
        token0.approve(address(zeroFeeCallback), type(uint256).max);
        vm.stopPrank();

        uint256 poolToken0Before = token0.balanceOf(address(poolZeroFee));
        uint256 poolToken1Before = token1.balanceOf(address(poolZeroFee));

        // Execute swap
        (uint256 amountIn, uint256 amountOut, uint256 fee) = poolZeroFee.swap(
            address(zeroFeeCallback),
            CALLBACK,
            bob,
            0,
            1,
            maxIn,
            0,
            0,
            false
        );

        // With zero fees, fee should be 0
        assertEq(fee, 0, "Fee should be zero in zero-fee pool");

        // Pool balances should match exactly
        uint256 poolToken0After = token0.balanceOf(address(poolZeroFee));
        uint256 poolToken1After = token1.balanceOf(address(poolZeroFee));

        assertEq(poolToken0After, poolToken0Before + amountIn, "Pool token0 balance exact");
        assertEq(poolToken1After, poolToken1Before - amountOut, "Pool token1 balance exact");
    }

    /// @notice Test that callback failure causes swap to revert
    function testSwapWithCallbackFailureReverts() public {
        uint256 maxIn = 10_000;

        callbackContract.setTokenSource(alice);
        callbackContract.setShouldFail(true); // Make callback fail

        vm.expectRevert();
        pool.swap(
            address(callbackContract),
            CALLBACK,
            bob,
            0,
            1,
            maxIn,
            0,
            0,
            false
        );
    }

    /* ----------------------
       Validation Against swapAmounts()
       ---------------------- */

    /// @notice Verify that pre-funded swap amounts match swapAmounts() view predictions
    function testPreFundingMatchesSwapAmountsView() public {
        uint256 maxIn = 10_000;

        // Perform a reference swap with USE_APPROVALS to get expected amounts
        vm.startPrank(bob);
        token0.approve(address(pool), type(uint256).max);
        
        (uint256 refAmountIn, uint256 refAmountOut, uint256 refFee) = pool.swap(
            bob,
            Funding.USE_APPROVALS,
            bob,
            0,
            1,
            maxIn,
            0,
            0,
            false
        );
        vm.stopPrank();

        // Reset pool state by creating a fresh pool with identical parameters
        IERC20[] memory tokens = new IERC20[](3);
        tokens[0] = IERC20(address(token0));
        tokens[1] = IERC20(address(token1));
        tokens[2] = IERC20(address(token2));
        
        uint256 feePpm = 1000;
        int128 kappa = LMSRStabilized.computeKappaFromSlippage(tokens.length, tradeFrac, targetSlippage);
        PartyPool testPool = Deploy.newPartyPool(address(this), "LP_TEST", "LP_TEST", tokens, kappa, feePpm, feePpm, false);
        
        token0.mint(address(this), INIT_BAL);
        token1.mint(address(this), INIT_BAL);
        token2.mint(address(this), INIT_BAL);
        
        token0.transfer(address(testPool), INIT_BAL);
        token1.transfer(address(testPool), INIT_BAL);
        token2.transfer(address(testPool), INIT_BAL);
        
        testPool.initialMint(address(this), INIT_BAL * tokens.length * 10**18);

        // Now test pre-funding with same initial state
        vm.startPrank(alice);
        token0.transfer(address(testPool), maxIn);
        
        (uint256 preAmountIn, uint256 preAmountOut, uint256 preFee) = testPool.swap(
            alice,
            Funding.USE_PREFUNDING,
            alice,
            0,
            1,
            maxIn,
            0,
            0,
            false
        );
        vm.stopPrank();

        // Pre-funded amounts should match reference swap amounts
        assertEq(preAmountIn, refAmountIn, "Pre-funded amountIn should match reference");
        assertEq(preAmountOut, refAmountOut, "Pre-funded amountOut should match reference");
        assertEq(preFee, refFee, "Pre-funded fee should match reference");
    }

    /// @notice Verify that callback swap amounts match swapAmounts() view predictions
    function testCallbackMatchesSwapAmountsView() public {
        uint256 maxIn = 10_000;

        // Perform a reference swap
        vm.startPrank(bob);
        token0.approve(address(pool), type(uint256).max);
        
        (uint256 refAmountIn, uint256 refAmountOut, uint256 refFee) = pool.swap(
            bob,
            Funding.USE_APPROVALS,
            bob,
            0,
            1,
            maxIn,
            0,
            0,
            false
        );
        vm.stopPrank();

        // Create fresh pool for callback test
        IERC20[] memory tokens = new IERC20[](3);
        tokens[0] = IERC20(address(token0));
        tokens[1] = IERC20(address(token1));
        tokens[2] = IERC20(address(token2));
        
        uint256 feePpm = 1000;
        int128 kappa = LMSRStabilized.computeKappaFromSlippage(tokens.length, tradeFrac, targetSlippage);
        PartyPool testPool = Deploy.newPartyPool(address(this), "LP_TEST2", "LP_TEST2", tokens, kappa, feePpm, feePpm, false);
        
        token0.mint(address(this), INIT_BAL);
        token1.mint(address(this), INIT_BAL);
        token2.mint(address(this), INIT_BAL);
        
        token0.transfer(address(testPool), INIT_BAL);
        token1.transfer(address(testPool), INIT_BAL);
        token2.transfer(address(testPool), INIT_BAL);
        
        testPool.initialMint(address(this), INIT_BAL * tokens.length * 10**18);

        // Setup callback for test pool
        SwapCallbackContract testCallback = new SwapCallbackContract(address(testPool));
        testCallback.setTokenSource(alice);
        testCallback.setShouldFail(false);

        vm.startPrank(alice);
        token0.approve(address(testCallback), type(uint256).max);
        vm.stopPrank();

        // Test callback with same initial state
        (uint256 cbAmountIn, uint256 cbAmountOut, uint256 cbFee) = testPool.swap(
            address(testCallback),
            CALLBACK,
            alice,
            0,
            1,
            maxIn,
            0,
            0,
            false
        );

        // Callback amounts should match reference swap amounts
        assertEq(cbAmountIn, refAmountIn, "Callback amountIn should match reference");
        assertEq(cbAmountOut, refAmountOut, "Callback amountOut should match reference");
        assertEq(cbFee, refFee, "Callback fee should match reference");
    }

    /// @notice Test multiple swaps in sequence with different funding methods produce consistent results
    function testMultipleSwapsFundingMethodsConsistency() public {
        uint256[] memory swapAmounts = new uint256[](3);
        swapAmounts[0] = 5_000;
        swapAmounts[1] = 7_500;
        swapAmounts[2] = 10_000;

        for (uint i = 0; i < swapAmounts.length; i++) {
            uint256 swapAmount = swapAmounts[i];

            // Create three identical pools
            IERC20[] memory tokens = new IERC20[](3);
            tokens[0] = IERC20(address(token0));
            tokens[1] = IERC20(address(token1));
            tokens[2] = IERC20(address(token2));
            
            int128 kappa = LMSRStabilized.computeKappaFromSlippage(tokens.length, tradeFrac, targetSlippage);
            
            PartyPool poolApproval = Deploy.newPartyPool(address(this), "LP_A", "LP_A", tokens, kappa, 0, 0, false);
            PartyPool poolPreFund = Deploy.newPartyPool(address(this), "LP_P", "LP_P", tokens, kappa, 0, 0, false);
            PartyPool poolCallback = Deploy.newPartyPool(address(this), "LP_C", "LP_C", tokens, kappa, 0, 0, false);

            // Initialize all three pools identically
            token0.mint(address(this), INIT_BAL * 3);
            token1.mint(address(this), INIT_BAL * 3);
            token2.mint(address(this), INIT_BAL * 3);

            token0.transfer(address(poolApproval), INIT_BAL);
            token1.transfer(address(poolApproval), INIT_BAL);
            token2.transfer(address(poolApproval), INIT_BAL);
            poolApproval.initialMint(address(this), 0);

            token0.transfer(address(poolPreFund), INIT_BAL);
            token1.transfer(address(poolPreFund), INIT_BAL);
            token2.transfer(address(poolPreFund), INIT_BAL);
            poolPreFund.initialMint(address(this), 0);

            token0.transfer(address(poolCallback), INIT_BAL);
            token1.transfer(address(poolCallback), INIT_BAL);
            token2.transfer(address(poolCallback), INIT_BAL);
            poolCallback.initialMint(address(this), 0);

            // Test with USE_APPROVALS (bytes4(0) with approvals)
            vm.startPrank(alice);
            token0.approve(address(poolApproval), type(uint256).max);
            (uint256 apprIn, uint256 apprOut, ) = poolApproval.swap(
                alice, Funding.USE_APPROVALS, alice, 0, 1, swapAmount, 0, 0, false
            );
            vm.stopPrank();

            // Test with PRE_FUNDED (bytes4(0) with pre-funding)
            vm.startPrank(alice);
            token0.transfer(address(poolPreFund), swapAmount);
            (uint256 preIn, uint256 preOut, ) = poolPreFund.swap(
                alice, Funding.USE_PREFUNDING, alice, 0, 1, swapAmount, 0, 0, false
            );
            vm.stopPrank();

            // Test with CALLBACK
            SwapCallbackContract cb = new SwapCallbackContract(address(poolCallback));
            cb.setTokenSource(alice);
            cb.setShouldFail(false);
            
            vm.startPrank(alice);
            token0.approve(address(cb), type(uint256).max);
            vm.stopPrank();
            
            (uint256 cbIn, uint256 cbOut, ) = poolCallback.swap(
                address(cb), CALLBACK, alice, 0, 1, swapAmount, 0, 0, false
            );

            // All three methods should produce identical results
            assertEq(preIn, apprIn, "Pre-funded input should match approval");
            assertEq(preOut, apprOut, "Pre-funded output should match approval");
            assertEq(cbIn, apprIn, "Callback input should match approval");
            assertEq(cbOut, apprOut, "Callback output should match approval");
        }
    }
}
/* solhint-enable */