waterpump-evm-pool-sdk 0.1.0

use alloy::{
    network::Ethereum,
    primitives::{Address, U256},
    providers::DynProvider,
    rpc::types::TransactionReceipt,
};
use anyhow::{Context, Result};
use async_trait::async_trait;
use tracing::{debug, info, instrument};
use uniswap_sdk_core::prelude::{BigInt, CurrencyAmount, ToBig};

use crate::{
    pool_swappers::common::{
        build_transaction_with_gas_prices, find_events, send_and_wait_for_transaction,
        MethodParameters,
    },
    traits::pool_harvester::{
        HarvestFeesAndRewardsParams, HarvestFeesAndRewardsResult, HarvestPositionResult,
        HarvestQuoteData, PoolHarvester,
    },
};

/// V3 Pool Harvester implementation for Uniswap V3
#[derive(Clone)]
pub struct V3PoolHarvester {
    pub position_manager_address: Address,
    pub sender_address: Address,
    pub chain_id: u64,
    pub provider: DynProvider<Ethereum>,
}

impl V3PoolHarvester {
    /// Create a new V3 pool harvester
    pub fn new(
        position_manager_address: Address,
        sender_address: Address,
        chain_id: u64,
        provider: DynProvider<Ethereum>,
    ) -> Self {
        Self { position_manager_address, sender_address, chain_id, provider }
    }

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

    /// Get the position manager address
    pub fn position_manager_address(&self) -> Address { self.position_manager_address }

    /// Get the chain ID
    pub fn chain_id(&self) -> u64 { self.chain_id }

    /// Decode batch collect results from transaction receipt
    async fn decode_batch_collect_results(
        &self,
        receipt: &TransactionReceipt,
        token_ids: &[U256],
    ) -> Result<Vec<HarvestPositionResult>> {
        use crate::pool_infusers::v3::decoder::INonfungiblePositionManagerEvents;

        // Get positions for all token_ids to determine token addresses
        let positions = waterpump_evm_uniswap_v3_client::get_positions(
            &self.provider,
            self.position_manager_address,
            token_ids,
            None,
        )
        .await
        .context("Failed to get positions for decoding")?;

        // Find all Collect events from position manager
        let collect_events: Vec<_> = find_events(receipt, |log, _| {
            if log.address() != self.position_manager_address {
                return Err(anyhow::anyhow!(
                    "Collect event from unexpected address: expected {:?}, got {:?}",
                    self.position_manager_address,
                    log.address()
                ));
            }

            log.log_decode::<INonfungiblePositionManagerEvents::Collect>()
                .map(|decoded| decoded.inner)
                .map_err(|e| anyhow::anyhow!("Not a Collect event: {:?}", e))
        })
        .context("Failed to find Collect events in transaction receipt")?;

        debug!(
            collect_events_count = collect_events.len(),
            token_ids_count = token_ids.len(),
            "Decoded Collect events"
        );

        // Build a map of token_id to event for efficient lookup
        use std::collections::HashMap;
        let event_map: HashMap<U256, _> =
            collect_events.iter().map(|event| (event.tokenId, event)).collect();

        // Match events to token IDs and build results
        let mut results = Vec::with_capacity(token_ids.len());
        for (token_id, position_data) in token_ids.iter().zip(positions.iter()) {
            // Get Currency objects from token addresses
            let currencies = waterpump_evm_uniswap_v3_client::get_currencies(
                &self.provider,
                &[position_data.token0, position_data.token1],
                self.chain_id,
            )
            .await
            .context("Failed to get currencies for decoding")?;

            let token0 = currencies.first().context("Expected token0 currency")?.clone();
            let token1 = currencies.get(1).context("Expected token1 currency")?.clone();

            if let Some(event) = event_map.get(token_id) {
                let amount0 = CurrencyAmount::from_raw_amount(
                    token0.clone(),
                    U256::from(event.amount0).to_big_int(),
                )
                .context("Failed to create CurrencyAmount for token0")?;
                let amount1 = CurrencyAmount::from_raw_amount(
                    token1.clone(),
                    U256::from(event.amount1).to_big_int(),
                )
                .context("Failed to create CurrencyAmount for token1")?;

                results.push(HarvestPositionResult {
                    token_id: *token_id,
                    amount0,
                    amount1,
                    reward_amounts: None,
                });
            } else {
                // If no event found for this token_id, create empty result
                results.push(HarvestPositionResult {
                    token_id: *token_id,
                    amount0: CurrencyAmount::from_raw_amount(token0, BigInt::from(0))
                        .context("Failed to create empty CurrencyAmount for token0")?,
                    amount1: CurrencyAmount::from_raw_amount(token1, BigInt::from(0))
                        .context("Failed to create empty CurrencyAmount for token1")?,
                    reward_amounts: None,
                });
            }
        }

        Ok(results)
    }
}

#[async_trait]
impl PoolHarvester for V3PoolHarvester {
    #[instrument(skip(self), fields(
        position_manager_address = ?self.position_manager_address,
        token_ids_count = token_ids.len()
    ))]
    async fn get_quote_data(&self, token_ids: &[U256]) -> Result<Vec<HarvestQuoteData>> {
        info!("Getting quote data for {} positions", token_ids.len());

        if token_ids.is_empty() {
            return Ok(Vec::new());
        }

        // Get positions for all token_ids
        let positions = waterpump_evm_uniswap_v3_client::get_positions(
            &self.provider,
            self.position_manager_address,
            token_ids,
            None,
        )
        .await
        .context("Failed to get positions")?;

        // Collect fees for all positions using collect_multiple (efficient multicall)
        let collect_results = waterpump_evm_uniswap_v3_client::collect_multiple(
            &self.provider,
            self.position_manager_address,
            token_ids,
            None,
        )
        .await
        .context("Failed to collect fees for positions")?;

        // Build a map of token_id to collect result for efficient lookup
        use std::collections::HashMap;
        let collect_map: HashMap<U256, &waterpump_evm_uniswap_v3_client::CollectResult> =
            collect_results.iter().map(|result| (result.token_id, result)).collect();

        let mut quote_data = Vec::with_capacity(token_ids.len());

        for (token_id, position_data) in token_ids.iter().zip(positions.iter()) {
            // Get Currency objects from token addresses
            let currencies = waterpump_evm_uniswap_v3_client::get_currencies(
                &self.provider,
                &[position_data.token0, position_data.token1],
                self.chain_id,
            )
            .await
            .context("Failed to get currencies")?;

            let token0 = currencies.first().context("Expected token0 currency")?.clone();
            let token1 = currencies.get(1).context("Expected token1 currency")?.clone();

            // Get collect result for this token_id
            let collect_result = collect_map
                .get(token_id)
                .context(format!("Missing collect result for token_id {}", token_id))?;

            // Convert to CurrencyAmount using tokens from position_data order
            let amount0 = CurrencyAmount::from_raw_amount(
                token0.clone(),
                collect_result.amount0.to_big_int(),
            )
            .context("Failed to create CurrencyAmount for token0")?;
            let amount1 = CurrencyAmount::from_raw_amount(
                token1.clone(),
                collect_result.amount1.to_big_int(),
            )
            .context("Failed to create CurrencyAmount for token1")?;

            quote_data.push(HarvestQuoteData {
                token_id: *token_id,
                amount0,
                amount1,
                reward_amounts: None,
            });
        }

        debug!(
            num_quotes = quote_data.len(),
            "Retrieved quote data for {} positions",
            quote_data.len()
        );

        Ok(quote_data)
    }

    #[instrument(skip(self), fields(
        position_manager_address = ?self.position_manager_address,
        token_ids_count = params.token_ids.len(),
        recipient = ?params.recipient
    ))]
    async fn harvest_fees_and_rewards(
        &self,
        params: HarvestFeesAndRewardsParams,
    ) -> Result<HarvestFeesAndRewardsResult> {
        info!("Harvesting fees and rewards for {} positions", params.token_ids.len());

        if params.token_ids.is_empty() {
            return Err(anyhow::anyhow!("No token IDs provided for harvest"));
        }

        // Build batch collect parameters using uniswap-v3-lib
        let uniswap_method_params =
            waterpump_evm_uniswap_v3_client::build_batch_collect_parameters(
                &params.token_ids,
                params.recipient,
            )?;

        // Convert to evm-pool-lib MethodParameters type
        let method_params = MethodParameters {
            calldata: uniswap_method_params.calldata,
            value: uniswap_method_params.value,
        };

        debug!(calldata_len = method_params.calldata.len(), "Built batch collect transaction");

        // Build and send transaction
        let tx = build_transaction_with_gas_prices(
            &self.provider,
            self.sender_address,
            self.position_manager_address,
            method_params,
            None::<crate::types::swap_params::GasPriceOptions>,
        )
        .await?;

        let receipt = send_and_wait_for_transaction(
            &self.provider,
            tx,
            Some(std::time::Duration::from_secs(60)),
            None::<fn(Box<dyn std::fmt::Display + Send + Sync>) -> anyhow::Error>,
        )
        .await?;

        info!(
            tx_hash = ?receipt.transaction_hash,
            block_number = ?receipt.block_number,
            gas_used = ?receipt.gas_used,
            status = ?receipt.status(),
            "Harvest transaction confirmed"
        );

        if !receipt.status() {
            return Err(anyhow::anyhow!("Harvest transaction failed"));
        }

        // Decode results
        let results = self.decode_batch_collect_results(&receipt, &params.token_ids).await?;

        Ok(HarvestFeesAndRewardsResult { results, tx_hash: receipt.transaction_hash, receipt })
    }
}