waterpump-evm-pool-sdk 0.1.0

use alloy::primitives::{
    address,
    aliases::{I24, U160},
    Address, Bytes, U256,
};
use anyhow::{Context, Result};
use uniswap_sdk_core::{
    entities::{BaseCurrency, BaseCurrencyCore, FractionBase},
    prelude::{Currency, Percent, Price},
    utils::FromBig,
};
use uniswap_v3_sdk::prelude::encode_multicall;
// Import V4 SDK position_manager functions
// Based on the reference: https://github.com/shuhuiluo/uniswap-v4-sdk-rs/blob/main/src/position_manager.rs
use uniswap_v4_sdk::position_manager::{
    encode_modify_liquidities, encode_permit_batch, AddLiquidityOptions as V4AddLiquidityOptions,
    AddLiquiditySpecificOptions as V4AddLiquiditySpecificOptions, CommonOptions,
    MintSpecificOptions, ModifyPositionSpecificOptions, OPEN_DELTA,
};

use crate::{
    common::{
        mint_amounts_utils::{calculate_mint_amounts, MintAmounts},
        slippage_utils::ratios_after_slippage,
    },
    pool_infusers::common::{
        calculate_position_amounts, calculate_slippage_context, V4PositionPlanner,
    },
    pool_swappers::common::MethodParameters,
    traits::pool_infuser::{
        AddBatchLiquidityParams, AddLiquidityOptions, AddLiquidityParams,
        AddLiquiditySpecificOptions, CollectOptions, CollectParams, RemoveBatchLiquidityParams,
        RemoveLiquidityOptions, RemoveLiquidityParams,
    },
    types::v4_pool_key::V4PoolKey,
};

// Constants from V4 SDK
const MSG_SENDER: Address = address!("0000000000000000000000000000000000000001");

/// Result of building add liquidity call parameters with amount estimates
#[derive(Debug, Clone)]
pub struct AddLiquidityCallResult {
    /// Transaction method parameters (calldata and value)
    pub method_parameters: MethodParameters,
    /// Estimated maximum amount of token0 needed
    pub amount0_max: U256,
    /// Estimated maximum amount of token1 needed
    pub amount1_max: U256,
    /// Estimated liquidity amount for the position
    pub liquidity: u128,
}

/// Result of building remove liquidity call parameters with amount estimates
#[derive(Debug, Clone)]
pub struct RemoveLiquidityCallResult {
    /// Transaction method parameters (calldata and value)
    pub method_parameters: MethodParameters,
    /// Estimated minimum amount of token0 to receive
    pub amount0_min: U256,
    /// Estimated minimum amount of token1 to receive
    pub amount1_min: U256,
    /// Estimated liquidity amount being removed
    pub liquidity: u128,
}

/// Convert our AddLiquiditySpecificOptions to V4 SDK's
/// AddLiquiditySpecificOptions
fn convert_add_liquidity_specific_options(
    opts: AddLiquiditySpecificOptions,
    sqrt_price_x96: U160,
) -> V4AddLiquiditySpecificOptions {
    match opts {
        AddLiquiditySpecificOptions::Mint(mint_opts) => {
            V4AddLiquiditySpecificOptions::Mint(MintSpecificOptions {
                recipient: mint_opts.recipient,
                create_pool: mint_opts.create_pool,
                sqrt_price_x96: if mint_opts.create_pool { Some(sqrt_price_x96) } else { None },
                migrate: false,
            })
        }
        AddLiquiditySpecificOptions::Increase(increase_opts) => {
            V4AddLiquiditySpecificOptions::Increase(ModifyPositionSpecificOptions {
                token_id: increase_opts.token_id,
            })
        }
    }
}

/// Convert our AddLiquidityOptions to V4 SDK's AddLiquidityOptions
fn convert_add_liquidity_options(
    pool_key: &V4PoolKey,
    options: AddLiquidityOptions,
    specific_opts: AddLiquiditySpecificOptions,
    sqrt_price_x96: U160,
) -> V4AddLiquidityOptions {
    V4AddLiquidityOptions {
        common_opts: CommonOptions {
            slippage_tolerance: options.slippage_tolerance,
            deadline: options.deadline,
            hook_data: Bytes::from(pool_key.hooks.into_word()),
        },
        use_native: options.use_native,
        batch_permit: None, // TODO: Convert permit options if needed
        specific_opts: convert_add_liquidity_specific_options(specific_opts, sqrt_price_x96),
    }
}

/// Build call parameters for adding liquidity (mint or increase)
///
/// Uses V4 SDK's Position::mint_amounts_with_slippage approach with
/// V4PositionPlanner to build the transaction.
/// Based on: https://github.com/shuhuiluo/uniswap-v4-sdk-rs/blob/main/src/position_manager.rs#L175
pub fn build_add_liquidity_call_parameters(
    pool_key: &V4PoolKey,
    params: AddLiquidityParams,
    options: AddLiquidityOptions,
) -> Result<AddLiquidityCallResult> {
    // Get current sqrt_price_x96 from token0_price
    let sqrt_price_x96 = encode_sqrt_price_from_price(&params.token0_price)?;

    // Calculate position amounts with slippage using common utility
    let amount0_desired = U256::from_big_int(params.amount0.quotient());
    let amount1_desired = U256::from_big_int(params.amount1.quotient());

    let slippage_ctx =
        calculate_slippage_context(&params.token0_price, &options.slippage_tolerance)?;
    let amounts = calculate_position_amounts(
        params.tick_lower,
        params.tick_upper,
        amount0_desired,
        amount1_desired,
        &slippage_ctx,
    )?;

    assert!(amounts.liquidity > 0, "ZERO_LIQUIDITY");

    // Convert options to V4 SDK format
    let v4_options = convert_add_liquidity_options(
        pool_key,
        options.clone(),
        params.specific_opts,
        sqrt_price_x96,
    );

    // Following the reference implementation from position_manager.rs:179-289
    let mut calldatas: Vec<Bytes> = Vec::with_capacity(3);
    let mut planner = V4PositionPlanner::default();

    // Encode initialize pool if needed
    // if let V4AddLiquiditySpecificOptions::Mint(mint_opts) =
    // v4_options.specific_opts {     if mint_opts.create_pool {
    //         let pool_key_v4 = Pool::get_pool_key(
    //             &pool_key.token_a,
    //             &pool_key.token_b,
    //             pool_key.fee,
    //             pool_key.tick_spacing,
    //             pool_key.hooks,
    //         )?;
    //         let init_sqrt_price =
    // mint_opts.sqrt_price_x96.unwrap_or(sqrt_price_x96);         calldatas.
    // push(encode_initialize_pool(pool_key_v4, init_sqrt_price));     }
    // }

    // position.pool.currency0 is native if and only if options.useNative is set
    assert!(
        if let Some(ether) = &v4_options.use_native {
            pool_key.token_a.equals(ether)
        } else {
            !pool_key.token_a.is_native()
        },
        "NATIVE_NOT_SET"
    );

    // We use permit2 to approve tokens to the position manager
    if let Some(batch_permit) = v4_options.batch_permit {
        calldatas.push(encode_permit_batch(
            batch_permit.owner,
            batch_permit.permit_batch,
            batch_permit.signature,
        ));
    }

    // Convert amounts to u128 for planner (use max amounts for mint/increase)
    let mint_amounts = mint_amounts_with_slippage(
        amounts.liquidity,
        params.tick_lower,
        params.tick_upper,
        &params.token0_price,
        &options.slippage_tolerance,
    )?;
    let amount0_max = mint_amounts.amount0.to::<u128>();
    let amount1_max = mint_amounts.amount1.to::<u128>();

    match v4_options.specific_opts {
        V4AddLiquiditySpecificOptions::Mint(mint_opts) => {
            planner.add_mint(
                pool_key,
                params.tick_lower,
                params.tick_upper,
                U256::from(amounts.liquidity),
                amount0_max,
                amount1_max,
                mint_opts.recipient,
            )?;
        }
        V4AddLiquiditySpecificOptions::Increase(increase_opts) => {
            planner.add_increase(
                increase_opts.token_id,
                U256::from(amounts.liquidity),
                amount0_max,
                amount1_max,
                v4_options.common_opts.hook_data,
            );
        }
    }

    let mut value = U256::ZERO;

    // If migrating, we need to settle and sweep both currencies individually
    match v4_options.specific_opts {
        V4AddLiquiditySpecificOptions::Mint(mint_opts) if mint_opts.migrate => {
            if v4_options.use_native.is_some() {
                // unwrap the exact amount needed to send to the pool manager
                planner.add_unwrap(OPEN_DELTA);
                // payer is v4 position manager
                planner.add_settle(&pool_key.token_a, false, None);
                planner.add_settle(&pool_key.token_b, false, None);
                // sweep any leftover wrapped native that was not unwrapped
                // recipient will be the same as the v4 lp token recipient
                planner.add_sweep(pool_key.token_a.wrapped(), mint_opts.recipient);
                planner.add_sweep(&pool_key.token_b, mint_opts.recipient);
            } else {
                // payer is v4 position manager
                planner.add_settle(&pool_key.token_a, false, None);
                planner.add_settle(&pool_key.token_b, false, None);
                // recipient will be the same as the v4 lp token recipient
                planner.add_sweep(&pool_key.token_a, mint_opts.recipient);
                planner.add_sweep(&pool_key.token_b, mint_opts.recipient);
            }
        }
        _ => {
            // need to settle both currencies when minting / adding liquidity (user is the
            // payer)
            planner.add_settle_pair(&pool_key.token_a, &pool_key.token_b);
            // When not migrating and adding native currency, add a final sweep
            if v4_options.use_native.is_some() {
                // Any sweeping must happen after the settling.
                // native currency will always be currency0 in v4
                value = U256::from(amount0_max);
                planner.add_sweep(&pool_key.token_a, MSG_SENDER);
            }
        }
    }

    calldatas
        .push(encode_modify_liquidities(planner.0.finalize(), v4_options.common_opts.deadline));

    Ok(AddLiquidityCallResult {
        method_parameters: MethodParameters { calldata: encode_multicall(calldatas), value },
        amount0_max: mint_amounts.amount0,
        amount1_max: mint_amounts.amount1,
        liquidity: amounts.liquidity,
    })
}

/// Calculate mint amounts with slippage tolerance
///
/// This function calculates the maximum token amounts needed to mint a position
/// with the given liquidity, accounting for price slippage. It uses the
/// worst-case price scenarios:
/// - amount0: maximum needed when price goes up (calculated at upper slippage
///   price)
/// - amount1: maximum needed when price goes down (calculated at lower slippage
///   price)
///
/// # Arguments
/// * `liquidity` - The liquidity amount for the position
/// * `tick_lower` - Lower tick of the position
/// * `tick_upper` - Upper tick of the position
/// * `token0_price` - Current price of token0 in terms of token1
/// * `slippage_tolerance` - Maximum acceptable price slippage
///
/// # Returns
/// `MintAmounts` containing the maximum amount0 and amount1 needed
fn mint_amounts_with_slippage(
    liquidity: u128,
    tick_lower: I24,
    tick_upper: I24,
    token0_price: &Price<Currency, Currency>,
    slippage_tolerance: &Percent,
) -> Result<MintAmounts> {
    // Calculate price bounds after slippage
    let (sqrt_ratio_x96_lower, sqrt_ratio_x96_upper) =
        ratios_after_slippage(token0_price, slippage_tolerance)
            .context("Failed to calculate ratios after slippage")?;

    // Calculate mint amounts at worst-case prices:
    // - amount0 at upper price (price goes up, needs more amount0)
    // - amount1 at lower price (price goes down, needs more amount1)
    let mint_amounts_upper =
        calculate_mint_amounts(tick_lower, tick_upper, liquidity, sqrt_ratio_x96_upper)
            .context("Failed to calculate mint amounts at upper price")?;

    let mint_amounts_lower =
        calculate_mint_amounts(tick_lower, tick_upper, liquidity, sqrt_ratio_x96_lower)
            .context("Failed to calculate mint amounts at lower price")?;

    Ok(MintAmounts { amount0: mint_amounts_upper.amount0, amount1: mint_amounts_lower.amount1 })
}

/// Calculate burn amounts with slippage tolerance
///
/// This function calculates the minimum token amounts that will be received
/// when removing liquidity, accounting for price slippage. It uses the
/// worst-case price scenarios:
/// - amount0_min: minimum received when price goes down (calculated at lower
///   slippage price)
/// - amount1_min: minimum received when price goes up (calculated at upper
///   slippage price)
///
/// # Arguments
/// * `liquidity` - The liquidity amount to remove
/// * `tick_lower` - Lower tick of the position
/// * `tick_upper` - Upper tick of the position
/// * `token0_price` - Current price of token0 in terms of token1
/// * `slippage_tolerance` - Maximum acceptable price slippage
///
/// # Returns
/// Tuple `(amount0_min, amount1_min)` containing the minimum amounts to receive
fn burn_amounts_with_slippage(
    liquidity: u128,
    tick_lower: I24,
    tick_upper: I24,
    token0_price: &Price<Currency, Currency>,
    slippage_tolerance: &Percent,
) -> Result<(U256, U256)> {
    // Calculate price bounds after slippage
    let (sqrt_ratio_x96_lower, sqrt_ratio_x96_upper) =
        ratios_after_slippage(token0_price, slippage_tolerance)
            .context("Failed to calculate ratios after slippage")?;

    // Calculate mint amounts at worst-case prices for minimum received:
    // - amount0_min at lower price (price goes down, receive less amount0)
    // - amount1_min at upper price (price goes up, receive less amount1)
    let mint_amounts_lower =
        calculate_mint_amounts(tick_lower, tick_upper, liquidity, sqrt_ratio_x96_lower)
            .context("Failed to calculate mint amounts at lower price")?;

    let mint_amounts_upper =
        calculate_mint_amounts(tick_lower, tick_upper, liquidity, sqrt_ratio_x96_upper)
            .context("Failed to calculate mint amounts at upper price")?;

    Ok((mint_amounts_lower.amount0, mint_amounts_upper.amount1))
}

/// Encode sqrt price from Price
fn encode_sqrt_price_from_price(price: &Price<Currency, Currency>) -> Result<U160> {
    use uniswap_sdk_core::utils::FromBig;
    use waterpump_evm_amm_math::encode_sqrt_ratio::encode_sqrt_ratio_x96;
    let fraction = price.as_fraction();
    let num = U256::from_big_int(fraction.numerator);
    let den = U256::from_big_int(fraction.denominator);
    let result = encode_sqrt_ratio_x96(num, den)?;
    Ok(U160::from(result))
}

/// Build call parameters for adding batch liquidity (multiple positions in one
/// transaction)
///
/// Uses V4 SDK's planner to batch multiple positions into a single
/// modifyLiquidities call. Uses a shared token0_price for all positions.
pub fn build_add_batch_liquidity_call_parameters(
    _pool_key: &V4PoolKey,
    _params: AddBatchLiquidityParams,
    _options: AddLiquidityOptions,
) -> Result<MethodParameters> {
    // Get sqrt_price_x96 from shared token0_price
    unimplemented!()
}

/// Build call parameters for removing liquidity
///
/// Based on the reference implementation from position_manager.rs:308-391
pub fn build_remove_liquidity_call_parameters(
    pool_key: &V4PoolKey,
    params: RemoveLiquidityParams,
    options: RemoveLiquidityOptions,
) -> Result<RemoveLiquidityCallResult> {
    // Get sqrt_price_x96 from token0_price

    // Create position representing the portion to remove
    // Note: params.liquidity is the absolute amount to remove
    let liquidity_u128 = params.liquidity.to::<u128>();

    let mut calldatas: Vec<Bytes> = Vec::with_capacity(2);
    let mut planner = V4PositionPlanner::default();

    let token_id = params.token_id;
    let hook_data = Bytes::from(pool_key.hooks.into_word());

    if options.burn_token {
        // if burnToken is true, the specified liquidity percentage must be 100%
        // Since we're using absolute liquidity, we assume params.liquidity represents
        // the full position In practice, the caller should ensure this matches
        // the full position liquidity

        // if there is a permit, encode the ERC721Permit permit call
        if let Some(permit) = options.permit {
            use uniswap_v4_sdk::position_manager::encode_erc721_permit;
            calldatas.push(encode_erc721_permit(
                permit.spender,
                token_id,
                permit.deadline,
                permit.nonce,
                permit.signature.as_bytes().into(),
            ));
        }

        // Slippage-adjusted amounts derived from position liquidity
        let (amount0_min, amount1_min) = burn_amounts_with_slippage(
            liquidity_u128,
            params.tick_lower,
            params.tick_upper,
            &params.token0_price,
            &options.slippage_tolerance,
        )?;
        planner.add_burn(token_id, amount0_min.to::<u128>(), amount1_min.to::<u128>(), hook_data);
    } else {
        // Partial removal: use the position as-is (it already represents the portion to
        // remove)
        assert!(liquidity_u128 > 0, "ZERO_LIQUIDITY");

        // Slippage-adjusted underlying amounts
        let (amount0_min, amount1_min) = burn_amounts_with_slippage(
            liquidity_u128,
            params.tick_lower,
            params.tick_upper,
            &params.token0_price,
            &options.slippage_tolerance,
        )?;
        planner.add_decrease(
            token_id,
            U256::from(liquidity_u128),
            amount0_min.to::<u128>(),
            amount1_min.to::<u128>(),
            hook_data,
        );
    }

    planner.add_take_pair(&pool_key.token_a, &pool_key.token_b, params.recipient);

    calldatas.push(encode_modify_liquidities(planner.0.finalize(), params.deadline));

    // Calculate burn amounts for the result (same calculation used above)
    let (amount0_min, amount1_min) = burn_amounts_with_slippage(
        liquidity_u128,
        params.tick_lower,
        params.tick_upper,
        &params.token0_price,
        &options.slippage_tolerance,
    )?;

    Ok(RemoveLiquidityCallResult {
        method_parameters: MethodParameters {
            calldata: encode_multicall(calldatas),
            value: U256::ZERO,
        },
        amount0_min,
        amount1_min,
        liquidity: liquidity_u128,
    })
}

/// Build call parameters for batch removing liquidity (multiple positions in
/// one transaction)
pub fn build_remove_batch_liquidity_call_parameters(
    _pool_key: &V4PoolKey,
    _params: RemoveBatchLiquidityParams,
    _options: RemoveLiquidityOptions,
) -> Result<MethodParameters> {
    unimplemented!()
}

/// Build call parameters for collecting fees
pub fn build_collect_call_parameters(
    pool_key: &V4PoolKey,
    params: &CollectParams,
    _options: &CollectOptions,
) -> Result<MethodParameters> {
    let mut planner = V4PositionPlanner::default();

    // To collect fees in V4, we need to:
    // - encode a decrease liquidity by 0
    // - and encode a TAKE_PAIR
    planner.add_decrease(
        params.token_id,
        U256::ZERO,
        0,
        0,
        Bytes::from(pool_key.hooks.into_word()),
    );

    planner.add_take_pair(&pool_key.token_a, &pool_key.token_b, params.recipient);

    let deadline = chrono::Utc::now().timestamp() + 1200;

    Ok(MethodParameters {
        calldata: encode_modify_liquidities(planner.0.finalize(), U256::from(deadline)),
        value: U256::ZERO,
    })
}