waterpump-evm-pool-sdk 0.1.0

//! Ramses V3 Pool Infuser
//!
//! Uses shadow-lib bindings for Ramses V3-specific interfaces.
use alloy::{
    network::Ethereum,
    primitives::{aliases::I24, Address},
    providers::{DynProvider, Provider},
};
use anyhow::Result;
use tracing::{debug, instrument};
use uniswap_sdk_core::entities::BaseCurrency;
use waterpump_evm_shadow_client::interfaces::IRamsesV3PoolImmutables::IRamsesV3PoolImmutablesInstance;

use crate::{common::pool_utils::get_pool_key_v3, types::v3_pool_key::V3PoolKey};

/// Result of claiming rewards for a position
#[derive(Clone, Debug)]
pub struct ClaimRewardsResult {
    /// Position ID (NFP token ID) for which rewards were claimed
    pub position_id: alloy::primitives::U256,
    /// Reward amounts claimed, grouped by token
    /// Each entry represents a reward token and its claimed amount
    pub reward_amounts:
        Vec<uniswap_sdk_core::prelude::CurrencyAmount<uniswap_sdk_core::prelude::Currency>>,
    /// Transaction hash
    pub tx_hash: alloy::primitives::TxHash,
}

#[derive(Clone, Debug)]
pub struct ShadowPoolInfuser {
    pub pool_key: V3PoolKey,
    pub pool_address: Address,
    pub position_manager_address: Address,
    pub voter_address: Address,
    pub sender_address: Address,
    /// Tick spacing for the pool (required for Ramses V3 mint operations)
    pub tick_spacing: I24,
    pub provider: DynProvider<Ethereum>,
}

impl ShadowPoolInfuser {
    #[instrument(skip(pool_key), fields(
        pool_address = ?pool_address,
        position_manager_address = ?position_manager_address,
        voter_address = ?voter_address,
        token_a = ?pool_key.token_a.address(),
        token_b = ?pool_key.token_b.address(),
        fee = ?pool_key.fee,
        tick_spacing = ?tick_spacing,
        sender_address = ?sender_address
    ))]
    pub fn new(
        pool_key: V3PoolKey,
        pool_address: Address,
        position_manager_address: Address,
        voter_address: Address,
        sender_address: Address,
        tick_spacing: I24,
        provider: DynProvider<Ethereum>,
    ) -> Self {
        debug!(
            pool_address = ?pool_address,
            position_manager_address = ?position_manager_address,
            voter_address = ?voter_address,
            tick_spacing = ?tick_spacing,
            "Creating ShadowPoolInfuser"
        );
        Self {
            pool_key,
            pool_address,
            position_manager_address,
            voter_address,
            sender_address,
            tick_spacing,
            provider,
        }
    }

    pub async fn with_pool_address(
        provider: DynProvider<Ethereum>,
        chain_id: u64,
        pool_address: Address,
        position_manager_address: Address,
        voter_address: Address,
        sender_address: Address,
    ) -> Result<Self> {
        // Get pool key from contract (get_pool_key_v3 calls the pool contract to get
        // fee)
        let pool_key = get_pool_key_v3(&provider, pool_address, chain_id).await?;

        println!("ShadowPoolInfuser pool_key token_a: {:?}", pool_key.token_a.address());
        println!("ShadowPoolInfuser pool_key token_b: {:?}", pool_key.token_b.address());
        println!("ShadowPoolInfuser pool_key fee: {:?}", pool_key.fee);

        let contract = IRamsesV3PoolImmutablesInstance::new(pool_address, provider.clone());
        let tick_spacing = contract.tickSpacing().call().await?;
        Ok(Self {
            pool_key,
            pool_address,
            position_manager_address,
            voter_address,
            sender_address,
            tick_spacing,
            provider,
        })
    }

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

    /// Get the pool address
    pub fn pool_address(&self) -> Address { self.pool_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 pool key
    pub fn pool_key(&self) -> &V3PoolKey { &self.pool_key }

    /// Get the tick spacing
    pub fn tick_spacing(&self) -> I24 { self.tick_spacing }

    /// Claim rewards for one or more position IDs
    ///
    /// This method:
    /// 1. Builds a claim rewards transaction using shadow-lib
    /// 2. Sends the transaction to the voter contract
    /// 3. Waits for confirmation
    /// 4. Returns the reward amounts grouped by position ID
    ///
    /// # Arguments
    /// * `position_ids` - Slice of position IDs (NFP token IDs) to claim
    ///   rewards for
    ///
    /// # Returns
    /// A vector of `ClaimRewardsResult`, one for each position ID that had
    /// rewards claimed
    #[instrument(skip(self), fields(
        voter_address = ?self.voter_address(),
        position_count = position_ids.len()
    ))]
    pub async fn claim_rewards(
        &self,
        position_ids: &[alloy::primitives::U256],
    ) -> Result<Vec<ClaimRewardsResult>> {
        use std::collections::{HashMap, HashSet};

        use alloy::primitives::{Address, U256};
        use anyhow::Context;
        use uniswap_sdk_core::prelude::*;
        use waterpump_evm_shadow_client::transactions::claim_rewards::claim_rewards;

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

        // Build claim rewards transaction
        let claim_tx = claim_rewards(
            &self.provider,
            self.voter_address(),
            self.position_manager_address(),
            self.sender_address(),
            position_ids,
        )
        .await
        .context("Failed to build claim rewards transaction")?;

        // Convert TransactionInput to Bytes
        let calldata_bytes: alloy::primitives::Bytes = claim_tx
            .request
            .input
            .into_input()
            .ok_or(anyhow::anyhow!("Failed to convert TransactionInput to Bytes"))?;
        let calldata_len = calldata_bytes.len();

        tracing::debug!(calldata_len = calldata_len, "Claim rewards transaction built");

        let call_parameters = crate::pool_swappers::common::MethodParameters {
            calldata: calldata_bytes,
            value: claim_tx.request.value.unwrap_or_default(),
        };

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

        let receipt = crate::pool_swappers::common::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?;

        tracing::info!(
            tx_hash = ?receipt.transaction_hash,
            block_number = ?receipt.block_number,
            gas_used = ?receipt.gas_used,
            status = ?receipt.status(),
            "Claim rewards transaction confirmed"
        );

        if !receipt.status() {
            tracing::error!(
                tx_hash = ?receipt.transaction_hash,
                block_number = ?receipt.block_number,
                "Claim rewards transaction failed"
            );
            return Err(anyhow::anyhow!("Claim rewards transaction failed"));
        }

        // Extract reward amounts from quotes, grouped by position ID
        let mut results = Vec::new();

        if !claim_tx.quote.is_empty() {
            // Filter quotes for non-zero amounts
            let relevant_quotes: Vec<_> = claim_tx
                .quote
                .iter()
                .filter(|q| position_ids.contains(&q.token_id) && q.amount > U256::ZERO)
                .collect();

            if !relevant_quotes.is_empty() {
                // Get unique reward token addresses across all positions
                let reward_token_addresses: Vec<Address> = relevant_quotes
                    .iter()
                    .map(|q| q.token)
                    .collect::<HashSet<_>>()
                    .into_iter()
                    .collect();

                // Get chain_id from provider
                let chain_id_raw = self
                    .provider
                    .get_chain_id()
                    .await
                    .context("Failed to get chain_id from provider")?;
                let chain_id: u64 = chain_id_raw;

                // Get Currency objects for reward tokens
                let reward_currencies = waterpump_evm_uniswap_v3_client::get_currencies(
                    &self.provider,
                    &reward_token_addresses,
                    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());
                }

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

                // Convert to ClaimRewardsResult for each position
                for (position_id, reward_map) in position_rewards {
                    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() {
                        results.push(ClaimRewardsResult {
                            position_id,
                            reward_amounts: reward_amounts_vec,
                            tx_hash: receipt.transaction_hash,
                        });
                    }
                }
            }
        }

        Ok(results)
    }
}

// Import macros from common module
use crate::impl_token_helper;

impl_token_helper!(ShadowPoolInfuser);

use crate::{
    impl_pool_base, impl_shadow_pool_infuser, impl_shadow_pool_state, impl_shadow_pool_viewer,
};

impl_shadow_pool_viewer!(ShadowPoolInfuser);
impl_shadow_pool_state!(ShadowPoolInfuser);
impl_shadow_pool_infuser!(ShadowPoolInfuser);

impl_pool_base!(ShadowPoolInfuser);