waterpump-evm-pool-sdk 0.1.0

use std::collections::HashSet;

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

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

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

impl ShadowPoolHarvester {
    /// Create a new Shadow pool harvester
    pub fn new(
        position_manager_address: Address,
        voter_address: Address,
        sender_address: Address,
        chain_id: u64,
        provider: DynProvider<Ethereum>,
    ) -> Self {
        Self { position_manager_address, voter_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 voter address
    pub fn voter_address(&self) -> Address { self.voter_address }

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

#[async_trait]
impl PoolHarvester for ShadowPoolHarvester {
    #[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 {} Shadow positions", token_ids.len());

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

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

        // Get reward quotes using shadow-lib
        use std::collections::{HashMap, HashSet};

        use waterpump_evm_shadow_client::common::reward;

        // Prepare gauge rewards items
        let items = reward::prepare_gauge_rewards_items(
            &self.provider,
            self.voter_address,
            &positions,
            token_ids,
        )
        .await
        .context("Failed to prepare gauge rewards items")?;

        // Quote rewards
        let reward_quotes = if !items.is_empty() {
            reward::quote_rewards(&self.provider, &items, None)
                .await
                .context("Failed to quote rewards")?
        } else {
            Vec::new()
        };

        // Build a map of (token_id, reward_token) -> amount for efficient lookup
        let mut reward_map: HashMap<(U256, Address), U256> = HashMap::new();
        for quote in &reward_quotes {
            let key = (quote.token_id, quote.token);
            *reward_map.entry(key).or_insert(U256::ZERO) += quote.amount;
        }

        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 fees from position data (tokensOwed0 and tokensOwed1)
            let fee_amount0 = CurrencyAmount::from_raw_amount(
                token0.clone(),
                BigInt::from(position_data.tokensOwed0),
            )
            .context("Failed to create CurrencyAmount for token0 fees")?;
            let fee_amount1 = CurrencyAmount::from_raw_amount(
                token1.clone(),
                BigInt::from(position_data.tokensOwed1),
            )
            .context("Failed to create CurrencyAmount for token1 fees")?;

            // Sum up rewards for this position (if any reward tokens match token0 or
            // token1) For now, we only include fees in the quote. Rewards are
            // separate tokens. The total harvestable amount is fees + rewards,
            // but rewards are in different tokens. We'll return fees as
            // amount0/amount1, and rewards would need to be handled separately.

            // Get reward amounts for this position from reward_quotes
            let position_rewards: Vec<_> = reward_quotes
                .iter()
                .filter(|q| q.token_id == *token_id && q.amount > U256::ZERO)
                .collect();

            let reward_amounts = if !position_rewards.is_empty() {
                // Get unique reward token addresses for this position
                let reward_token_addresses: Vec<Address> = position_rewards
                    .iter()
                    .map(|q| q.token)
                    .collect::<HashSet<_>>()
                    .into_iter()
                    .collect();

                // Get Currency objects for reward tokens
                let reward_currencies = waterpump_evm_uniswap_v3_client::get_currencies(
                    &self.provider,
                    &reward_token_addresses,
                    self.chain_id,
                )
                .await
                .context("Failed to get currencies for reward tokens")?;

                // Create a map from token address to Currency
                let mut token_to_currency: HashMap<Address, Currency> = HashMap::new();
                for (addr, currency) in reward_token_addresses.iter().zip(reward_currencies.iter())
                {
                    token_to_currency.insert(*addr, currency.clone());
                }

                // Convert quotes to CurrencyAmount, grouping by token
                let mut reward_map: HashMap<Address, U256> = HashMap::new();
                for quote in position_rewards {
                    *reward_map.entry(quote.token).or_insert(U256::ZERO) += quote.amount;
                }

                // Convert to Vec<CurrencyAmount>
                let mut reward_amounts_vec = Vec::new();
                for (token_addr, total_amount) in reward_map {
                    if let Some(currency) = token_to_currency.get(&token_addr) {
                        if let Ok(amount) = CurrencyAmount::from_raw_amount(
                            currency.clone(),
                            total_amount.to_big_int(),
                        ) {
                            reward_amounts_vec.push(amount);
                        }
                    }
                }

                if reward_amounts_vec.is_empty() {
                    None
                } else {
                    Some(reward_amounts_vec)
                }
            } else {
                None
            };

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

        debug!(
            num_quotes = quote_data.len(),
            "Retrieved quote data for {} Shadow 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 {} Shadow positions", params.token_ids.len());

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

        // For Shadow, we need to:
        // 1. Collect fees using position manager's collect function (batch via
        //    multicall)
        // 2. Claim rewards using voter contract's claimRewards function
        // We'll do both in separate transactions

        // First, collect fees for all positions using multicall
        use waterpump_evm_shadow_client::transactions::claim_rewards::claim_rewards;

        // Build claim rewards transaction using shadow-lib
        let claim_tx = claim_rewards(
            &self.provider,
            self.voter_address,
            self.position_manager_address,
            self.sender_address,
            &params.token_ids,
        )
        .await
        .context("Failed to build claim rewards transaction")?;

        // Convert claim rewards transaction input to bytes
        let claim_rewards_calldata: alloy::primitives::Bytes = claim_tx
            .request
            .input
            .into_input()
            .ok_or(anyhow::anyhow!("Failed to convert TransactionInput to Bytes"))?;

        // Now send reward claiming transaction
        let reward_method_params = MethodParameters {
            calldata: claim_rewards_calldata,
            value: claim_tx.request.value.unwrap_or_default(),
        };

        let reward_tx = build_transaction_with_gas_prices(
            &self.provider,
            self.sender_address,
            self.voter_address,
            reward_method_params,
            None::<crate::types::swap_params::GasPriceOptions>,
        )
        .await?;

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

        info!(
            tx_hash = ?reward_receipt.transaction_hash,
            block_number = ?reward_receipt.block_number,
            gas_used = ?reward_receipt.gas_used,
            status = ?reward_receipt.status(),
            "Reward claiming transaction confirmed"
        );

        if !reward_receipt.status() {
            return Err(anyhow::anyhow!("Reward claiming transaction failed"));
        }

        // Build results from claim_tx quotes
        let positions = waterpump_evm_shadow_client::common::position::get_positions(
            &self.provider,
            self.position_manager_address,
            &params.token_ids,
            None,
        )
        .await
        .context("Failed to get positions for results")?;

        // Extract reward amounts from claim_tx.quote, grouped by position ID
        use std::collections::HashMap;

        // Group quotes by position ID and token address
        let mut position_rewards: HashMap<U256, HashMap<Address, U256>> = HashMap::new();
        for quote in &claim_tx.quote {
            if params.token_ids.contains(&quote.token_id) && quote.amount > U256::ZERO {
                let position_map = position_rewards.entry(quote.token_id).or_default();
                *position_map.entry(quote.token).or_insert(U256::ZERO) += quote.amount;
            }
        }

        // Get unique reward token addresses across all positions
        let all_reward_tokens: Vec<Address> = position_rewards
            .values()
            .flat_map(|reward_map| reward_map.keys())
            .copied()
            .collect::<HashSet<_>>()
            .into_iter()
            .collect();

        // Get Currency objects for all reward tokens (batch fetch)
        let reward_currencies = if !all_reward_tokens.is_empty() {
            waterpump_evm_uniswap_v3_client::get_currencies(
                &self.provider,
                &all_reward_tokens,
                self.chain_id,
            )
            .await
            .context("Failed to get currencies for reward tokens")?
        } else {
            Vec::new()
        };

        // Create a map from token address to Currency
        let mut token_to_currency: HashMap<Address, Currency> = HashMap::new();
        for (addr, currency) in all_reward_tokens.iter().zip(reward_currencies.iter()) {
            token_to_currency.insert(*addr, currency.clone());
        }

        // Build results for each position
        let mut results = Vec::new();
        for (token_id, position_data) in params.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 results")?;

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

            // Get reward amounts for this position from claim_tx
            let reward_amounts = if let Some(reward_map) = position_rewards.get(token_id) {
                // Convert to Vec<CurrencyAmount>
                let mut reward_amounts_vec = Vec::new();
                for (token_addr, total_amount) in reward_map {
                    if let Some(currency) = token_to_currency.get(token_addr) {
                        if let Ok(amount) = CurrencyAmount::from_raw_amount(
                            currency.clone(),
                            total_amount.to_big_int(),
                        ) {
                            reward_amounts_vec.push(amount);
                        }
                    }
                }

                if reward_amounts_vec.is_empty() {
                    None
                } else {
                    Some(reward_amounts_vec)
                }
            } else {
                None
            };

            // Create result with zero fees (since we only claimed rewards)
            results.push(HarvestPositionResult {
                token_id: *token_id,
                amount0: CurrencyAmount::from_raw_amount(token0, BigInt::from(0))
                    .context("Failed to create CurrencyAmount for token0")?,
                amount1: CurrencyAmount::from_raw_amount(token1, BigInt::from(0))
                    .context("Failed to create CurrencyAmount for token1")?,
                reward_amounts,
            });
        }

        debug!(
            reward_tx_hash = ?reward_receipt.transaction_hash,
            "Rewards claimed for {} positions",
            params.token_ids.len()
        );

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