waterpump-evm-pool-sdk 0.1.0

use std::time::Duration;

use alloy::{
    network::Ethereum,
    primitives::{
        aliases::{I24, U160},
        Address, Bytes, TxHash, I256, U256,
    },
    providers::{PendingTransactionConfig, Provider},
    rpc::types::{TransactionReceipt, TransactionRequest},
};
use anyhow::{Context, Result};
use tracing::{debug, error, info, trace};

/// Represents method call parameters including calldata and value.
#[derive(Debug, Clone)]
pub struct MethodParameters {
    pub calldata: Bytes,
    pub value: U256,
}

use super::{adjust_gas_prices, AdjustedGasPrices};
use crate::types::{
    swap_params::GasPriceOptions,
    swap_results::{V3SwapResult, V3SwapWithIntermediateResult},
};

/// Send a transaction, wait for confirmation, and retrieve the receipt.
///
/// This function:
/// 1. Sends the transaction to the network
/// 2. Waits for confirmation with an optional timeout
/// 3. Retrieves and returns the transaction receipt
///
/// # Arguments
///
/// * `provider` - The Ethereum provider to send transactions through
/// * `tx` - The transaction request to send
/// * `timeout` - Optional timeout for waiting for transaction confirmation
///   (defaults to 30 seconds)
/// * `error_handler` - Optional function to handle errors when sending the
///   transaction (must accept a displayable error)
///
/// # Returns
///
/// Returns the transaction receipt if successful, or an error if any step
/// fails.
///
/// # Example
///
/// ```ignore
/// let receipt = send_and_wait_for_transaction(
///     &provider,
///     tx,
///     Some(Duration::from_secs(30)),
///     Some(handle_uniswap_v3_error)
/// ).await?;
/// ```
#[tracing::instrument(skip(provider, tx, error_handler))]
pub async fn send_and_wait_for_transaction<P, F>(
    provider: &P,
    tx: TransactionRequest,
    timeout: Option<Duration>,
    error_handler: Option<F>,
) -> Result<TransactionReceipt>
where
    P: Provider<Ethereum> + ?Sized,
    F: FnOnce(Box<dyn std::fmt::Display + Send + Sync>) -> anyhow::Error,
{
    let pending_tx = if let Some(handler) = error_handler {
        provider.send_transaction(tx).await.map_err(|e| handler(Box::new(e)))?
    } else {
        provider
            .send_transaction(tx)
            .await
            .map_err(|e| anyhow::anyhow!("Failed to send transaction: {}", e))?
    };

    let tx_hash = pending_tx.tx_hash();
    info!(tx_hash = ?tx_hash, "Transaction sent, waiting for confirmation");

    let timeout_duration = timeout.unwrap_or(Duration::from_secs(120));
    let pending_tx_config =
        PendingTransactionConfig::new(*tx_hash).with_timeout(Some(timeout_duration));

    match provider.watch_pending_transaction(pending_tx_config).await?.await {
        Ok(tx_hash) => info!(tx_hash = ?tx_hash, "Transaction confirmed"),
        Err(e) => {
            error!(error = ?e, "Failed to watch pending transaction");
        }
    };

    info!(
        tx_hash = ?tx_hash,
        timeout_seconds = timeout_duration.as_secs(),
        "Transaction confirmed or timed out"
    );

    let error_message = format!("Failed to get transaction receipt: {}", tx_hash);
    let receipt = provider
        .get_transaction_receipt(*tx_hash)
        .await
        .context(error_message.clone())?
        .context(error_message)?;

    info!(
        tx_hash = ?receipt.transaction_hash,
        block_number = ?receipt.block_number,
        gas_used = ?receipt.gas_used,
        status = ?receipt.status(),
        "Transaction receipt retrieved"
    );

    Ok(receipt)
}

/// Build a transaction request with gas price configuration.
///
/// This function:
/// 1. Gets the current nonce for the sender address
/// 2. Adjusts gas prices based on provided options and network conditions
/// 3. Creates a transaction request with all parameters set
///
/// # Arguments
///
/// * `provider` - The Ethereum provider
/// * `sender_address` - The address sending the transaction
/// * `to_address` - The address to send the transaction to
/// * `method_params` - The method parameters containing calldata and value
/// * `gas_price_options` - Optional gas price configuration
///
/// # Returns
///
/// Returns a `TransactionRequest` with nonce and gas prices configured.
///
/// # Example
///
/// ```ignore
/// let tx = build_transaction_with_gas_prices(
///     &provider,
///     sender_address,
///     router_address,
///     method_params,
///     gas_price_options
/// ).await?;
/// ```
#[tracing::instrument(skip(provider, method_params, gas_price_options), fields(
    sender_address = ?sender_address,
    to_address = ?to_address,
    value = ?method_params.value
))]
pub async fn build_transaction_with_gas_prices<P: Provider<Ethereum> + ?Sized>(
    provider: &P,
    sender_address: Address,
    to_address: Address,
    method_params: MethodParameters,
    gas_price_options: Option<GasPriceOptions>,
) -> Result<TransactionRequest> {
    // Extract gas fee fields from gas_price_options parameter
    let provided_max_priority_fee_per_gas =
        gas_price_options.as_ref().and_then(|g| g.max_priority_fee_per_gas);
    let provided_max_fee_per_gas = gas_price_options.as_ref().and_then(|g| g.max_fee_per_gas);
    let gas_price_buffer_multiplier =
        gas_price_options.as_ref().and_then(|g| g.gas_price_buffer_multiplier);

    // Get the current nonce for the sender address
    let nonce = provider.get_transaction_count(sender_address).await?;
    info!(nonce = ?nonce, "Fetched transaction nonce");

    // Dynamically adjust gas prices based on current network conditions
    let AdjustedGasPrices { max_priority_fee_per_gas, max_fee_per_gas } =
        if let Some(buffer_multiplier) = gas_price_buffer_multiplier {
            adjust_gas_prices(
                provider,
                provided_max_priority_fee_per_gas,
                provided_max_fee_per_gas,
                buffer_multiplier,
            )
            .await?
        } else {
            AdjustedGasPrices {
                max_priority_fee_per_gas: provided_max_priority_fee_per_gas,
                max_fee_per_gas: provided_max_fee_per_gas,
            }
        };

    // Build the transaction request
    let mut tx = TransactionRequest::default()
        .from(sender_address)
        .to(to_address)
        .input(method_params.calldata.into())
        .value(method_params.value)
        .nonce(nonce);

    // Set priority gas fee (maxPriorityFeePerGas) if provided
    if let Some(max_priority_fee) = max_priority_fee_per_gas {
        // Convert U256 to u128 (gas fees fit in u128)
        let priority_fee_u128 = max_priority_fee.to::<u128>();
        tx = tx.max_priority_fee_per_gas(priority_fee_u128);
        debug!(max_priority_fee_per_gas = ?max_priority_fee, "Priority gas fee set");
    }

    // Set max fee per gas if provided
    if let Some(max_fee) = max_fee_per_gas {
        // Convert U256 to u128 (gas fees fit in u128)
        let max_fee_u128 = max_fee.to::<u128>();
        tx = tx.max_fee_per_gas(max_fee_u128);
        debug!(max_fee_per_gas = ?max_fee, "Max fee per gas set");
    }

    info!(value = ?method_params.value, "Transaction built with gas prices");
    Ok(tx)
}

/// Swap event data structure
pub struct SwapEventData {
    pub sender: Address,
    pub recipient: Address,
    pub amount0: I256,
    pub amount1: I256,
    pub sqrt_price_x96: U160,
    pub liquidity: u128,
    pub tick: I24,
}

/// Find and decode event from transaction receipt logs
///
/// This function:
/// 1. Iterates through all logs in the transaction receipt
/// 2. Attempts to decode each log using the provided decoder function
/// 3. Returns the result if found
///
/// # Arguments
///
/// * `receipt` - The transaction receipt containing logs
/// * `decode_fn` - A function that attempts to decode a log and returns the
///   result type
///
/// # Returns
///
/// Returns the result type `R` if an event is found, or an error if not found.
///
/// # Example
///
/// ```ignore
/// let swap_result = find_event(
///     &receipt,
///     |log, tx_hash| {
///         let decoded = log.log_decode::<Swap>()?;
///         let swap_data = SwapEventData {
///             sender: decoded.inner.sender,
///             recipient: decoded.inner.recipient,
///             amount0: decoded.inner.amount0,
///             amount1: decoded.inner.amount1,
///             sqrt_price_x96: decoded.inner.sqrtPriceX96,
///             liquidity: decoded.inner.liquidity,
///             tick: decoded.inner.tick,
///         };
///         let (amount_in, amount_out) = if is_to_b {
///             (swap_data.amount0.unsigned_abs(), swap_data.amount1.unsigned_abs())
///         } else {
///             (swap_data.amount1.unsigned_abs(), swap_data.amount0.unsigned_abs())
///         };
///         Ok(V3SwapResult {
///             tx_hash,
///             sender: swap_data.sender,
///             recipient: swap_data.recipient,
///             amount_in,
///             amount_out,
///             sqrt_price_x96: swap_data.sqrt_price_x96,
///             liquidity: U128::from(swap_data.liquidity),
///             tick: swap_data.tick,
///         })
///     }
/// )?;
/// ```
#[tracing::instrument(skip(receipt, decode_fn), fields(
    tx_hash = ?receipt.transaction_hash,
    log_count = receipt.logs().len()
))]
pub fn find_event<F, R>(receipt: &TransactionReceipt, decode_fn: F) -> Result<R>
where
    F: Fn(&alloy::rpc::types::Log, TxHash) -> Result<R>,
{
    let tx_hash = receipt.transaction_hash;
    info!(
        tx_hash = ?tx_hash,
        log_count = receipt.logs().len(),
        "Processing transaction logs"
    );

    for log in receipt.logs() {
        match decode_fn(log, tx_hash) {
            Ok(result) => {
                info!(
                    tx_hash = ?tx_hash,
                    "Event decoded successfully"
                );
                return Ok(result);
            }
            Err(e) => {
                trace!(error = ?e, "Log is not a matching event, skipping");
                continue;
            }
        }
    }

    error!(
        tx_hash = ?tx_hash,
        log_count = receipt.logs().len(),
        "No matching event found in transaction receipt"
    );
    Err(anyhow::anyhow!("No matching event found in transaction receipt"))
}

/// Find and decode events from transaction receipt logs using a decoder
///
/// This function applies a decoder function to all logs and collects all
/// successful results. For each log in the receipt, it attempts to decode using
/// the provided decoder function, collecting all successful decodings into an
/// array.
///
/// This function:
/// 1. Iterates through all logs in the transaction receipt
/// 2. For each log, attempts to decode using the decoder function
/// 3. Collects all successful results into a vector
/// 4. Returns the vector of all successfully decoded events
///
/// # Arguments
///
/// * `receipt` - The transaction receipt containing logs
/// * `decoder` - A decoder function that attempts to decode a log
///
/// # Returns
///
/// Returns a vector of result type `R` containing all successfully decoded
/// events. Returns an error if no events are found.
///
/// # Example
///
/// ```ignore
/// use waterpump_evm_pool_sdk::pool_swappers::v3::decoder::decode_swap_event_to_result;
///
/// let swap_results = find_events(
///     &receipt,
///     |log, tx_hash| decode_swap_event_to_result(log, tx_hash, is_to_b)
/// )?;
/// // swap_results is a Vec<R> containing all successfully decoded events
/// ```
#[tracing::instrument(skip(receipt, decoder), fields(
    tx_hash = ?receipt.transaction_hash,
    log_count = receipt.logs().len()
))]
pub fn find_events<F, R>(receipt: &TransactionReceipt, decoder: F) -> Result<Vec<R>>
where
    F: Fn(&alloy::rpc::types::Log, TxHash) -> Result<R>,
{
    let tx_hash = receipt.transaction_hash;
    info!(
        tx_hash = ?tx_hash,
        log_count = receipt.logs().len(),
        "Processing transaction logs with decoder"
    );

    let mut results = Vec::new();

    for log in receipt.logs() {
        match decoder(log, tx_hash) {
            Ok(result) => {
                info!(
                    tx_hash = ?tx_hash,
                    "Event decoded successfully"
                );
                results.push(result);
            }
            Err(e) => {
                trace!(
                    error = ?e,
                    "Failed to decode log, skipping"
                );
                continue;
            }
        }
    }

    if results.is_empty() {
        error!(
            tx_hash = ?tx_hash,
            log_count = receipt.logs().len(),
            "No matching event found in transaction receipt"
        );
        Err(anyhow::anyhow!("No matching event found in transaction receipt"))
    } else {
        info!(
            tx_hash = ?tx_hash,
            result_count = results.len(),
            "Found {} events using decoder",
            results.len()
        );
        Ok(results)
    }
}

/// Merge two swap results from a two-hop swap into a single result
///
/// This function combines two swap results representing a two-hop swap:
/// - First swap: input token -> intermediate token (is_to_intermediate = true)
/// - Second swap: intermediate token -> output token (is_to_intermediate =
///   false)
///
/// The merged result represents the complete two-hop swap with:
/// - Total input amount from the first swap
/// - Total output amount from the second swap
/// - Final pool state (sqrt_price_x96, liquidity, tick) from the second swap
///
/// # Arguments
///
/// * `swap_results` - A vector containing exactly 2 swap results with
///   intermediate flags
/// * `tx_hash` - The transaction hash for the merged result
///
/// # Returns
///
/// Returns a merged `V3SwapResult` representing the complete two-hop swap.
/// Returns an error if the vector doesn't contain exactly 2 swap results or if
/// the intermediate flags don't identify first and second hops correctly.
///
/// # Example
///
/// ```ignore
/// let swap_results = find_events(&receipt, &decoder)?;
/// let merged_result = merge_two_hop_swap_results(swap_results, receipt.transaction_hash)?;
/// ```
#[tracing::instrument(skip(swap_results), fields(
    tx_hash = ?tx_hash,
    swap_count = swap_results.len()
))]
pub fn merge_two_hop_swap_results(
    swap_results: Vec<V3SwapWithIntermediateResult>,
    tx_hash: TxHash,
) -> Result<V3SwapResult> {
    if swap_results.len() < 2 {
        error!(
            tx_hash = ?tx_hash,
            found_events = swap_results.len(),
            "Expected 2 swap events for two-hop swap, found {}",
            swap_results.len()
        );
        return Err(anyhow::anyhow!(
            "Expected 2 swap events for two-hop swap, found {}",
            swap_results.len()
        ));
    }

    // Identify first and second swaps using is_to_intermediate flag
    // First swap: is_to_intermediate = true (input token -> intermediate token)
    // Second swap: is_to_intermediate = false (intermediate token -> output token)
    let (first_swap, second_swap) = if swap_results[0].is_to_intermediate {
        if swap_results[1].is_to_intermediate {
            error!(
                tx_hash = ?tx_hash,
                "Both swap results have is_to_intermediate = true, expected one true and one false"
            );
            return Err(anyhow::anyhow!(
                "Both swap results have is_to_intermediate = true, expected one true and one false"
            ));
        }
        (&swap_results[0], &swap_results[1])
    } else {
        if !swap_results[1].is_to_intermediate {
            error!(
                tx_hash = ?tx_hash,
                "Both swap results have is_to_intermediate = false, expected one true and one false"
            );
            return Err(anyhow::anyhow!(
                "Both swap results have is_to_intermediate = false, expected one true and one \
                 false"
            ));
        }
        (&swap_results[1], &swap_results[0])
    };

    // The total amount_in is from the first swap
    // The total amount_out is from the second swap
    // Use the final state (sqrt_price_x96, liquidity, tick) from the second swap
    info!(
        tx_hash = ?tx_hash,
        first_amount_in = ?first_swap.amount_in,
        first_amount_out = ?first_swap.amount_out,
        second_amount_in = ?second_swap.amount_in,
        second_amount_out = ?second_swap.amount_out,
        first_is_to_intermediate = first_swap.is_to_intermediate,
        second_is_to_intermediate = second_swap.is_to_intermediate,
        "Merging two-hop swap results"
    );

    Ok(V3SwapResult {
        sender: first_swap.sender,
        recipient: second_swap.recipient,
        amount_in: first_swap.amount_in.clone(),
        amount_out: second_swap.amount_out.clone(),
        sqrt_price_x96: second_swap.sqrt_price_x96,
        liquidity: second_swap.liquidity,
        tick: second_swap.tick,
        tx_hash,
    })
}