oxidized-builder 0.1.0-delta

use crate::common::constants::LOW_BALANCE_THRESHOLD_WEI;
use crate::common::error::AppError;
use crate::common::retry::retry_async;
use crate::network::provider::HttpProvider;
use alloy::primitives::{Address, U256};
use alloy::providers::Provider;
use dashmap::DashMap;
use std::time::Duration;

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum BalanceTier {
    Emergency, // < 0.01 ETH
    Low,       // < 0.05 ETH
    Medium,   // < 1.0 ETH
    Healthy,   // Normal
    Whale,     // > 5.0 ETH
}

pub struct PortfolioManager {
    provider: HttpProvider,
    wallet_address: Address,
    token_balances: DashMap<(u64, Address), U256>,
    eth_balance: DashMap<u64, U256>, 
    total_profit_eth: DashMap<u64, f64>,
    total_gas_spent_eth: DashMap<u64, f64>,
    token_profit: DashMap<(u64, Address), f64>,
}

impl PortfolioManager {
    pub fn new(provider: HttpProvider, wallet_address: Address) -> Self {
        Self {
            provider,
            wallet_address,
            token_balances: DashMap::new(),
            eth_balance: DashMap::new(),
            total_profit_eth: DashMap::new(),
            total_gas_spent_eth: DashMap::new(),
            token_profit: DashMap::new(),
        }
    }

    pub async fn update_eth_balance(&self, chain_id: u64) -> Result<U256, AppError> {
        let provider = self.provider.clone();
        let addr = self.wallet_address;
        let bal = retry_async(
            move |_| {
                let provider = provider.clone();
                async move { provider.get_balance(addr).await }
            },
            3,
            Duration::from_millis(100),
        )
        .await
        .map_err(|e| AppError::Connection(format!("Balance check failed: {}", e)))?;

        let previous = self.eth_balance.insert(chain_id, bal);
        if let Some(prev) = previous {
            if bal > prev {
                let delta = bal - prev;
                self.record_profit(chain_id, wei_to_eth(delta), 0.0);
            }
        }
        Ok(bal)
    }

    pub fn get_tier(&self, chain_id: u64) -> BalanceTier {
        let bal = self
            .eth_balance
            .get(&chain_id)
            .map(|v| *v)
            .unwrap_or(U256::ZERO);
        
        let emergency = U256::from(10_000_000_000_000_000u64); // 0.01 ETH
        let whale = U256::from(5_000_000_000_000_000_000u64); // 5.0 ETH

        if bal < emergency {
            BalanceTier::Emergency
        } else if bal < *LOW_BALANCE_THRESHOLD_WEI {
            BalanceTier::Low
        } else if bal > whale {
            BalanceTier::Whale
        } else {
            BalanceTier::Healthy
        }
    }

    pub fn ensure_funding(&self, chain_id: u64, amount_needed: U256) -> Result<(), AppError> {
        let bal = self
            .eth_balance
            .get(&chain_id)
            .map(|v| *v)
            .unwrap_or(U256::ZERO);
        let gas_reserve = U256::from(5_000_000_000_000_000u64); // 0.005 ETH

        if bal < amount_needed + gas_reserve {
            return Err(AppError::InsufficientFunds {
                required: amount_needed.to_string(),
                available: bal.to_string(),
            });
        }
        Ok(())
    }

    pub async fn update_token_balance(
        &self,
        chain_id: u64,
        token: Address,
    ) -> Result<U256, AppError> {
        alloy::sol! {
            #[derive(Debug, PartialEq, Eq)]
            #[sol(rpc)]
            contract ERC20 {
                function balanceOf(address) external view returns (uint256);
            }
        }

        let contract = ERC20::new(token, self.provider.clone());
        let bal: U256 = retry_async(
            move |_| {
                let contract = contract.clone();
                async move { contract.balanceOf(self.wallet_address).call().await }
            },
            3,
            Duration::from_millis(100),
        )
        .await
        .map_err(|e| AppError::Connection(format!("Token balance failed: {}", e)))?;

        let previous = self.token_balances.insert((chain_id, token), bal);
        if let Some(prev) = previous {
            let delta = bal.saturating_sub(prev);
            let delta_f = wei_to_eth(delta); // assume 18 decimals; acceptable for most ERC20s
            self.record_token_profit(chain_id, token, delta_f);
        }
        Ok(bal)
    }

    pub fn get_token_balance(&self, chain_id: u64, token: Address) -> U256 {
        self.token_balances
            .get(&(chain_id, token))
            .map(|v| *v)
            .unwrap_or(U256::ZERO)
    }

    pub fn record_profit(&self, chain_id: u64, profit_eth: f64, gas_cost_eth: f64) {
        self.total_profit_eth
            .entry(chain_id)
            .and_modify(|v| *v += profit_eth)
            .or_insert(profit_eth);

        self.total_gas_spent_eth
            .entry(chain_id)
            .and_modify(|v| *v += gas_cost_eth)
            .or_insert(gas_cost_eth);
    }

    pub fn record_token_profit(&self, chain_id: u64, token: Address, profit_token: f64) {
        self.token_profit
            .entry((chain_id, token))
            .and_modify(|v| *v += profit_token)
            .or_insert(profit_token);
    }

    pub fn net_profit_eth(&self, chain_id: u64) -> f64 {
        let profit = self
            .total_profit_eth
            .get(&chain_id)
            .map(|v| *v)
            .unwrap_or(0.0);
        let gas = self
            .total_gas_spent_eth
            .get(&chain_id)
            .map(|v| *v)
            .unwrap_or(0.0);
        profit - gas
    }

    pub fn net_profit_all(&self) -> Vec<(u64, f64)> {
        self.total_profit_eth
            .iter()
            .map(|entry| {
                let chain = *entry.key();
                (chain, self.net_profit_eth(chain))
            })
            .collect()
    }

    pub fn token_profit_all(&self) -> Vec<(u64, Address, f64)> {
        self.token_profit
            .iter()
            .map(|entry| {
                let (chain, token) = *entry.key();
                (chain, token, *entry.value())
            })
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use url::Url;

    #[tokio::test]
    async fn records_profit_and_net() {
        let dummy_provider = HttpProvider::new_http(Url::parse("http://localhost:8545").unwrap());
        let pm = PortfolioManager::new(dummy_provider, Address::ZERO);
        pm.record_profit(1, 1.5, 0.4);
        pm.record_profit(1, 0.5, 0.1);
        assert!((pm.net_profit_eth(1) - 1.5).abs() < 1e-9);
    }

    #[tokio::test]
    async fn token_profit_tracks_delta() {
        let dummy_provider = HttpProvider::new_http(Url::parse("http://localhost:8545").unwrap());
        let pm = PortfolioManager::new(dummy_provider, Address::ZERO);
        let token = Address::from([1u8; 20]);
        // inject initial balance cache
        pm.token_balances
            .insert((1, token), U256::from(1_000_000_000_000_000_000u128));
        pm.record_token_profit(1, token, 0.5);
        assert_eq!(pm.token_profit_all().len(), 1);
    }
}

fn wei_to_eth(wei: U256) -> f64 {
    let base = U256::from(1_000_000_000_000_000_000u128);
    let num: u128 = wei.try_into().unwrap_or(u128::MAX);
    (num as f64) / (base.try_into().unwrap_or(1e18 as u128) as f64)
}