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

import "forge-std/console2.sol";
import "@abdk/ABDKMath64x64.sol";
import "@openzeppelin/contracts/token/ERC20/ERC20.sol";
import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
import "@openzeppelin/contracts/token/ERC20/utils/SafeERC20.sol";
import "@openzeppelin/contracts/utils/ReentrancyGuard.sol";
import "./LMSRStabilized.sol";
import "./IPartyPool.sol";
import "./IPartyFlashCallback.sol";

/// @title PartyPool - LMSR-backed multi-asset pool with LP ERC20 token
/// @notice Uses LMSRStabilized library; stores per-token uint bases to convert to/from 64.64 fixed point.
/// - Caches qInternal[] (int128 64.64) and cachedUintBalances[] to minimize balanceOf() calls.
/// - swap and swapToLimit mimic core lib; mint/burn call updateForProportionalChange() and manage LP tokens.
contract PartyPool is IPartyPool, ERC20, ReentrancyGuard {
    using ABDKMath64x64 for int128;
    using LMSRStabilized for LMSRStabilized.State;
    using SafeERC20 for IERC20;


    //
    // Immutable pool configuration
    //

    address[] public tokens; // effectively immutable since there is no interface to change the tokens
    function numTokens() external view returns (uint256) { return tokens.length; }
    function allTokens() external view returns (address[] memory) { return tokens; }

    // NOTE that the slippage target is only exactly achieved in completely balanced pools where all assets are
    // priced the same. This target is actually a minimum slippage that the pool imposes on traders, and the actual
    // slippage cost can be multiples bigger in practice due to pool inventory imbalances.
    int128 public immutable tradeFrac; // slippage target trade size as a fraction of one asset's inventory
    int128 public immutable targetSlippage; // target slippage applied to that trade size

    // fee in parts-per-million (ppm), taken from inputs before swaps
    uint256 public immutable swapFeePpm;

    // flash loan fee in parts-per-million (ppm)
    uint256 public immutable flashFeePpm;

    //
    // Internal state
    //
    
    LMSRStabilized.State internal lmsr;

    // Cached on-chain balances (uint) and internal 64.64 representation
    // balance / base = internal
    uint256[] internal cachedUintBalances;
    uint256[] internal bases; // per-token uint base used to scale token amounts <-> internal
    
    mapping(address=>uint) public tokenAddressToIndexPlusOne; // Uses index+1 so a result of 0 indicates a failed lookup
    
    uint256 public constant LP_SCALE = 1e18; // Scale used to convert LMSR lastTotal (Q64.64) into LP token units (uint)

    /// @param name_ LP token name
    /// @param symbol_ LP token symbol
    /// @param _tokens token addresses (n)
    /// @param _bases scaling bases for each token (n) - used when converting to/from internal 64.64 amounts
    /// @param _tradeFrac trade fraction in 64.64 fixed-point (as used by LMSR)
    /// @param _targetSlippage target slippage in 64.64 fixed-point (as used by LMSR)
    /// @param _swapFeePpm fee in parts-per-million, taken from swap input amounts before LMSR calculations
    /// @param _flashFeePpm fee in parts-per-million, taken for flash loans
    constructor(
        string memory name_,
        string memory symbol_,
        address[] memory _tokens,
        uint256[] memory _bases,
        int128 _tradeFrac,
        int128 _targetSlippage,
        uint256 _swapFeePpm,
        uint256 _flashFeePpm
    ) ERC20(name_, symbol_) {
        require(_tokens.length > 1, "Pool: need >1 asset");
        require(_tokens.length == _bases.length, "Pool: lengths mismatch");
        tokens = _tokens;
        bases = _bases;
        tradeFrac = _tradeFrac;
        targetSlippage = _targetSlippage;
        require(_swapFeePpm < 1_000_000, "Pool: fee >= ppm");
        swapFeePpm = _swapFeePpm;
        require(_flashFeePpm < 1_000_000, "Pool: flash fee >= ppm");
        flashFeePpm = _flashFeePpm;

        uint256 n = _tokens.length;

        // Initialize LMSR state nAssets; full init occurs on first mint when quantities are known.
        lmsr.nAssets = n;

        // Initialize token address to index mapping
        for (uint i = 0; i < n;) {
            tokenAddressToIndexPlusOne[_tokens[i]] = i + 1;
            unchecked {i++;}
        }

        // Initialize caches to zero
        cachedUintBalances = new uint256[](n);
    }


    /* ----------------------
       Initialization / Mint / Burn (LP token managed)
       ---------------------- */

    /// @notice Calculate the proportional deposit amounts required for a given LP token amount
    /// @param lpTokenAmount The amount of LP tokens desired
    /// @return depositAmounts Array of token amounts to deposit (rounded up)
    function mintDepositAmounts(uint256 lpTokenAmount) public view returns (uint256[] memory depositAmounts) {
        uint256 n = tokens.length;
        depositAmounts = new uint256[](n);

        // If this is the first mint or pool is empty, return zeros
        // For first mint, tokens should already be transferred to the pool
        if (totalSupply() == 0 || lmsr.nAssets == 0) {
            return depositAmounts; // Return zeros, initial deposit handled differently
        }

        // Calculate deposit based on current proportions
        uint256 totalLpSupply = totalSupply();

        // lpTokenAmount / totalLpSupply = depositAmount / currentBalance
        // Therefore: depositAmount = (lpTokenAmount * currentBalance) / totalLpSupply
        // We round up to protect the pool
        for (uint i = 0; i < n; i++) {
            uint256 currentBalance = cachedUintBalances[i];
            // Calculate with rounding up: (a * b + c - 1) / c
            depositAmounts[i] = (lpTokenAmount * currentBalance + totalLpSupply - 1) / totalLpSupply;
        }

        return depositAmounts;
    }

    /// @notice Calculate the proportional withdrawal amounts for a given LP token amount
    /// @param lpTokenAmount The amount of LP tokens to burn
    /// @return withdrawAmounts Array of token amounts to withdraw (rounded down)
    function burnReceiveAmounts(uint256 lpTokenAmount) external view returns (uint256[] memory withdrawAmounts) {
        return _burnReceiveAmounts(lpTokenAmount);
    }

    function _burnReceiveAmounts(uint256 lpTokenAmount) internal view returns (uint256[] memory withdrawAmounts) {
        uint256 n = tokens.length;
        withdrawAmounts = new uint256[](n);

        // If supply is zero or pool uninitialized, return zeros
        if (totalSupply() == 0 || lmsr.nAssets == 0) {
            return withdrawAmounts; // Return zeros, nothing to withdraw
        }

        // Calculate withdrawal amounts based on current proportions
        uint256 totalLpSupply = totalSupply();

        // withdrawAmount = floor(lpTokenAmount * currentBalance / totalLpSupply)
        for (uint i = 0; i < n; i++) {
            uint256 currentBalance = cachedUintBalances[i];
            withdrawAmounts[i] = (lpTokenAmount * currentBalance) / totalLpSupply;
        }

        return withdrawAmounts;
    }

    /// @notice Proportional mint (or initial supply if first call).
    /// For initial supply: assumes tokens have already been transferred to the pool
    /// For subsequent mints: payer must approve tokens beforehand, receiver gets the LP tokens
    /// @param payer address that provides the input tokens (ignored for initial deposit)
    /// @param receiver address that receives the LP tokens
    /// @param lpTokenAmount desired amount of LP tokens to mint (ignored for initial deposit)
    /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore.
    function mint(address payer, address receiver, uint256 lpTokenAmount, uint256 deadline) external nonReentrant {
        require(deadline == 0 || block.timestamp <= deadline, "mint: deadline exceeded");
        uint256 n = tokens.length;
        // Check if this is initial deposit
        bool isInitialDeposit = totalSupply() == 0 || lmsr.nAssets == 0;

        require(lpTokenAmount > 0 || isInitialDeposit, "mint: zero LP amount");

        // Capture old pool size metric (scaled) by computing from current balances
        uint256 oldScaled = 0;
        if (!isInitialDeposit) {
            int128 oldTotal = _computeSizeMetric(lmsr.qInternal);
            oldScaled = ABDKMath64x64.mulu(oldTotal, LP_SCALE);
        }

        // For non-initial deposits, transfer tokens from payer
        uint256[] memory depositAmounts = new uint256[](n);

        if (!isInitialDeposit) {
            // Calculate required deposit amounts for the desired LP tokens
            depositAmounts = mintDepositAmounts(lpTokenAmount);

            // Transfer in all token amounts
            for (uint i = 0; i < n; ) {
                if (depositAmounts[i] > 0) {
                    _safeTransferFrom(tokens[i], payer, address(this), depositAmounts[i]);
                }
                unchecked { i++; }
            }
        }

        // Update cached balances for all assets
        int128[] memory newQInternal = new int128[](n);
        for (uint i = 0; i < n; ) {
            uint256 bal = IERC20(tokens[i]).balanceOf(address(this));
            cachedUintBalances[i] = bal;
            newQInternal[i] = _uintToInternalFloor(bal, bases[i]);

            // For initial deposit, record the actual deposited amounts
            if (isInitialDeposit) {
                depositAmounts[i] = bal;
            }

            unchecked { i++; }
        }

        // If first time, call init, otherwise update proportional change.
        if (isInitialDeposit) {
            // Initialize the stabilized LMSR state
            lmsr.init(newQInternal, tradeFrac, targetSlippage);
        } else {
            // Update for proportional change
            lmsr.updateForProportionalChange(newQInternal);
        }

        // Compute actual LP tokens to mint based on change in size metric (scaled)
        // floor truncation rounds in favor of the pool
        int128 newTotal = _computeSizeMetric(newQInternal);
        uint256 newScaled = ABDKMath64x64.mulu(newTotal, LP_SCALE);
        uint256 actualLpToMint;

        if (isInitialDeposit) {
            // Initial provisioning: mint newScaled (as LP units)
            actualLpToMint = newScaled;
        } else {
            require(oldScaled > 0, "mint: oldScaled zero");
            uint256 delta = (newScaled > oldScaled) ? (newScaled - oldScaled) : 0;
            // Proportional issuance: totalSupply * delta / oldScaled
            if (delta > 0) {
                // floor truncation rounds in favor of the pool
                actualLpToMint = (totalSupply() * delta) / oldScaled;
            } else {
                actualLpToMint = 0;
            }
        }

        // For subsequent mints, ensure the calculated LP amount is not too different from requested
        if (!isInitialDeposit) {
            // Allow for some rounding error but ensure we're not far off from requested amount
            require(actualLpToMint > 0, "mint: zero LP minted");

            // Allow actual amount to be at most 0.00001% less than requested
            // This accounts for rounding in deposit calculations
            uint256 minAcceptable = lpTokenAmount * 99_999 / 100_000;
            require(actualLpToMint >= minAcceptable, "mint: insufficient LP minted");
        }

        require( actualLpToMint > 0, "mint: zero LP amount");
        _mint(receiver, actualLpToMint);
        emit Mint(payer, receiver, depositAmounts, actualLpToMint);
    }

    /// @notice Burn LP tokens and withdraw the proportional basket to receiver.
    /// Payer must own the LP tokens; withdraw amounts are computed from current proportions.
    /// @param payer address that provides the LP tokens to burn
    /// @param receiver address that receives the withdrawn tokens
    /// @param lpAmount amount of LP tokens to burn (proportional withdrawal)
    /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore.
    function burn(address payer, address receiver, uint256 lpAmount, uint256 deadline) external nonReentrant {
        require(deadline == 0 || block.timestamp <= deadline, "burn: deadline exceeded");
        uint256 n = tokens.length;
        require(lpAmount > 0, "burn: zero lp");

        uint256 supply = totalSupply();
        require(supply > 0, "burn: empty supply");
        require(lmsr.nAssets > 0, "burn: uninit pool");
        require(balanceOf(payer) >= lpAmount, "burn: insufficient LP");

        // Refresh cached balances to reflect current on-chain balances before computing withdrawal amounts
        for (uint i = 0; i < n; ) {
            uint256 bal = IERC20(tokens[i]).balanceOf(address(this));
            cachedUintBalances[i] = bal;
            unchecked { i++; }
        }

        // Compute proportional withdrawal amounts for the requested LP amount (rounded down)
        uint256[] memory withdrawAmounts = _burnReceiveAmounts(lpAmount);

        // Transfer underlying tokens out to receiver according to computed proportions
        for (uint i = 0; i < n; ) {
            if (withdrawAmounts[i] > 0) {
                _safeTransfer(tokens[i], receiver, withdrawAmounts[i]);
            }
            unchecked { i++; }
        }

        // Update cached balances and internal q for all assets
        int128[] memory newQInternal = new int128[](n);
        for (uint i = 0; i < n; ) {
            uint256 bal = IERC20(tokens[i]).balanceOf(address(this));
            cachedUintBalances[i] = bal;
            newQInternal[i] = _uintToInternalFloor(bal, bases[i]);
            unchecked { i++; }
        }

        // Apply proportional update or deinitialize if drained
        bool allZero = true;
        for (uint i = 0; i < n; ) {
            if (newQInternal[i] != int128(0)) {
                allZero = false;
                break;
            }
            unchecked { i++; }
        }

        if (allZero) {
            lmsr.deinit();
        } else {
            lmsr.updateForProportionalChange(newQInternal);
        }

        // Burn exactly the requested LP amount from payer (authorization via allowance)
        if (msg.sender != payer) {
            uint256 allowed = allowance(payer, msg.sender);
            require(allowed >= lpAmount, "burn: allowance insufficient");
            _approve(payer, msg.sender, allowed - lpAmount);
        }
        _burn(payer, lpAmount);

        emit Burn(payer, receiver, withdrawAmounts, lpAmount);
    }

    /* ----------------------
       Swaps
       ---------------------- */

    /// @notice Internal quote for exact-input swap that mirrors swap() rounding and fee application
    /// @return grossIn amount to transfer in (inclusive of fee), amountOutUint output amount (uint),
    ///         amountInInternalUsed and amountOutInternal (64.64), amountInUintNoFee input amount excluding fee (uint),
    ///         feeUint fee taken from the gross input (uint)
    function _quoteSwapExactIn(
        uint256 i,
        uint256 j,
        uint256 maxAmountIn,
        int128 limitPrice
    )
        internal
        view
        returns (
            uint256 grossIn,
            uint256 amountOutUint,
            int128 amountInInternalUsed,
            int128 amountOutInternal,
            uint256 amountInUintNoFee,
            uint256 feeUint
        )
    {
        uint256 n = tokens.length;
        require(i < n && j < n, "swap: idx");
        require(maxAmountIn > 0, "swap: input zero");
        require(lmsr.nAssets > 0, "swap: empty pool");

        // Estimate max net input (fee on gross rounded up, then subtract)
        (, uint256 netUintForSwap) = _computeFee(maxAmountIn);

        // Convert to internal (floor)
        int128 deltaInternalI = _uintToInternalFloor(netUintForSwap, bases[i]);
        require(deltaInternalI > int128(0), "swap: input too small after fee");

        // Compute internal amounts using LMSR (exact-input with price limit)
        (amountInInternalUsed, amountOutInternal) = lmsr.swapAmountsForExactInput(i, j, deltaInternalI, limitPrice);

        // Convert actual used input internal -> uint (ceil)
        amountInUintNoFee = _internalToUintCeil(amountInInternalUsed, bases[i]);
        require(amountInUintNoFee > 0, "swap: input zero");

        // Compute gross transfer including fee on the used input (ceil)
        feeUint = 0;
        grossIn = amountInUintNoFee;
        if (swapFeePpm > 0) {
            feeUint = _ceilFee(amountInUintNoFee, swapFeePpm);
            grossIn += feeUint;
        }

        // Ensure within user max
        require(grossIn <= maxAmountIn, "swap: transfer exceeds max");

        // Compute output (floor)
        amountOutUint = _internalToUintFloor(amountOutInternal, bases[j]);
        require(amountOutUint > 0, "swap: output zero");
    }

    /// @notice Internal quote for swap-to-limit that mirrors swapToLimit() rounding and fee application
    /// @return grossIn amount to transfer in (inclusive of fee), amountOutUint output amount (uint),
    ///         amountInInternal and amountOutInternal (64.64), amountInUintNoFee input amount excluding fee (uint),
    ///         feeUint fee taken from the gross input (uint)
    function _quoteSwapToLimit(
        uint256 i,
        uint256 j,
        int128 limitPrice
    )
        internal
        view
        returns (
            uint256 grossIn,
            uint256 amountOutUint,
            int128 amountInInternal,
            int128 amountOutInternal,
            uint256 amountInUintNoFee,
            uint256 feeUint
        )
    {
        uint256 n = tokens.length;
        require(i < n && j < n, "swapToLimit: idx");
        require(limitPrice > int128(0), "swapToLimit: limit <= 0");
        require(lmsr.nAssets > 0, "swapToLimit: pool uninitialized");

        // Compute internal maxima at the price limit
        (amountInInternal, amountOutInternal) = lmsr.swapAmountsForPriceLimit(i, j, limitPrice);

        // Convert input to uint (ceil) and output to uint (floor)
        amountInUintNoFee = _internalToUintCeil(amountInInternal, bases[i]);
        require(amountInUintNoFee > 0, "swapToLimit: input zero");

        feeUint = 0;
        grossIn = amountInUintNoFee;
        if (swapFeePpm > 0) {
            feeUint = _ceilFee(amountInUintNoFee, swapFeePpm);
            grossIn += feeUint;
        }

        amountOutUint = _internalToUintFloor(amountOutInternal, bases[j]);
        require(amountOutUint > 0, "swapToLimit: output zero");
    }

    /// @notice External view to quote exact-in swap amounts (gross input incl. fee and output), matching swap() computations
    function swapAmounts(
        uint256 i,
        uint256 j,
        uint256 maxAmountIn,
        int128 limitPrice
    ) external view returns (uint256 amountIn, uint256 amountOut, uint256 fee) {
        (uint256 grossIn, uint256 outUint,,,, uint256 feeUint) = _quoteSwapExactIn(i, j, maxAmountIn, limitPrice);
        return (grossIn, outUint, feeUint);
    }

    /// @notice External view to quote swap-to-limit amounts (gross input incl. fee and output), matching swapToLimit() computations
    function swapToLimitAmounts(
        uint256 i,
        uint256 j,
        int128 limitPrice
    ) external view returns (uint256 amountIn, uint256 amountOut, uint256 fee) {
        (uint256 grossIn, uint256 outUint,,,, uint256 feeUint) = _quoteSwapToLimit(i, j, limitPrice);
        return (grossIn, outUint, feeUint);
    }


    /// @notice Swap input token i -> token j. Payer must approve token i.
    /// @param payer address of the account that pays for the swap
    /// @param receiver address that will receive the output tokens
    /// @param i index of input asset
    /// @param j index of output asset
    /// @param maxAmountIn maximum amount of token i (uint256) to transfer in (inclusive of fees)
    /// @param limitPrice maximum acceptable marginal price (64.64 fixed point). Pass 0 to ignore.
    /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore.
    /// @return amountIn actual input used (uint256), amountOut actual output sent (uint256), fee fee taken from the input (uint256)
    function swap(
        address payer,
        address receiver,
        uint256 i,
        uint256 j,
        uint256 maxAmountIn,
        int128 limitPrice,
        uint256 deadline
    ) external nonReentrant returns (uint256 amountIn, uint256 amountOut, uint256 fee) {
        uint256 n = tokens.length;
        require(i < n && j < n, "swap: idx");
        require(maxAmountIn > 0, "swap: input zero");
        require(deadline == 0 || block.timestamp <= deadline, "swap: deadline exceeded");

        // Read previous balances for affected assets
        uint256 prevBalI = IERC20(tokens[i]).balanceOf(address(this));
        uint256 prevBalJ = IERC20(tokens[j]).balanceOf(address(this));

        // Compute amounts using the same path as views
        (uint256 totalTransferAmount, uint256 amountOutUint, int128 amountInInternalUsed, int128 amountOutInternal, , uint256 feeUint) =
            _quoteSwapExactIn(i, j, maxAmountIn, limitPrice);

        // Transfer the exact amount from payer and require exact receipt (revert on fee-on-transfer)
        _safeTransferFrom(tokens[i], payer, address(this), totalTransferAmount);
        uint256 balIAfter = IERC20(tokens[i]).balanceOf(address(this));
        require(balIAfter == prevBalI + totalTransferAmount, "swap: non-standard tokenIn");

        // Transfer output to receiver and verify exact decrease
        _safeTransfer(tokens[j], receiver, amountOutUint);
        uint256 balJAfter = IERC20(tokens[j]).balanceOf(address(this));
        require(balJAfter == prevBalJ - amountOutUint, "swap: non-standard tokenOut");

        // Update cached uint balances for i and j using actual balances
        cachedUintBalances[i] = balIAfter;
        cachedUintBalances[j] = balJAfter;

        // Apply swap to LMSR state with the internal amounts actually used
        lmsr.applySwap(i, j, amountInInternalUsed, amountOutInternal);

        emit Swap(payer, receiver, tokens[i], tokens[j], totalTransferAmount, amountOutUint);

        return (totalTransferAmount, amountOutUint, feeUint);
    }

    /// @notice Swap up to the price limit; computes max input to reach limit then performs swap.
    /// If the pool can't fill entirely because of balances, it caps appropriately and returns actuals.
    /// Payer must approve token i for the exact computed input amount.
    /// @param deadline timestamp after which the transaction will revert. Pass 0 to ignore.
    function swapToLimit(
        address payer,
        address receiver,
        uint256 i,
        uint256 j,
        int128 limitPrice,
        uint256 deadline
    ) external returns (uint256 amountInUsed, uint256 amountOut, uint256 fee) {
        uint256 n = tokens.length;
        require(i < n && j < n, "swapToLimit: idx");
        require(limitPrice > int128(0), "swapToLimit: limit <= 0");
        require(deadline == 0 || block.timestamp <= deadline, "swapToLimit: deadline exceeded");

        // Read previous balances for affected assets
        uint256 prevBalI = IERC20(tokens[i]).balanceOf(address(this));
        uint256 prevBalJ = IERC20(tokens[j]).balanceOf(address(this));

        // Compute amounts using the same path as views
        (uint256 totalTransferAmount, uint256 amountOutUint, int128 amountInInternalMax, int128 amountOutInternal, uint256 amountInUsedUint, uint256 feeUint) =
            _quoteSwapToLimit(i, j, limitPrice);

        // Transfer the exact amount needed from payer and require exact receipt (revert on fee-on-transfer)
        _safeTransferFrom(tokens[i], payer, address(this), totalTransferAmount);
        uint256 balIAfter = IERC20(tokens[i]).balanceOf(address(this));
        require(balIAfter == prevBalI + totalTransferAmount, "swapToLimit: non-standard tokenIn");

        // Transfer output to receiver and verify exact decrease
        _safeTransfer(tokens[j], receiver, amountOutUint);
        uint256 balJAfter = IERC20(tokens[j]).balanceOf(address(this));
        require(balJAfter == prevBalJ - amountOutUint, "swapToLimit: non-standard tokenOut");

        // Update caches to actual balances
        cachedUintBalances[i] = balIAfter;
        cachedUintBalances[j] = balJAfter;

        // Apply swap to LMSR state with the internal amounts
        lmsr.applySwap(i, j, amountInInternalMax, amountOutInternal);

        // Maintain original event semantics (logs input without fee)
        emit Swap(payer, receiver, tokens[i], tokens[j], amountInUsedUint, amountOutUint);

        return (amountInUsedUint, amountOutUint, feeUint);
    }

    /// @notice Ceiling fee helper: computes ceil(x * feePpm / 1_000_000)
    function _ceilFee(uint256 x, uint256 feePpm) internal pure returns (uint256) {
        if (feePpm == 0) return 0;
        // ceil division: (num + denom - 1) / denom
        return (x * feePpm + 1_000_000 - 1) / 1_000_000;
    }

    /// @notice Compute fee and net amounts for a gross input (fee rounded up to favor the pool).
    /// @return feeUint fee taken (uint) and netUint remaining for protocol use (uint)
    function _computeFee(uint256 gross) internal view returns (uint256 feeUint, uint256 netUint) {
        if (swapFeePpm == 0) {
            return (0, gross);
        }
        feeUint = _ceilFee(gross, swapFeePpm);
        netUint = gross - feeUint;
    }

    /// @notice Convenience: return gross = net + fee(net) using ceiling for fee.
    function _addFee(uint256 netUint) internal view returns (uint256 gross) {
        if (swapFeePpm == 0) return netUint;
        uint256 fee = _ceilFee(netUint, swapFeePpm);
        return netUint + fee;
    }

    // --- New events for single-token mint/burn flows ---
    // Note: events intentionally avoid exposing internal ΔS and avoid duplicating LP mint/burn data
    // which is already present in the standard Mint/Burn events.

    /// @notice Single-token mint: deposit a single token, charge swap-LMSR cost, and mint LP.
    /// @param payer who transfers the input token
    /// @param receiver who receives the minted LP tokens
    /// @param i index of the input token
    /// @param maxAmountIn maximum uint token input (inclusive of fee)
    /// @param deadline optional deadline
    /// @return lpMinted actual LP minted (uint)
    function swapMint(
        address payer,
        address receiver,
        uint256 i,
        uint256 maxAmountIn,
        uint256 deadline
    ) external nonReentrant returns (uint256 lpMinted) {
        uint256 n = tokens.length;
        require(i < n, "swapMint: idx");
        require(maxAmountIn > 0, "swapMint: input zero");
        require(deadline == 0 || block.timestamp <= deadline, "swapMint: deadline");

        // Ensure pool initialized
        require(lmsr.nAssets > 0, "swapMint: uninit pool");

        // compute fee on gross maxAmountIn to get an initial net estimate (we'll recompute based on actual used)
        (, uint256 netUintGuess) = _computeFee(maxAmountIn);

        // Convert the net guess to internal (floor)
        int128 netInternalGuess = _uintToInternalFloor(netUintGuess, bases[i]);
        require(netInternalGuess > int128(0), "swapMint: input too small after fee");

        // Use LMSR view to determine actual internal consumed and size-increase (ΔS) for mint
        (int128 amountInInternalUsed, int128 sizeIncreaseInternal) = lmsr.swapAmountsForMint(i, netInternalGuess);

        // amountInInternalUsed may be <= netInternalGuess. Convert to uint (ceil) to determine actual transfer
        uint256 amountInUint = _internalToUintCeil(amountInInternalUsed, bases[i]);
        require(amountInUint > 0, "swapMint: input zero after internal conversion");

        // Compute fee on the actual used input and total transfer amount (ceiling)
        uint256 feeUintActual = _ceilFee(amountInUint, swapFeePpm);
        uint256 totalTransfer = amountInUint + feeUintActual;
        require(totalTransfer > 0 && totalTransfer <= maxAmountIn, "swapMint: transfer exceeds max");

        // Record pre-balance and transfer tokens from payer, require exact receipt (revert on fee-on-transfer)
        uint256 prevBalI = IERC20(tokens[i]).balanceOf(address(this));
        _safeTransferFrom(tokens[i], payer, address(this), totalTransfer);
        uint256 balIAfter = IERC20(tokens[i]).balanceOf(address(this));
        require(balIAfter == prevBalI + totalTransfer, "swapMint: non-standard tokenIn");

        // Update cached uint balances for token i (only i changed externally)
        cachedUintBalances[i] = balIAfter;

        // Compute old and new scaled size metrics to determine LP minted
        int128 oldTotal = _computeSizeMetric(lmsr.qInternal);
        require(oldTotal > int128(0), "swapMint: zero total");
        uint256 oldScaled = ABDKMath64x64.mulu(oldTotal, LP_SCALE);

        int128 newTotal = oldTotal.add(sizeIncreaseInternal);
        uint256 newScaled = ABDKMath64x64.mulu(newTotal, LP_SCALE);

        uint256 actualLpToMint;
        if (totalSupply() == 0) {
            // If somehow supply zero (shouldn't happen as lmsr.nAssets>0), mint newScaled
            actualLpToMint = newScaled;
        } else {
            require(oldScaled > 0, "swapMint: oldScaled zero");
            uint256 delta = (newScaled > oldScaled) ? (newScaled - oldScaled) : 0;
            if (delta > 0) {
                // floor truncation rounds in favor of pool
                actualLpToMint = (totalSupply() * delta) / oldScaled;
            } else {
                actualLpToMint = 0;
            }
        }

        require(actualLpToMint > 0, "swapMint: zero LP minted");

        // Update LMSR internal state: scale qInternal proportionally by newTotal/oldTotal
        int128[] memory newQInternal = new int128[](n);
        for (uint256 idx = 0; idx < n; idx++) {
            // newQInternal[idx] = qInternal[idx] * (newTotal / oldTotal)
            newQInternal[idx] = lmsr.qInternal[idx].mul(newTotal).div(oldTotal);
        }

        // Update cached internal and kappa via updateForProportionalChange
        lmsr.updateForProportionalChange(newQInternal);

        // Note: we updated cachedUintBalances[i] above via reading balance; other token uint balances did not
        // change externally (they were not transferred in). We keep cachedUintBalances for others unchanged.
        // Mint LP tokens to receiver
        _mint(receiver, actualLpToMint);

        // Emit SwapMint event with gross transfer, net input and fee (planned exact-in)
        emit SwapMint(payer, receiver, i, totalTransfer, amountInUint, feeUintActual);

        // Emit standard Mint event which records deposit amounts and LP minted
        emit Mint(payer, receiver, new uint256[](n), actualLpToMint);
        // Note: depositAmounts array omitted (empty) since swapMint uses single-token input

        return actualLpToMint;
    }

    /// @notice Burn LP tokens then swap the redeemed proportional basket into a single asset `i` and send to receiver.
    /// @param payer who burns LP tokens
    /// @param receiver who receives the single asset
    /// @param lpAmount amount of LP tokens to burn
    /// @param i index of target asset to receive
    /// @param deadline optional deadline
    /// @return amountOutUint uint amount of asset i sent to receiver
    function burnSwap(
        address payer,
        address receiver,
        uint256 lpAmount,
        uint256 i,
        uint256 deadline
    ) external nonReentrant returns (uint256 amountOutUint) {
        uint256 n = tokens.length;
        require(i < n, "burnSwap: idx");
        require(lpAmount > 0, "burnSwap: zero lp");
        require(deadline == 0 || block.timestamp <= deadline, "burnSwap: deadline");

        uint256 supply = totalSupply();
        require(supply > 0, "burnSwap: empty supply");
        require(balanceOf(payer) >= lpAmount, "burnSwap: insufficient LP");

        // alpha = lpAmount / supply as Q64.64
        int128 alpha = ABDKMath64x64.divu(lpAmount, supply);

        // Use LMSR view to compute single-asset payout and burned size-metric
        (int128 payoutInternal, ) = lmsr.swapAmountsForBurn(i, alpha);

        // Convert payoutInternal -> uint (floor) to favor pool
        amountOutUint = _internalToUintFloor(payoutInternal, bases[i]);
        require(amountOutUint > 0, "burnSwap: output zero");

        // Transfer the payout to receiver
        _safeTransfer(tokens[i], receiver, amountOutUint);

        // Burn LP tokens from payer (authorization via allowance)
        if (msg.sender != payer) {
            uint256 allowed = allowance(payer, msg.sender);
            require(allowed >= lpAmount, "burnSwap: allowance insufficient");
            _approve(payer, msg.sender, allowed - lpAmount);
        }
        _burn(payer, lpAmount);

        // Update cached balances by reading on-chain balances for all tokens
        int128[] memory newQInternal = new int128[](n);
        for (uint256 idx = 0; idx < n; idx++) {
            uint256 bal = IERC20(tokens[idx]).balanceOf(address(this));
            cachedUintBalances[idx] = bal;
            newQInternal[idx] = _uintToInternalFloor(bal, bases[idx]);
        }

        // Emit BurnSwap with public-facing info only (do not expose ΔS or LP burned)
        emit BurnSwap(payer, receiver, i, amountOutUint);

        // If entire pool drained, deinit; else update proportionally
        bool allZero = true;
        for (uint256 idx = 0; idx < n; idx++) {
            if (newQInternal[idx] != int128(0)) { allZero = false; break; }
        }
        if (allZero) {
            lmsr.deinit();
        } else {
            lmsr.updateForProportionalChange(newQInternal);
        }

        emit Burn(payer, receiver, new uint256[](n), lpAmount);
        return amountOutUint;
    }


    function flashRepaymentAmounts(uint256[] memory loanAmounts) external view
    returns (uint256[] memory repaymentAmounts) {
        repaymentAmounts = new uint256[](tokens.length);
        for (uint256 i = 0; i < tokens.length; i++) {
            uint256 amount = loanAmounts[i];
            if (amount > 0) {
                repaymentAmounts[i] = amount + _ceilFee(amount, flashFeePpm);
            }
        }
    }
    
    
    /// @notice Receive token0 and/or token1 and pay it back, plus a fee, in the callback
    /// @dev The caller of this method receives a callback in the form of IPartyFlashCallback#partyFlashCallback
    /// @param recipient The address which will receive the token amounts
    /// @param amounts The amount of each token to send
    /// @param data Any data to be passed through to the callback
    function flash(
        address recipient,
        uint256[] memory amounts,
        bytes calldata data
    ) external nonReentrant {
        require(recipient != address(0), "flash: zero recipient");
        require(amounts.length == tokens.length, "flash: amounts length mismatch");

        // Calculate repayment amounts for each token including fee
        uint256[] memory repaymentAmounts = new uint256[](tokens.length);

        // Store initial balances to verify repayment later
        uint256[] memory initialBalances = new uint256[](tokens.length);

        // Track if any token amount is non-zero
        bool hasNonZeroAmount = false;

        // Process each token, skipping those with zero amounts
        for (uint256 i = 0; i < tokens.length; i++) {
            uint256 amount = amounts[i];

            if (amount > 0) {
                hasNonZeroAmount = true;

                // Calculate repayment amount with fee (ceiling)
                repaymentAmounts[i] = amount + _ceilFee(amount, flashFeePpm);

                // Record initial balance
                initialBalances[i] = IERC20(tokens[i]).balanceOf(address(this));

                // Transfer token to recipient
                _safeTransfer(tokens[i], recipient, amount);
            }
        }

        // Ensure at least one token is being borrowed
        require(hasNonZeroAmount, "flash: no tokens requested");

        // Call flash callback with expected repayment amounts
        IPartyFlashCallback(msg.sender).partyFlashCallback(amounts, repaymentAmounts, data);

        // Verify repayment amounts for tokens that were borrowed
        for (uint256 i = 0; i < tokens.length; i++) {
            if (amounts[i] > 0) {
                uint256 currentBalance = IERC20(tokens[i]).balanceOf(address(this));

                // Verify repayment: current balance must be at least (initial balance + fee)
                require(
                    currentBalance >= initialBalances[i] + _ceilFee(amounts[i], flashFeePpm),
                    "flash: repayment failed"
                );

                // Update cached balance
                cachedUintBalances[i] = currentBalance;
            }
        }
    }

    
    /* ----------------------
       Conversion helpers
       ---------------------- */

    // Convert uint token amount -> internal 64.64 (floor). Uses ABDKMath64x64.divu which truncates.
    function _uintToInternalFloor(uint256 amount, uint256 base) internal pure returns (int128) {
        // internal = amount / base  (as Q64.64)
        return ABDKMath64x64.divu(amount, base);
    }

    // Convert internal 64.64 -> uint token amount (floor). Uses ABDKMath64x64.mulu which floors the product.
    function _internalToUintFloor(int128 internalAmount, uint256 base) internal pure returns (uint256) {
        // uint = internal * base (floored)
        return ABDKMath64x64.mulu(internalAmount, base);
    }

    // Convert internal 64.64 -> uint token amount (ceiling). Rounds up to protect the pool.
    function _internalToUintCeil(int128 internalAmount, uint256 base) internal pure returns (uint256) {
        // Get the floor value first
        uint256 floorValue = ABDKMath64x64.mulu(internalAmount, base);

        // Check if there was any fractional part by comparing to a reconstruction of the original
        int128 reconstructed = ABDKMath64x64.divu(floorValue, base);

        // If reconstructed is less than original, there was a fractional part that was truncated
        if (reconstructed < internalAmount) {
            return floorValue + 1;
        }

        return floorValue;
    }

    /* ----------------------
       ERC20 helpers (minimal)
       ---------------------- */

    function _safeTransferFrom(address token, address from, address to, uint256 amt) internal {
        IERC20(token).safeTransferFrom(from, to, amt);
    }

    function _safeTransfer(address token, address to, uint256 amt) internal {
        IERC20(token).safeTransfer(to, amt);
    }

    /// @notice Helper to compute size metric (sum of all asset quantities) from internal balances
    function _computeSizeMetric(int128[] memory qInternal_) private pure returns (int128) {
        int128 total = int128(0);
        for (uint i = 0; i < qInternal_.length; ) {
            total = total.add(qInternal_[i]);
            unchecked { i++; }
        }
        return total;
    }

}