use alloy::primitives::{Address, Uint};
use alloy::providers::{Provider, ProviderBuilder};
use alloy_chains::NamedChain;
use comfy_table::{presets::UTF8_BORDERS_ONLY, Table};
use eyre::Result;
use std::collections::HashMap;
use tracing::{info, warn};
use url::Url;

use crate::chain_client::{ChainClient, ARCH_SOLANA};
use crate::commands::config::config_pb::{Chain, Configuration, GetConfigResponse};
use crate::evm::rpc::{MidribV3, IERC20};
#[cfg(test)]
use crate::wallet::CurveType;
use crate::wallet::Wallet;

/// Represents a unique token across all chains
#[derive(Debug, Clone)]
struct TokenInfo {
    symbol: String,
    decimals: u32,
}

/// Balance information for a token on a specific chain
#[derive(Debug)]
struct ChainBalance {
    chain_network: String,
    wallet_balance: String,
    available_balance: String,
    locked_balance: String,
}

/// Native gas token balance for a chain
#[derive(Debug)]
struct NativeBalance {
    chain_network: String,
    balance: String,
}

/// Aggregated balance for a single token across all chains
#[derive(Debug)]
struct TokenBalance {
    token_info: TokenInfo,
    chain_balances: Vec<ChainBalance>,
}

/// Extract all unique tokens from configuration chains
fn extract_all_tokens_from_config(config: &Configuration) -> HashMap<String, TokenInfo> {
    let mut tokens: HashMap<String, TokenInfo> = HashMap::new();

    for chain in &config.chains {
        for (symbol, token) in &chain.tokens {
            tokens.entry(symbol.clone()).or_insert_with(|| TokenInfo {
                symbol: symbol.clone(),
                decimals: token.decimals,
            });
        }
    }

    tokens
}

/// Format balance with decimals for human-readable display
fn format_balance_with_decimals(balance_str: &str, decimals: u32) -> String {
    if balance_str == "error" || balance_str == "not deployed" {
        return balance_str.to_string();
    }

    match balance_str.parse::<u128>() {
        Ok(balance) => {
            let divisor = 10_u128.pow(decimals);
            let integer_part = balance / divisor;
            let fractional_part = balance % divisor;

            // Format with proper decimal places
            format!(
                "{}.{:0width$}",
                integer_part,
                fractional_part,
                width = decimals as usize
            )
        }
        Err(_) => balance_str.to_string(),
    }
}

/// Display all token balances in a single table grouped by chain
fn display_all_token_balances(
    all_token_balances: &[TokenBalance],
    native_balances: &[NativeBalance],
) -> String {
    if all_token_balances.is_empty() {
        return String::new();
    }

    // Collect all unique chain networks across all tokens and sort them
    let mut all_chains: Vec<String> = Vec::new();
    for token_balance in all_token_balances {
        for chain_balance in &token_balance.chain_balances {
            if !all_chains.contains(&chain_balance.chain_network) {
                all_chains.push(chain_balance.chain_network.clone());
            }
        }
    }
    all_chains.sort();

    let mut output = String::new();
    output.push('\n');

    output
        .push_str("═══════════════════════════════════════════════════════════════════════════\n");
    output.push_str("                                 BALANCES\n");
    output.push_str(
        "═══════════════════════════════════════════════════════════════════════════\n\n",
    );

    let mut table = Table::new();
    table.load_preset(UTF8_BORDERS_ONLY);

    // Set header: Token, Wallet, Deposited, Locked
    table.set_header(vec!["Token", "Wallet", "Deposited", "Locked"]);

    // Group by chain - add rows for each chain section
    for (chain_idx, chain) in all_chains.iter().enumerate() {
        // Add chain header row (spans conceptually as a section header)
        if chain_idx > 0 {
            // Add separator row between chains
            table.add_row(vec!["", "", "", ""]);
        }

        // Add chain name as a section header
        table.add_row(vec![
            format!("── {} ──", chain),
            String::new(),
            String::new(),
            String::new(),
        ]);

        // Add native gas balance for this chain
        if let Some(native) = native_balances.iter().find(|nb| nb.chain_network == *chain) {
            let gas_balance = format_balance_with_decimals(&native.balance, 18);
            table.add_row(vec![
                "GAS".to_string(),
                gas_balance,
                String::new(),
                String::new(),
            ]);
        }

        // Add tokens for this chain
        for token_balance in all_token_balances {
            // Find if this token exists on this chain
            if let Some(chain_balance) = token_balance
                .chain_balances
                .iter()
                .find(|cb| cb.chain_network == *chain)
            {
                let decimals = token_balance.token_info.decimals;
                let symbol = &token_balance.token_info.symbol;

                let wallet = format_balance_with_decimals(&chain_balance.wallet_balance, decimals);
                let deposited =
                    format_balance_with_decimals(&chain_balance.available_balance, decimals);
                let locked = format_balance_with_decimals(&chain_balance.locked_balance, decimals);

                table.add_row(vec![symbol.clone(), wallet, deposited, locked]);
            }
        }
    }

    output.push_str(&table.to_string());
    output.push('\n');

    output
}

/// Query a token balance on a specific chain using a curve-aware client.
///
/// Used by `balance_from_config_with_wallet`. For Solana chains, locked/deposited
/// balances come from the trade program; for Solana it queries the on-chain
/// `UserBalance` PDA via the Midrib program.
#[cfg(feature = "solana")]
async fn solana_user_balance(
    chain: &Chain,
    token: &crate::commands::config::config_pb::Token,
    owner_address: &str,
) -> (String, String) {
    use solana_sdk::pubkey::Pubkey;
    use std::str::FromStr;

    let (program_id, instance) = match crate::solana::client::resolve_program_and_instance(chain) {
        Ok(v) => v,
        Err(_) => return ("not deployed".to_string(), "not deployed".to_string()),
    };
    let user = match Pubkey::from_str(owner_address) {
        Ok(p) => p,
        Err(_) => return ("bad address".to_string(), "bad address".to_string()),
    };
    let mint = match Pubkey::from_str(&token.address) {
        Ok(p) => p,
        Err(_) => return ("bad mint".to_string(), "bad mint".to_string()),
    };
    match crate::solana::client::fetch_user_balance(
        &chain.rpc_url,
        &instance,
        &user,
        &mint,
        &program_id,
    )
    .await
    {
        Ok((deposited, locked)) => {
            let available = deposited.saturating_sub(locked);
            (available.to_string(), locked.to_string())
        }
        Err(e) => {
            warn!(
                "Solana UserBalance fetch failed on {}: {}",
                chain.network, e
            );
            ("error".to_string(), "error".to_string())
        }
    }
}

#[cfg(not(feature = "solana"))]
async fn solana_user_balance(
    _chain: &Chain,
    _token: &crate::commands::config::config_pb::Token,
    _owner_address: &str,
) -> (String, String) {
    (
        "solana feature disabled".to_string(),
        "solana feature disabled".to_string(),
    )
}

async fn query_token_balance_via_client(
    chain: &Chain,
    token_symbol: &str,
    owner_address: &str,
) -> ChainBalance {
    let chain_network = chain.network.clone();
    let token = match chain.tokens.get(token_symbol) {
        Some(t) => t,
        None => {
            return ChainBalance {
                chain_network,
                wallet_balance: "missing token".to_string(),
                available_balance: "missing token".to_string(),
                locked_balance: "missing token".to_string(),
            };
        }
    };

    let client = match ChainClient::from_chain_config(chain) {
        Ok(c) => c,
        Err(e) => {
            warn!("Failed to build client for {}: {}", chain_network, e);
            return ChainBalance {
                chain_network,
                wallet_balance: "error".to_string(),
                available_balance: "error".to_string(),
                locked_balance: "error".to_string(),
            };
        }
    };

    let wallet_balance = client
        .token_balance(token, owner_address)
        .await
        .map_or_else(
            |e| {
                warn!("Failed to get wallet balance on {}: {}", chain_network, e);
                "error".to_string()
            },
            |v| v.to_string(),
        );

    // Locked/deposited balances come from the trade contract.
    // Solana side is not yet wired up; EVM uses MidribV3.
    let (available_balance, locked_balance) =
        if chain.architecture.eq_ignore_ascii_case(ARCH_SOLANA) {
            solana_user_balance(chain, token, owner_address).await
        } else {
            let contract_address = chain
                .trade_contract
                .as_ref()
                .map(|tc| tc.address.clone())
                .unwrap_or_default();

            if contract_address.is_empty() {
                ("not deployed".to_string(), "not deployed".to_string())
            } else {
                let named_chain =
                    NamedChain::try_from(chain.chain_id as u64).unwrap_or(NamedChain::BaseSepolia);
                // Reuse existing EVM helpers — they need a privkey to derive the address,
                // but we already have the address. Use *_for_address variants.
                let owner: Address = match owner_address.parse() {
                    Ok(a) => a,
                    Err(_) => {
                        return ChainBalance {
                            chain_network,
                            wallet_balance,
                            available_balance: "bad address".to_string(),
                            locked_balance: "bad address".to_string(),
                        };
                    }
                };
                let available = call_get_balance_for_address(
                    named_chain,
                    &chain.rpc_url,
                    &token.address,
                    &contract_address,
                    owner,
                )
                .await
                .map_or_else(
                    |e| {
                        warn!(
                            "Failed to get available balance on {}: {}",
                            chain_network, e
                        );
                        "error".to_string()
                    },
                    |v| v.to_string(),
                );
                let locked = call_get_locked_balance_for_address(
                    &chain.rpc_url,
                    &token.address,
                    &contract_address,
                    owner,
                )
                .await
                .map_or_else(
                    |e| {
                        warn!("Failed to get locked balance on {}: {}", chain_network, e);
                        "error".to_string()
                    },
                    |v| v.to_string(),
                );
                (available, locked)
            }
        };

    ChainBalance {
        chain_network,
        wallet_balance,
        available_balance,
        locked_balance,
    }
}

/// Curve-agnostic config-driven balance function.
///
/// Dispatches per-chain based on `chain.architecture`:
/// - EVM chains query via Alloy + MidribV3
/// - Solana chains query via solana-client (SOL + SPL); deposited/locked
///   balances are scaffolded as "not deployed" until the on-chain program lands
pub async fn balance_from_config_with_wallet(
    config: GetConfigResponse,
    wallet: &Wallet,
) -> Result<()> {
    balance_from_config_with_wallets(config, &[wallet]).await
}

/// Pick the first wallet matching `chain`'s curve. Returns `None` when no
/// caller-supplied wallet matches — the chain's rows are reported as
/// "no wallet" rather than erroring out.
fn select_wallet_for_chain<'a>(chain: &Chain, wallets: &'a [&'a Wallet]) -> Option<&'a Wallet> {
    let wanted = crate::wallet::chain_curve(chain);
    wallets.iter().copied().find(|w| w.curve() == wanted)
}

/// Curve-aware, multi-wallet config-driven balance function.
///
/// Each chain in the config gets matched to a wallet with the same curve
/// (EVM chains → `Wallet::Evm`, Solana chains → `Wallet::Solana`). A chain
/// with no matching wallet is reported as "no wallet" rather than failing
/// the whole call.
pub async fn balance_from_config_with_wallets(
    config: GetConfigResponse,
    wallets: &[&Wallet],
) -> Result<()> {
    let configuration = config
        .config
        .ok_or_else(|| eyre::eyre!("No configuration found in response"))?;

    let tokens = extract_all_tokens_from_config(&configuration);

    if tokens.is_empty() {
        info!("No tokens found in configuration");
        return Ok(());
    }

    info!("Found {} unique token(s) across all chains", tokens.len());

    let mut all_token_balances: Vec<TokenBalance> = Vec::new();

    for (symbol, token_info) in tokens {
        let mut chain_balances = Vec::new();

        for chain in &configuration.chains {
            if !chain.tokens.contains_key(&symbol) {
                continue;
            }
            let cb = match select_wallet_for_chain(chain, wallets) {
                Some(w) => query_token_balance_via_client(chain, &symbol, &w.address()).await,
                None => ChainBalance {
                    chain_network: chain.network.clone(),
                    wallet_balance: "no wallet".to_string(),
                    available_balance: "no wallet".to_string(),
                    locked_balance: "no wallet".to_string(),
                },
            };
            chain_balances.push(cb);
        }

        all_token_balances.push(TokenBalance {
            token_info: token_info.clone(),
            chain_balances,
        });
    }

    all_token_balances.sort_by(|a, b| a.token_info.symbol.cmp(&b.token_info.symbol));

    // Native gas balances per chain
    let mut native_balances: Vec<NativeBalance> = Vec::new();
    let mut seen: HashMap<String, ()> = HashMap::new();
    for chain in &configuration.chains {
        if seen.contains_key(&chain.network) {
            continue;
        }
        seen.insert(chain.network.clone(), ());

        let owner_address = match select_wallet_for_chain(chain, wallets) {
            Some(w) => w.address(),
            None => {
                native_balances.push(NativeBalance {
                    chain_network: chain.network.clone(),
                    balance: "no wallet".to_string(),
                });
                continue;
            }
        };

        let client = match ChainClient::from_chain_config(chain) {
            Ok(c) => c,
            Err(e) => {
                warn!("Failed to build client for {}: {}", chain.network, e);
                native_balances.push(NativeBalance {
                    chain_network: chain.network.clone(),
                    balance: "error".to_string(),
                });
                continue;
            }
        };
        let balance = client.native_balance(&owner_address).await.map_or_else(
            |e| {
                warn!("Failed to get native balance on {}: {}", chain.network, e);
                "error".to_string()
            },
            |v| v.to_string(),
        );
        native_balances.push(NativeBalance {
            chain_network: chain.network.clone(),
            balance,
        });
    }

    let output = display_all_token_balances(&all_token_balances, &native_balances);
    info!("{}", output);

    Ok(())
}

/// Read the trader's available trade balance from MidribV3's
/// `tradeBalance(owner, token)` accessor.
pub async fn call_get_balance_for_address(
    chain: NamedChain,
    rpc_url: &str,
    token_address: &str,
    contract_address: &str,
    depositer_address: Address,
) -> Result<Uint<256, 4>> {
    let contract_addr: Address = contract_address.parse()?;
    let token_addr: Address = token_address.parse()?;
    let rpc_url = Url::parse(rpc_url)?;
    let provider = ProviderBuilder::new()
        .with_chain(chain)
        .connect_http(rpc_url);
    let contract = MidribV3::new(contract_addr, &provider);
    let result = contract
        .tradeBalance(depositer_address, token_addr)
        .call()
        .await?;
    Ok(result)
}

/// Variant of `call_get_locked_balance` that takes an `Address` directly.
///
/// MidribV3 has no on-chain locked balance: under the optimistic shadow
/// ledger, order reservations live off-chain in the TEE (the chain only sees
/// deposits, net settlement, and voucher withdrawals). There is nothing on
/// the contract to query, so the EVM locked balance is always zero.
pub async fn call_get_locked_balance_for_address(
    _rpc_url: &str,
    _token_address: &str,
    _contract_address: &str,
    _depositer_address: Address,
) -> Result<Uint<256, 4>> {
    Ok(Uint::<256, 4>::ZERO)
}

/// Read the native gas balance (wei) for an explicit `address`.
pub async fn call_get_native_balance_for_address(
    rpc_url: &str,
    address: Address,
) -> Result<Uint<256, 4>> {
    let rpc_url = Url::parse(rpc_url)?;
    let provider = ProviderBuilder::new().connect_http(rpc_url);
    let balance = provider.get_balance(address).await?;
    Ok(balance)
}

/// Read the ERC-20 balance held by `holder`, without any signer setup.
pub async fn call_get_erc20_balance_for_address(
    rpc_url: &str,
    token_address: &str,
    holder: Address,
) -> Result<Uint<256, 4>> {
    let token_addr: Address = token_address.parse()?;
    let rpc_url = Url::parse(rpc_url)?;
    let provider = ProviderBuilder::new().connect_http(rpc_url);
    let contract = IERC20::new(token_addr, &provider);
    let result = contract.balanceOf(holder).call().await?;
    Ok(result)
}

/// Format a u256 raw balance as a decimal string with `decimals` digits
/// after the point (e.g. `1500000` with `decimals = 6` → `"1.500000"`).
pub fn format_balance(value: Uint<256, 4>, decimals: u32) -> String {
    let s = value.to_string();
    if decimals == 0 {
        return s;
    }
    let dec = decimals as usize;
    if s.len() <= dec {
        let padded = format!("{:0>width$}", s, width = dec + 1);
        let (int_part, frac_part) = padded.split_at(padded.len() - dec);
        format!("{}.{}", int_part, frac_part)
    } else {
        let (int_part, frac_part) = s.split_at(s.len() - dec);
        format!("{}.{}", int_part, frac_part)
    }
}

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

    // -- format_balance_with_decimals ------------------------------------

    #[test]
    fn format_balance_with_decimals_passes_sentinels_through() {
        // "error" and "not deployed" are sentinel strings the RPC layer
        // emits when a balance call fails; they must reach the user
        // verbatim, not get mangled by the formatter.
        assert_eq!(format_balance_with_decimals("error", 6), "error");
        assert_eq!(
            format_balance_with_decimals("not deployed", 18),
            "not deployed"
        );
    }

    #[test]
    fn format_balance_with_decimals_renders_typical_amounts() {
        // 1 USDC (6 decimals).
        assert_eq!(format_balance_with_decimals("1000000", 6), "1.000000");
        // 1.5 USDC.
        assert_eq!(format_balance_with_decimals("1500000", 6), "1.500000");
        // Sub-unit amount must keep the zero-padded fractional part.
        assert_eq!(
            format_balance_with_decimals("1", 18),
            "0.000000000000000001"
        );
        // Zero.
        assert_eq!(format_balance_with_decimals("0", 6), "0.000000");
    }

    #[test]
    fn format_balance_with_decimals_falls_back_on_unparseable_input() {
        // Anything that isn't a u128 (negative numbers, hex, garbage)
        // round-trips unchanged rather than panicking.
        assert_eq!(format_balance_with_decimals("0xdeadbeef", 6), "0xdeadbeef");
        assert_eq!(format_balance_with_decimals("-1", 6), "-1");
        assert_eq!(format_balance_with_decimals("garbage", 6), "garbage");
    }

    // -- format_balance (Uint<256, 4> variant) ---------------------------

    #[test]
    fn format_balance_handles_zero_decimals() {
        // decimals=0 → integer-only output, no decimal point.
        assert_eq!(format_balance(Uint::<256, 4>::from(0u64), 0), "0");
        assert_eq!(format_balance(Uint::<256, 4>::from(12345u64), 0), "12345");
    }

    #[test]
    fn format_balance_pads_subunit_values() {
        // Value shorter than `decimals` must left-pad with zeros after
        // the implicit "0." prefix.
        assert_eq!(
            format_balance(Uint::<256, 4>::from(1u64), 18),
            "0.000000000000000001"
        );
        assert_eq!(
            format_balance(Uint::<256, 4>::from(1500000u64), 6),
            "1.500000"
        );
    }

    #[test]
    fn format_balance_splits_above_unit_values() {
        // Value longer than `decimals` splits at len - decimals.
        assert_eq!(
            format_balance(Uint::<256, 4>::from(123456789u64), 6),
            "123.456789"
        );
    }

    // -- select_wallet_for_chain -----------------------------------------

    fn evm_chain() -> Chain {
        Chain {
            architecture: "EVM".to_string(),
            ..Default::default()
        }
    }

    fn solana_chain() -> Chain {
        Chain {
            architecture: ARCH_SOLANA.to_string(),
            ..Default::default()
        }
    }

    fn evm_wallet() -> Wallet {
        // Anvil test key #0.
        Wallet::from_evm_hex("0xac0974bec39a17e36ba4a6b4d238ff944bacb478cbed5efcae784d7bf4f2ff80")
            .unwrap()
    }

    #[cfg(feature = "solana")]
    fn solana_wallet() -> Wallet {
        use solana_keypair::Keypair;
        let kp = Keypair::new();
        Wallet::from_solana_base58(&bs58::encode(kp.to_bytes()).into_string()).unwrap()
    }

    #[test]
    fn select_wallet_for_chain_routes_evm_to_secp256k1() {
        let w = evm_wallet();
        let chain = evm_chain();
        let wallets: [&Wallet; 1] = [&w];
        let picked = select_wallet_for_chain(&chain, &wallets);
        assert!(picked.is_some());
        assert_eq!(picked.unwrap().curve(), CurveType::Secp256k1);
    }

    #[cfg(feature = "solana")]
    #[test]
    fn select_wallet_for_chain_routes_solana_to_ed25519() {
        let w = solana_wallet();
        let chain = solana_chain();
        let wallets: [&Wallet; 1] = [&w];
        let picked = select_wallet_for_chain(&chain, &wallets);
        assert!(picked.is_some());
        assert_eq!(picked.unwrap().curve(), CurveType::Ed25519);
    }

    #[cfg(feature = "solana")]
    #[test]
    fn select_wallet_for_chain_picks_matching_curve_from_mixed_set() {
        // Caller passes both wallets; the helper must pick by curve,
        // not by list order. Regression guard for the multi-wallet
        // cross-chain order flow.
        let evm = evm_wallet();
        let sol = solana_wallet();
        let sol_addr = sol.address();

        let mixed_a: [&Wallet; 2] = [&evm, &sol];
        let pick = select_wallet_for_chain(&solana_chain(), &mixed_a).unwrap();
        assert_eq!(pick.curve(), CurveType::Ed25519);
        assert_eq!(pick.address(), sol_addr);

        let mixed_b: [&Wallet; 2] = [&sol, &evm];
        let pick = select_wallet_for_chain(&evm_chain(), &mixed_b).unwrap();
        assert_eq!(pick.curve(), CurveType::Secp256k1);
    }

    #[test]
    fn select_wallet_for_chain_returns_none_when_no_curve_matches() {
        // EVM-only caller, Solana chain → no match (reported as
        // "no wallet" downstream rather than erroring out).
        let w = evm_wallet();
        let wallets: [&Wallet; 1] = [&w];
        assert!(select_wallet_for_chain(&solana_chain(), &wallets).is_none());
    }
}
