waterpump-evm-pool-sdk 0.1.0

use std::time::Duration;

use alloy::{
    network::Ethereum,
    primitives::{Address, U256},
    providers::{DynProvider, Provider},
    rpc::types::TransactionRequest,
};
use anyhow::Result;
use async_trait::async_trait;
use tracing::{debug, error, info, instrument};
use uniswap_sdk_core::entities::BaseCurrency;
use waterpump_evm_uniswap_v3_client;

use super::{
    build_two_hop_quote_call_parameters, build_two_hop_swap_call_parameters, decode_quote_amount,
    decode_swap_event_to_result_with_intermediate, handle_uniswap_v3_error,
};
// Import macros from common module
use crate::{
    impl_two_hop_swapper_methods, impl_v3_multi_pool_state, impl_v3_two_pool_viewer,
    pool_swappers::common::{
        build_transaction_with_gas_prices, find_events, merge_two_hop_swap_results,
        send_and_wait_for_transaction,
    },
    traits::{pool_base::PoolBase, pool_swapper::PoolSwapper},
    types::{
        swap_params::{GasPriceOptions, QuoteOptions, SwapOptions, SwapParams},
        swap_results::V3SwapResult,
        v3_pool_key::V3PoolKey,
    },
};

// Helper functions to convert between evm-pool-lib types and uniswap-v3-lib
// types
fn convert_swap_params(params: SwapParams) -> waterpump_evm_uniswap_v3_client::SwapParams {
    waterpump_evm_uniswap_v3_client::SwapParams {
        amount: params.amount,
        trade_type: params.trade_type,
        is_to_b: params.is_to_b,
    }
}

fn convert_swap_options(options: SwapOptions) -> waterpump_evm_uniswap_v3_client::SwapOptions {
    waterpump_evm_uniswap_v3_client::SwapOptions {
        amount_in_maximum: options.amount_in_maximum,
        amount_out_minimum: options.amount_out_minimum,
        recipient: options.recipient,
        input_token_permit: options.input_token_permit,
        sqrt_price_limit_x96: options.sqrt_price_limit_x96,
        fee: options.fee,
    }
}

fn convert_quote_options(options: QuoteOptions) -> waterpump_evm_uniswap_v3_client::QuoteOptions {
    waterpump_evm_uniswap_v3_client::QuoteOptions {
        sqrt_price_limit_x96: options.sqrt_price_limit_x96,
        use_quoter_v2: options.use_quoter_v2,
    }
}

#[derive(Clone, Debug)]
pub struct V3TwoHopSwapper {
    pub pool_a: (V3PoolKey, Address),
    pub pool_b: (V3PoolKey, Address),

    pub router_address: Address,
    pub quoter_address: Address,
    pub sender_address: Address,
    pub provider: DynProvider<Ethereum>,
}

impl V3TwoHopSwapper {
    #[instrument(skip(pool_a, pool_b), fields(pool_a_address = ?pool_a.1, pool_b_address = ?pool_b.1, token_a = ?pool_a.0.token_a.address(), token_b = ?pool_a.0.token_b.address(), token_c = ?pool_b.0.token_b.address(), router_address = ?router_address, quoter_address = ?quoter_address, sender_address = ?sender_address))]
    pub fn new(
        pool_a: (V3PoolKey, Address),
        pool_b: (V3PoolKey, Address),
        router_address: Address,
        quoter_address: Address,
        sender_address: Address,
        provider: DynProvider<Ethereum>,
    ) -> Self {
        debug!("Creating V3TwoHopSwapper");
        Self { pool_a, pool_b, router_address, quoter_address, sender_address, provider }
    }

    /// Get the sender address
    pub fn sender_address(&self) -> Address { self.sender_address }

    /// Get the router address
    pub fn router_address(&self) -> Address { self.router_address }

    /// Get the quoter address
    pub fn quoter_address(&self) -> Address { self.quoter_address }

    /// Get the first pool key and address
    pub fn pool_a(&self) -> &(V3PoolKey, Address) { &self.pool_a }

    /// Get the second pool key and address
    pub fn pool_b(&self) -> &(V3PoolKey, Address) { &self.pool_b }

    pub fn pools(&self) -> Vec<(V3PoolKey, Address)> {
        vec![self.pool_a.clone(), self.pool_b.clone()]
    }
}

impl PoolBase for V3TwoHopSwapper {
    fn currency0(&self) -> uniswap_sdk_core::prelude::Currency { self.pool_a().0.token_a.clone() }

    fn currency1(&self) -> uniswap_sdk_core::prelude::Currency { self.pool_b().0.token_b.clone() }
}

impl_two_hop_swapper_methods!(V3TwoHopSwapper);

impl_v3_two_pool_viewer!(V3TwoHopSwapper);
impl_v3_multi_pool_state!(V3TwoHopSwapper);

use crate::traits::two_hop_swapper::TwoHopSwapper;
impl TwoHopSwapper for V3TwoHopSwapper {
    fn intermediate_token(&self) -> anyhow::Result<uniswap_sdk_core::prelude::Currency> {
        self.intermediate_token()
    }

    fn determine_two_hop_path_tokens(
        &self,
        is_to_b: bool,
    ) -> anyhow::Result<(
        uniswap_sdk_core::prelude::Currency,
        uniswap_sdk_core::prelude::Currency,
        uniswap_sdk_core::prelude::Currency,
    )> {
        self.determine_two_hop_path_tokens(is_to_b)
    }

    fn get_swap_directions(&self, is_to_b: bool) -> (bool, bool) {
        self.get_swap_directions(is_to_b)
    }
}

use crate::impl_token_helper;
impl_token_helper!(V3TwoHopSwapper);

#[async_trait]
impl PoolSwapper for V3TwoHopSwapper {
    #[instrument(skip(self, swap_params, swap_options, gas_price_options), fields(pool_a_address = ?self.pool_a.1, pool_b_address = ?self.pool_b.1, router_address = ?self.router_address(), is_to_b = ?swap_params.is_to_b, trade_type = ?swap_params.trade_type, amount = ?swap_params.amount))]
    async fn swap(
        &self,
        swap_params: SwapParams,
        swap_options: SwapOptions,
        gas_price_options: Option<GasPriceOptions>,
    ) -> Result<V3SwapResult> {
        // Extract is_to_b before swap_params is moved
        let overall_is_to_b = swap_params.is_to_b;

        let swap_call_parameters = build_two_hop_swap_call_parameters(
            &self.pool_a.0,
            &self.pool_b.0,
            convert_swap_params(swap_params),
            convert_swap_options(swap_options.clone()),
        )?;
        debug!(
            calldata_len = swap_call_parameters.calldata.len(),
            value = ?swap_call_parameters.value,
            "Two-hop swap call parameters built"
        );

        // Convert MethodParameters from uniswap-v3-lib to evm-pool-lib
        let method_params = crate::pool_swappers::common::MethodParameters {
            calldata: swap_call_parameters.calldata,
            value: swap_call_parameters.value,
        };

        let tx = build_transaction_with_gas_prices(
            &self.provider,
            self.sender_address(),
            self.router_address(),
            method_params,
            gas_price_options,
        )
        .await?;

        let receipt = send_and_wait_for_transaction(
            &self.provider,
            tx,
            Some(Duration::from_secs(120)),
            Some(|e| handle_uniswap_v3_error(format!("{}", e))),
        )
        .await?;

        let tx_hash = receipt.transaction_hash;
        info!(
            tx_hash = ?tx_hash,
            block_number = ?receipt.block_number,
            gas_used = ?receipt.gas_used,
            status = ?receipt.status(),
            "Two-hop transaction confirmed"
        );

        if !receipt.status() {
            error!(
                tx_hash = ?tx_hash,
                block_number = ?receipt.block_number,
                "Two-hop transaction failed"
            );
            return Err(anyhow::anyhow!("Two-hop transaction failed"));
        }

        // Determine token direction (A->B or B->A) for each pool
        // Note: This is the token direction within each pool, needed for decoding
        // events. The overall swap direction (towards pool_b) is given by
        // overall_is_to_b
        let (pool_a_token_is_to_b, pool_b_token_is_to_b) =
            self.get_swap_directions(overall_is_to_b);

        // Determine is_to_intermediate flags for each pool based on overall swap
        // direction
        let (pool_a_is_to_intermediate, pool_b_is_to_intermediate) =
            self.get_intermediate_flags(overall_is_to_b);

        debug!(
            pool_a_token_is_to_b = pool_a_token_is_to_b,
            pool_b_token_is_to_b = pool_b_token_is_to_b,
            pool_a_is_to_intermediate = pool_a_is_to_intermediate,
            pool_b_is_to_intermediate = pool_b_is_to_intermediate,
            overall_is_to_b = overall_is_to_b,
            "Determined token directions and intermediate flags for two-hop swap"
        );

        for log in receipt.logs().iter() {
            info!(
                log = ?log,
                "Two-hop transaction log"
            );
        }

        // Find events for both pools with the correct token direction and intermediate
        // flag We need to find exactly one event from pool_a and one from
        // pool_b Filter by pool address to ensure we get the correct event from
        // each pool
        let pool_a_address = self.pool_a.1;
        let pool_b_address = self.pool_b.1;
        let pool_a_currency0 = self.pool_a.0.token_a.clone();
        let pool_a_currency1 = self.pool_a.0.token_b.clone();
        let pool_b_currency0 = self.pool_b.0.token_a.clone();
        let pool_b_currency1 = self.pool_b.0.token_b.clone();
        let pool_a_token_is_to_b_captured = pool_a_token_is_to_b;
        let pool_a_is_to_intermediate_captured = pool_a_is_to_intermediate;
        let pool_b_token_is_to_b_captured = pool_b_token_is_to_b;
        let pool_b_is_to_intermediate_captured = pool_b_is_to_intermediate;
        let swap_results = find_events(&receipt, move |log, tx_hash| {
            // Check which pool this log is from and use the appropriate decoder
            if log.address() == pool_a_address {
                decode_swap_event_to_result_with_intermediate(
                    log,
                    tx_hash,
                    pool_a_token_is_to_b_captured,
                    pool_a_is_to_intermediate_captured,
                    pool_a_currency0.clone(),
                    pool_a_currency1.clone(),
                )
            } else if log.address() == pool_b_address {
                decode_swap_event_to_result_with_intermediate(
                    log,
                    tx_hash,
                    pool_b_token_is_to_b_captured,
                    pool_b_is_to_intermediate_captured,
                    pool_b_currency0.clone(),
                    pool_b_currency1.clone(),
                )
            } else {
                Err(anyhow::anyhow!("Log is not from pool_a or pool_b"))
            }
        })?;

        merge_two_hop_swap_results(swap_results, tx_hash)
    }

    #[instrument(skip(self, swap_params, swap_options), fields(pool_a_address = ?self.pool_a.1, pool_b_address = ?self.pool_b.1, router_address = ?self.router_address(), is_to_b = ?swap_params.is_to_b, trade_type = ?swap_params.trade_type, amount = ?swap_params.amount))]
    async fn simulate_swap(
        &self,
        swap_params: SwapParams,
        swap_options: SwapOptions,
    ) -> Result<()> {
        let swap_call_parameters = build_two_hop_swap_call_parameters(
            &self.pool_a.0,
            &self.pool_b.0,
            convert_swap_params(swap_params),
            convert_swap_options(swap_options.clone()),
        )?;
        debug!(
            calldata_len = swap_call_parameters.calldata.len(),
            value = ?swap_call_parameters.value,
            "Two-hop simulation call parameters built"
        );
        let tx = TransactionRequest::default()
            .from(self.sender_address())
            .to(self.router_address())
            .input(swap_call_parameters.calldata.into())
            .value(swap_call_parameters.value);

        let res = Provider::call(&self.provider, tx).await?;
        info!(result_len = res.len(), "Two-hop simulation completed successfully");
        debug!(simulation_result = ?res, "Two-hop simulation result details");

        Ok(())
    }

    #[instrument(skip(self, swap_params, options), fields(pool_a_address = ?self.pool_a.1, pool_b_address = ?self.pool_b.1, quoter_address = ?self.quoter_address, is_to_b = ?swap_params.is_to_b, trade_type = ?swap_params.trade_type, amount = ?swap_params.amount, use_quoter_v2 = ?options.use_quoter_v2))]
    async fn quote(&self, swap_params: SwapParams, options: QuoteOptions) -> Result<U256> {
        let use_quoter_v2 = options.use_quoter_v2;
        let swap_call_parameters = build_two_hop_quote_call_parameters(
            &self.pool_a.0,
            &self.pool_b.0,
            convert_swap_params(swap_params.clone()),
            convert_quote_options(options),
        )?;
        debug!(
            calldata_len = swap_call_parameters.calldata.len(),
            value = ?swap_call_parameters.value,
            "Two-hop quote call parameters built"
        );
        let tx = TransactionRequest::default()
            .from(self.sender_address())
            .to(self.quoter_address)
            .input(swap_call_parameters.calldata.into())
            .value(swap_call_parameters.value);
        let res = Provider::call(&self.provider, tx).await?;
        let amount = decode_quote_amount(res, &swap_params.trade_type, use_quoter_v2)?;
        info!(
            quoted_amount = ?amount,
            trade_type = ?swap_params.trade_type,
            "Two-hop quote completed successfully"
        );
        Ok(amount)
    }
}