waterpump-evm-pool-sdk 0.1.0

use std::time::{SystemTime, UNIX_EPOCH};

use alloy::primitives::{aliases::I24, Address, Bytes, U160, U256};
use alloy_sol_types::SolCall;
use anyhow::{Context, Result};
use tracing::{debug, info};
use uniswap_sdk_core::{
    entities::{BaseCurrency, BaseCurrencyCore, FractionBase},
    prelude::{BigInt, TradeType},
    utils::FromBig,
};
use uniswap_v3_sdk::prelude::{encode_multicall, encode_permit, encode_refund_eth};
use waterpump_evm_core::pool_key::SlipstreamPoolKey;
use waterpump_evm_slipstream_client::interfaces::{IQuoter, IQuoterV2, ISwapRouter};

use crate::{
    pool_swappers::common::{determine_two_hop_path_tokens, MethodParameters},
    types::swap_params::{QuoteOptions, SwapOptions, SwapParams},
};

/// Encode a two-hop path for Slipstream quoter
/// Path format: token0 (20 bytes) + tickSpacing0 (3 bytes, packed as int24) +
/// token1 (20 bytes) + tickSpacing1 (3 bytes) + token2 (20 bytes)
/// For ExactOutput, the path should be in reverse order
fn encode_two_hop_path(
    token0: Address,
    tick_spacing0: I24,
    token1: Address,
    tick_spacing1: I24,
    token2: Address,
) -> Bytes {
    let mut path = Vec::with_capacity(20 + 3 + 20 + 3 + 20);

    // token0 (20 bytes)
    path.extend_from_slice(token0.as_slice());

    // tickSpacing0 (3 bytes as int24)
    path.extend_from_slice(&tick_spacing0.to_be_bytes::<3>());

    // token1 (20 bytes)
    path.extend_from_slice(token1.as_slice());

    // tickSpacing1 (3 bytes as int24)
    path.extend_from_slice(&tick_spacing1.to_be_bytes::<3>());

    // token2 (20 bytes)
    path.extend_from_slice(token2.as_slice());

    path.into()
}

#[inline]
#[tracing::instrument(skip(pool_key, swap_params, swap_options), fields(token_a = ?pool_key.token_a.address(), token_b = ?pool_key.token_b.address(), tick_spacing = ?pool_key.tick_spacing, is_to_b = ?swap_params.is_to_b, trade_type = ?swap_params.trade_type))]
pub fn build_swap_call_parameters(
    pool_key: &SlipstreamPoolKey,
    swap_params: SwapParams,
    swap_options: SwapOptions,
) -> Result<MethodParameters> {
    let SwapOptions {
        amount_in_maximum,
        amount_out_minimum,
        recipient,
        input_token_permit, // Slipstream doesn't use permit in the same way
        sqrt_price_limit_x96,
        fee: _fee, // Not used for Slipstream
    } = swap_options;

    if pool_key.token_a == pool_key.token_b {
        return Err(anyhow::anyhow!("Token A and Token B cannot be the same"));
    }
    let (token_in, token_out) = if swap_params.is_to_b {
        (pool_key.token_a.clone(), pool_key.token_b.clone())
    } else {
        (pool_key.token_b.clone(), pool_key.token_a.clone())
    };

    debug!(
        token_in = ?token_in.address(),
        token_out = ?token_out.address(),
        is_to_b = ?swap_params.is_to_b,
        "Determined swap direction"
    );

    let mut calldatas: Vec<Bytes> = Vec::with_capacity(1 + 3);

    // encode permit if necessary
    if let Some(input_token_permit) = input_token_permit {
        assert!(!token_in.is_native(), "NON_TOKEN_PERMIT");
        calldatas.push(encode_permit(&token_in, input_token_permit));
    }

    let trade_type = swap_params.trade_type;

    // Get deadline (current time + 20 minutes)
    let deadline = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .context("Failed to get current time")?
        .as_secs()
        + 1200; // 20 minutes

    let amount_in_max_clone = amount_in_maximum.clone();
    let amount_out_min_clone = amount_out_minimum.clone();

    let swap_calldata: Bytes = match trade_type {
        TradeType::ExactInput => {
            let amount_out_minimum = if let Some(amount_out_minimum) = amount_out_min_clone {
                U256::from_big_int(amount_out_minimum.quotient())
            } else {
                U256::ZERO
            };
            let amount_in = U256::from_big_int(swap_params.amount.quotient());
            debug!(
                amount_in = ?amount_in,
                sqrt_price_limit_x96 = ?sqrt_price_limit_x96,
                amount_out_minimum = ?amount_out_minimum,
                "Building ExactInput swap/n/n"
            );

            info!(
                token_in = ?token_in.address(),
                token_out = ?token_out.address(),
                tick_spacing = ?pool_key.tick_spacing,
                recipient = ?recipient,
                deadline = ?deadline,
                amount_in = ?amount_in,
                amount_out_minimum = ?amount_out_minimum,
                sqrt_price_limit_x96 = ?sqrt_price_limit_x96,
                "Building ExactInput swap/n/n"
            );
            ISwapRouter::exactInputSingleCall {
                params: ISwapRouter::ExactInputSingleParams {
                    tokenIn: token_in.address(),
                    tokenOut: token_out.address(),
                    tickSpacing: pool_key.tick_spacing,
                    recipient,
                    deadline: U256::from(deadline),
                    amountIn: amount_in,
                    amountOutMinimum: amount_out_minimum,
                    sqrtPriceLimitX96: sqrt_price_limit_x96.unwrap_or_default(),
                },
            }
            .abi_encode()
            .into()
        }
        TradeType::ExactOutput => {
            let amount_in_maximum = if let Some(amount_in_maximum) = amount_in_max_clone {
                U256::from_big_int(amount_in_maximum.quotient())
            } else {
                U256::MAX
            };

            info!(
                amount_in_maximum = ?amount_in_maximum,
                amount_out = ?U256::from_big_int(swap_params.amount.quotient()),
                "Building ExactOutput swap/n/n"
            );
            info!(
                token_in = ?token_in.address(),
                token_out = ?token_out.address(),
                tick_spacing = ?pool_key.tick_spacing,
                recipient = ?recipient,
                deadline = ?deadline,
                amount_out = ?U256::from_big_int(swap_params.amount.quotient()),
                amount_in_maximum = ?amount_in_maximum,
                sqrt_price_limit_x96 = ?sqrt_price_limit_x96,
                "Building ExactOutput swap/n/n"
            );
            ISwapRouter::exactOutputSingleCall {
                params: ISwapRouter::ExactOutputSingleParams {
                    tokenIn: token_in.address(),
                    tokenOut: token_out.address(),
                    tickSpacing: pool_key.tick_spacing,
                    recipient,
                    deadline: U256::from(deadline),
                    amountOut: U256::from_big_int(swap_params.amount.quotient()),
                    amountInMaximum: amount_in_maximum,
                    sqrtPriceLimitX96: U160::from(0),
                },
            }
            .abi_encode()
            .into()
        }
    };

    info!(
        token_in = ?token_in.is_native(),
        token_out = ?token_out.is_native(),
        trade_type = ?trade_type,
        "Building swap call parameters"
    );

    calldatas.push(swap_calldata);

    // flag for whether a refund needs to happen
    let must_refund = token_in.is_native() && trade_type == TradeType::ExactOutput;
    // refund
    if must_refund {
        calldatas.push(encode_refund_eth());
    }

    let value = if token_in.is_native() {
        match trade_type {
            TradeType::ExactOutput => {
                amount_in_maximum.context("Failed to get amount in maximum")?.quotient()
            }
            TradeType::ExactInput => swap_params.amount.quotient(),
        }
    } else {
        BigInt::ZERO
    };

    let calldata_count = calldatas.len();
    let final_calldata = encode_multicall(calldatas);
    info!(
        calldata_count = calldata_count,
        final_calldata_len = final_calldata.len(),
        value = ?U256::from_big_int(value),
        trade_type = ?trade_type,
        "Swap call parameters finalized"
    );
    Ok(MethodParameters { calldata: final_calldata, value: U256::from_big_int(value) })
}

#[inline]
#[tracing::instrument(skip(pool_key, swap_params, options), fields(token_a = ?pool_key.token_a.address(), token_b = ?pool_key.token_b.address(), tick_spacing = ?pool_key.tick_spacing, is_to_b = ?swap_params.is_to_b, trade_type = ?swap_params.trade_type, use_quoter_v2 = ?options.use_quoter_v2))]
pub fn build_quote_call_parameters(
    pool_key: &SlipstreamPoolKey,
    swap_params: SwapParams,
    options: QuoteOptions,
) -> Result<MethodParameters> {
    let QuoteOptions { sqrt_price_limit_x96, use_quoter_v2 } = options;
    let sqrt_price_limit_x96 = sqrt_price_limit_x96.unwrap_or_default();

    let quote_amount = U256::from_big_int(swap_params.amount.quotient());

    let (token_in, token_out) = if swap_params.is_to_b {
        (pool_key.token_a.clone(), pool_key.token_b.clone())
    } else {
        (pool_key.token_b.clone(), pool_key.token_a.clone())
    };
    debug!(
        token_in = ?token_in.address(),
        token_out = ?token_out.address(),
        is_to_b = ?swap_params.is_to_b,
        trade_type = ?swap_params.trade_type,
        use_quoter_v2 = ?use_quoter_v2,
        "Building quote call parameters"
    );
    let calldata = match swap_params.trade_type {
        TradeType::ExactInput => {
            if use_quoter_v2 {
                IQuoterV2::quoteExactInputSingleCall {
                    params: IQuoterV2::QuoteExactInputSingleParams {
                        tokenIn: token_in.address(),
                        tokenOut: token_out.address(),
                        amountIn: quote_amount,
                        tickSpacing: pool_key.tick_spacing,
                        sqrtPriceLimitX96: sqrt_price_limit_x96,
                    },
                }
                .abi_encode()
            } else {
                IQuoter::quoteExactInputSingleCall {
                    tokenIn: token_in.address(),
                    tokenOut: token_out.address(),
                    tickSpacing: pool_key.tick_spacing,
                    amountIn: quote_amount,
                    sqrtPriceLimitX96: sqrt_price_limit_x96,
                }
                .abi_encode()
            }
        }
        TradeType::ExactOutput => {
            if use_quoter_v2 {
                IQuoterV2::quoteExactOutputSingleCall {
                    params: IQuoterV2::QuoteExactOutputSingleParams {
                        tokenIn: token_in.address(),
                        tokenOut: token_out.address(),
                        amount: quote_amount,
                        tickSpacing: pool_key.tick_spacing,
                        sqrtPriceLimitX96: sqrt_price_limit_x96,
                    },
                }
                .abi_encode()
            } else {
                IQuoter::quoteExactOutputSingleCall {
                    tokenIn: token_in.address(),
                    tokenOut: token_out.address(),
                    tickSpacing: pool_key.tick_spacing,
                    amountOut: quote_amount,
                    sqrtPriceLimitX96: sqrt_price_limit_x96,
                }
                .abi_encode()
            }
        }
    };

    debug!(calldata_len = calldata.len(), "Quote call parameters finalized");

    Ok(MethodParameters { calldata: calldata.into(), value: U256::ZERO })
}

/// Build swap call parameters for a two-hop swap
#[inline]
#[tracing::instrument(skip(pool1_key, pool2_key, swap_params, swap_options), fields(token_a = ?pool1_key.token_a.address(), token_b = ?pool1_key.token_b.address(), token_c = ?pool2_key.token_b.address(), is_to_b = ?swap_params.is_to_b, trade_type = ?swap_params.trade_type))]
pub fn build_two_hop_swap_call_parameters(
    pool1_key: &SlipstreamPoolKey,
    pool2_key: &SlipstreamPoolKey,
    swap_params: SwapParams,
    swap_options: SwapOptions,
) -> Result<MethodParameters> {
    // Forbid ExactOutput for two-hop swaps
    if swap_params.trade_type == TradeType::ExactOutput {
        return Err(anyhow::anyhow!(
            "ExactOutput is not supported for two-hop swaps. Use ExactInput instead."
        ));
    }

    let SwapOptions {
        amount_in_maximum: _amount_in_maximum, // Not used for ExactInput
        amount_out_minimum,
        recipient,
        input_token_permit,
        sqrt_price_limit_x96: _sqrt_price_limit_x96, // Not used for multi-hop
        fee: _fee,
    } = swap_options;

    // Determine tokens for the path
    let (token0, token1, token2) =
        determine_two_hop_path_tokens(pool1_key, pool2_key, swap_params.is_to_b)?;

    debug!(
        token0 = ?token0.address(),
        token1 = ?token1.address(),
        token2 = ?token2.address(),
        "Determined two-hop swap path"
    );

    // Get deadline (current time + 20 minutes)
    let deadline = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .context("Failed to get current time")?
        .as_secs()
        + 1200; // 20 minutes

    // Determine which tick spacing to use for each hop
    let (tick_spacing0, tick_spacing1) = if swap_params.is_to_b {
        (pool1_key.tick_spacing, pool2_key.tick_spacing)
    } else {
        (pool2_key.tick_spacing, pool1_key.tick_spacing)
    };

    // Only ExactInput is supported for two-hop swaps
    let path = encode_two_hop_path(
        token0.address(),
        tick_spacing0,
        token1.address(),
        tick_spacing1,
        token2.address(),
    );

    // Convert amounts for ExactInput
    let amount_in = U256::from_big_int(swap_params.amount.quotient());
    let amount_out_minimum = amount_out_minimum
        .as_ref()
        .map(|amt| U256::from_big_int(amt.quotient()))
        .unwrap_or(U256::ZERO);

    // Calculate ETH value for native tokens
    let eth_value = if token0.is_native() { swap_params.amount.quotient() } else { BigInt::ZERO };

    info!(
        token0 = ?token0.address(),
        token1 = ?token1.address(),
        token2 = ?token2.address(),
        recipient = ?recipient,
        deadline = ?deadline,
        amount_in = ?amount_in,
        amount_out_minimum = ?amount_out_minimum,
        "Building two-hop swap with exactInput"
    );

    let swap_calldata = ISwapRouter::exactInputCall {
        params: ISwapRouter::ExactInputParams {
            path,
            recipient,
            deadline: U256::from(deadline),
            amountIn: amount_in,
            amountOutMinimum: amount_out_minimum,
        },
    }
    .abi_encode();

    let mut calldatas: Vec<Bytes> = Vec::with_capacity(1 + 3);

    // encode permit if necessary
    if let Some(input_token_permit) = input_token_permit {
        assert!(!token0.is_native(), "NON_TOKEN_PERMIT");
        calldatas.push(encode_permit(&token0, input_token_permit));
    }

    calldatas.push(swap_calldata.into());

    let calldata_count = calldatas.len();
    let final_calldata = encode_multicall(calldatas);
    info!(
        calldata_count = calldata_count,
        final_calldata_len = final_calldata.len(),
        value = ?U256::from_big_int(eth_value),
        "Swap call parameters finalized"
    );
    Ok(MethodParameters { calldata: final_calldata, value: U256::from_big_int(eth_value) })
}

/// Build quote call parameters for a two-hop swap
#[inline]
#[tracing::instrument(skip(pool1_key, pool2_key, swap_params, options), fields(token_a = ?pool1_key.token_a.address(), token_b = ?pool1_key.token_b.address(), token_c = ?pool2_key.token_b.address(), is_to_b = ?swap_params.is_to_b, trade_type = ?swap_params.trade_type, use_quoter_v2 = ?options.use_quoter_v2))]
pub fn build_two_hop_quote_call_parameters(
    pool1_key: &SlipstreamPoolKey,
    pool2_key: &SlipstreamPoolKey,
    swap_params: SwapParams,
    options: QuoteOptions,
) -> Result<MethodParameters> {
    let QuoteOptions {
        sqrt_price_limit_x96: _sqrt_price_limit_x96, // Not used for multi-hop quotes
        use_quoter_v2,
    } = options;

    let quote_amount = U256::from_big_int(swap_params.amount.quotient());

    // Determine tokens for the path
    let (token0, token1, token2) =
        determine_two_hop_path_tokens(pool1_key, pool2_key, swap_params.is_to_b)?;

    info!(
        token0 = ?token0.address(),
        token1 = ?token1.address(),
        token2 = ?token2.address(),
        "Determined two-hop path tokens"
    );

    info!(amount = ?quote_amount, "Quote amount");

    let (tick_spacing0, tick_spacing1) = if swap_params.is_to_b {
        (pool1_key.tick_spacing, pool2_key.tick_spacing)
    } else {
        (pool2_key.tick_spacing, pool1_key.tick_spacing)
    };

    debug!(
        token0 = ?token0.address(),
        token1 = ?token1.address(),
        token2 = ?token2.address(),
        is_to_b = ?swap_params.is_to_b,
        trade_type = ?swap_params.trade_type,
        use_quoter_v2 = ?use_quoter_v2,
        "Building two-hop quote call parameters with two calldatas"
    );

    let quote_calldata = match swap_params.trade_type {
        TradeType::ExactInput => {
            // Encode path: token0 + tickSpacing0 + token1 + tickSpacing1 + token2
            let path = encode_two_hop_path(
                token0.address(),
                tick_spacing0,
                token1.address(),
                tick_spacing1,
                token2.address(),
            );

            if use_quoter_v2 {
                IQuoterV2::quoteExactInputCall { path: path.clone(), amountIn: quote_amount }
                    .abi_encode()
            } else {
                IQuoter::quoteExactInputCall { path: path.clone(), amountIn: quote_amount }
                    .abi_encode()
            }
        }
        TradeType::ExactOutput => {
            // For ExactOutput, path must be in reverse order: token2 + tickSpacing1 +
            // token1 + tickSpacing0 + token0
            let path = encode_two_hop_path(
                token2.address(),
                tick_spacing1,
                token1.address(),
                tick_spacing0,
                token0.address(),
            );

            if use_quoter_v2 {
                IQuoterV2::quoteExactOutputCall { path: path.clone(), amountOut: quote_amount }
                    .abi_encode()
            } else {
                IQuoter::quoteExactOutputCall { path: path.clone(), amountOut: quote_amount }
                    .abi_encode()
            }
        }
    };

    let final_calldata: Bytes = quote_calldata.into();

    debug!(calldata_len = final_calldata.len(), "Two-hop quote call parameters finalized");

    Ok(MethodParameters { calldata: final_calldata, value: U256::ZERO })
}