polymarket 0.1.0

//! Main OnchainClient for Polymarket on-chain operations

use alloy::network::{EthereumWallet, AnyNetwork};
use alloy::primitives::{Address, U256};
use alloy::providers::{ProviderBuilder, Provider as _};
use std::sync::Arc;
use tokio::sync::OnceCell;

use crate::onchain::{
    contracts::{get_addresses, Addresses, IERC20},
    NetworkConfig, OnchainError, Result, WalletType,
};

use super::provider::OnchainProvider;
use super::signer::OnchainSigner;

/// Main client for on-chain operations on Polygon
///
/// This client provides a unified interface for:
/// - Token operations (ERC20, ERC1155)
/// - CTF operations (split, merge, redeem)
/// - Proxy wallet interactions
/// - Safe wallet interactions
///
/// # Example
/// ```no_run
/// # use polymarket_sdk::onchain::OnchainClientBuilder;
/// # #[tokio::main]
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// // Create client with private key
/// let client = OnchainClientBuilder::new()
///     .mainnet()
///     .private_key("0x1234...")?
///     .build()?;
///
/// // Get USDC balance
/// let balance = client.get_usdc_balance().await?;
/// println!("USDC Balance: {}", balance);
/// # Ok(())
/// # }
/// ```
#[derive(Clone)]
pub struct OnchainClient {
    provider: Arc<OnchainProvider>,
    signer: Arc<OnchainSigner>,
    addresses: Addresses,
    /// Provider with wallet and nonce management for sending transactions
    /// Lazily initialized on first use to avoid issues during client construction
    wallet_provider: Arc<OnceCell<alloy::providers::DynProvider<AnyNetwork>>>,
}

impl OnchainClient {
    /// Create a new OnchainClient
    ///
    /// Most users should use `OnchainClientBuilder` instead.
    pub fn new(provider: OnchainProvider, signer: OnchainSigner) -> Self {
        let addresses = get_addresses(provider.chain_id())
            .expect("Provider should have valid chain ID");

        Self {
            provider: Arc::new(provider),
            signer: Arc::new(signer),
            addresses,
            wallet_provider: Arc::new(OnceCell::new()),
        }
    }

    /// Get or create the wallet provider with nonce management
    /// This provider is created once and reused for all transactions
    pub(crate) async fn wallet_provider(&self) -> Result<&alloy::providers::DynProvider<AnyNetwork>> {
        self.wallet_provider
            .get_or_try_init(|| async {
                let wallet = self.ethereum_wallet();

                // Build the provider  with nonce management FIRST, then wallet
                // The order matters: network -> fillers -> transport
                let url = self.network().rpc_url.parse().map_err(|e| {
                    OnchainError::NetworkError(format!("Invalid RPC URL: {}", e))
                })?;

                let provider = ProviderBuilder::new()
                    .network::<AnyNetwork>()
                    .with_cached_nonce_management()  // Alloy's built-in nonce management
                    .wallet(wallet)
                    .connect_http(url);

                // Type erase to DynProvider
                Ok::<_, OnchainError>(provider.erased())
            })
            .await
    }

    /// Get the provider
    pub fn provider(&self) -> &OnchainProvider {
        &self.provider
    }

    /// Get the signer
    pub fn signer(&self) -> &OnchainSigner {
        &self.signer
    }

    /// Get the wallet address
    ///
    /// For smart wallets (Proxy/Safe), this returns the smart wallet address.
    /// For EOA, this returns the EOA address.
    pub fn address(&self) -> Address {
        self.signer.address()
    }

    /// Get the signer's EOA address
    ///
    /// This always returns the underlying EOA, even for smart wallets.
    pub fn signer_address(&self) -> Address {
        self.signer.signer_address()
    }

    /// Get the wallet type
    pub fn wallet_type(&self) -> WalletType {
        self.signer.wallet_type()
    }

    /// Get the network configuration
    pub fn network(&self) -> &NetworkConfig {
        self.provider.network()
    }

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

    /// Get contract addresses for this network
    pub fn addresses(&self) -> &Addresses {
        &self.addresses
    }

    /// Get USDC contract address
    ///
    /// Returns USDC.e address (bridged USDC) on mainnet since that's what Polymarket uses.
    /// Returns test USDC on testnet.
    pub fn usdc_address(&self) -> Address {
        self.addresses.usdc
    }

    /// Get CTF contract address
    pub fn ctf_address(&self) -> Address {
        self.addresses.ctf
    }

    /// Get Proxy Factory contract address
    pub fn proxy_factory_address(&self) -> Address {
        self.addresses.proxy_factory
    }

    /// Get Safe Factory contract address
    pub fn safe_factory_address(&self) -> Address {
        self.addresses.safe_factory
    }

    /// Get the current gas price
    pub async fn get_gas_price(&self) -> Result<U256> {
        self.provider.get_gas_price().await
    }

    /// Get the latest block number
    pub async fn get_block_number(&self) -> Result<u64> {
        self.provider.get_block_number().await
    }

    /// Check if the client is connected to the network
    pub async fn is_connected(&self) -> bool {
        self.provider.is_connected().await
    }

    /// Get USDC balance for the wallet
    pub async fn get_usdc_balance(&self) -> Result<U256> {
        self.get_usdc_balance_of(self.address()).await
    }

    /// Get USDC balance for a specific address
    ///
    /// Note: Polymarket uses USDC.e (bridged USDC) exclusively on mainnet.
    ///
    /// Background: The original "USDC" on Polygon was bridged from Ethereum.
    /// When Circle launched native USDC on Polygon, the original was renamed to USDC.e.
    /// Polymarket still uses USDC.e for all operations.
    pub async fn get_usdc_balance_of(&self, address: Address) -> Result<U256> {
        // Polymarket uses USDC.e exclusively (the original bridged USDC)
        let usdc_address = self.addresses.usdc;
        let usdc = IERC20::new(usdc_address, self.provider.provider());

        let balance_result = usdc
            .balanceOf(address)
            .call()
            .await
            .map_err(|e| OnchainError::ContractError(format!("Failed to get USDC balance: {}", e)))?;

        // In Alloy 1.0, balanceOf returns the U256 value directly
        Ok(balance_result)
    }

    /// Sign a message with the underlying EOA
    pub async fn sign_message(&self, message: &[u8]) -> Result<Vec<u8>> {
        self.signer.sign_message(message).await
    }

    /// Create an EthereumWallet for contract interactions
    pub(crate) fn ethereum_wallet(&self) -> EthereumWallet {
        self.signer.to_ethereum_wallet(self.chain_id())
    }

    /// Check if this is a smart contract wallet
    pub fn is_smart_wallet(&self) -> bool {
        self.signer.is_smart_wallet()
    }

    /// Get proxy address (if proxy wallet)
    pub fn proxy_address(&self) -> Option<Address> {
        self.signer.proxy_address()
    }

    /// Get Safe address (if Safe wallet)
    pub fn safe_address(&self) -> Option<Address> {
        self.signer.safe_address()
    }
}

impl std::fmt::Debug for OnchainClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("OnchainClient")
            .field("chain_id", &self.chain_id())
            .field("wallet_type", &self.wallet_type())
            .field("address", &self.address())
            .field("signer_address", &self.signer_address())
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloy::network::AnyNetwork;
    use alloy::providers::{ProviderBuilder, Provider};

    const TEST_PRIVATE_KEY: &str =
        "0x0123456789abcdef0123456789abcdef0123456789abcdef0123456789abcdef";

    #[test]
    fn test_client_new() {
        let network = NetworkConfig::polygon_mainnet();
        let provider = OnchainProvider {
            provider: {
                let url = network.rpc_url.parse().unwrap();
                ProviderBuilder::new()
                    .network::<AnyNetwork>()
                    .connect_http(url)
                    .erased()
            },
            network: network.clone(),
        };
        let signer = OnchainSigner::from_private_key(TEST_PRIVATE_KEY).unwrap();

        let client = OnchainClient::new(provider, signer);

        assert_eq!(client.chain_id(), 137);
        assert_eq!(client.wallet_type(), WalletType::PrivateKey);
        assert!(!client.is_smart_wallet());
    }

    #[test]
    fn test_client_addresses() {
        let network = NetworkConfig::polygon_mainnet();
        let provider = OnchainProvider {
            provider: {
                let url = network.rpc_url.parse().unwrap();
                ProviderBuilder::new()
                    .network::<AnyNetwork>()
                    .connect_http(url)
                    .erased()
            },
            network: network.clone(),
        };
        let signer = OnchainSigner::from_private_key(TEST_PRIVATE_KEY).unwrap();

        let client = OnchainClient::new(provider, signer);

        // Verify addresses are accessible
        assert_ne!(client.usdc_address(), Address::ZERO);
        assert_ne!(client.ctf_address(), Address::ZERO);
        assert_ne!(client.proxy_factory_address(), Address::ZERO);
        assert_ne!(client.safe_factory_address(), Address::ZERO);
    }

    #[test]
    fn test_client_proxy_wallet() {
        let network = NetworkConfig::polygon_mainnet();
        let provider = OnchainProvider {
            provider: {
                let url = network.rpc_url.parse().unwrap();
                ProviderBuilder::new()
                    .network::<AnyNetwork>()
                    .connect_http(url)
                    .erased()
            },
            network: network.clone(),
        };

        let proxy_address = Address::repeat_byte(0x42);
        let signer = OnchainSigner::from_proxy(TEST_PRIVATE_KEY, proxy_address).unwrap();

        let client = OnchainClient::new(provider, signer);

        assert_eq!(client.wallet_type(), WalletType::Proxy);
        assert!(client.is_smart_wallet());
        assert_eq!(client.address(), proxy_address);
        assert_eq!(client.proxy_address(), Some(proxy_address));
        assert_ne!(client.signer_address(), proxy_address);
    }

    #[test]
    fn test_client_safe_wallet() {
        let network = NetworkConfig::polygon_mainnet();
        let provider = OnchainProvider {
            provider: {
                let url = network.rpc_url.parse().unwrap();
                ProviderBuilder::new()
                    .network::<AnyNetwork>()
                    .connect_http(url)
                    .erased()
            },
            network: network.clone(),
        };

        let safe_address = Address::repeat_byte(0x99);
        let signer = OnchainSigner::from_safe(TEST_PRIVATE_KEY, safe_address, Some(2)).unwrap();

        let client = OnchainClient::new(provider, signer);

        assert_eq!(client.wallet_type(), WalletType::Safe);
        assert!(client.is_smart_wallet());
        assert_eq!(client.address(), safe_address);
        assert_eq!(client.safe_address(), Some(safe_address));
        assert_ne!(client.signer_address(), safe_address);
    }

    #[test]
    fn test_client_debug() {
        let network = NetworkConfig::polygon_mainnet();
        let provider = OnchainProvider {
            provider: {
                let url = network.rpc_url.parse().unwrap();
                ProviderBuilder::new()
                    .network::<AnyNetwork>()
                    .connect_http(url)
                    .erased()
            },
            network: network.clone(),
        };
        let signer = OnchainSigner::from_private_key(TEST_PRIVATE_KEY).unwrap();

        let client = OnchainClient::new(provider, signer);

        let debug_str = format!("{:?}", client);
        assert!(debug_str.contains("OnchainClient"));
        assert!(debug_str.contains("chain_id"));
        assert!(debug_str.contains("137"));
    }
}