use clap::{Arg, Command};
use reqwest::Client;
use std::collections::HashMap;
use std::error::Error;
use serde::{Deserialize,Serialize};
use futures::future::join_all;
use tokio::task::JoinHandle;
use toml;
use dirs;
use indicatif::{ProgressBar, ProgressStyle};
use colored::*;
use alloy::providers::{ProviderBuilder};
use alloy_ccip_read::CCIPReader;

#[derive(Serialize)]
struct JsonRpcRequest {
  jsonrpc: String,
  method: String,
  params: Vec<String>,
  id: u32,
}

#[derive(Deserialize)]
struct JsonRpcResponse {
  result: String,
}


#[derive(Deserialize)]
struct Config {
  address: Option<String>,
  networks: Option<HashMap<String, NetworkConfig>>,
}

#[derive(Deserialize)]
struct TokenConfig {
  address: String,
  decimals: u8,
}

#[derive(Deserialize)]
struct NetworkConfig {
  name: String,
  rpc_url: String,
  tokens: Option<HashMap<String, TokenConfig>>,
}

#[derive(Clone)]
struct TokenInfo {
  symbol: String,
  address: String,
  decimals: u8,
}

#[derive(Clone)]
struct Network {
  chain_id: u64,
  name: String,
  rpc_url: String,
  tokens: Vec<TokenInfo>,
}

#[derive(Clone)]
struct TokenBalance {
  network_name: String,
  symbol: String,
  balance: f64,
}

enum BalanceResult {
  Native(f64, String),
  Token(TokenBalance),
}

fn read_config() -> Result<Config, Box<dyn Error>> {
  let home_dir = dirs::home_dir().ok_or("Could not find home directory")?;
  let config_dir = home_dir.join(".config").join("walletfetch");
  let config_path = config_dir.join("config.toml");

  if !config_path.exists(){

    std::fs::create_dir_all(&config_dir)?;

    let default_config = r#"# WalletFetch Configuration
# You can set a default address here (optional)
# address = "0x..."

[networks.1]
name = "Mainnet"
rpc_url = "https://eth.drpc.org"

[networks.1.tokens]
USDC = { address = "0xA0b86991c6218b36c1d19D4a2e9Eb0cE3606eB48", decimals = 6 }

[networks.42161]
name = "Arbitrum"
rpc_url = "https://arbitrum.drpc.org"

[networks.42161.tokens]
USDC = { address = "0xaf88d065e77c8cC2239327C5EDb3A432268e5831", decimals = 6 }

[networks.8453]
name = "Base"
rpc_url = "https://base.drpc.org"

[networks.8453.tokens]
USDC = { address = "0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913", decimals = 6 }
"#;

    std::fs::write(&config_path, default_config)?;

    println!("Created default config at: {}", config_path.display());
    println!("You can edit this file to customize your networks and tokens.");
  }

  let config_content = std::fs::read_to_string(config_path)?;
  let config: Config = toml::from_str(&config_content)?;

  Ok(config)
}

fn collect_rpc_urls(config: &Config) -> HashMap<u64, Network> {
  let mut networks = HashMap::new();

  if let Some(network_configs) = &config.networks {
    for (chain_id_str, network_config) in network_configs {
      if let Ok(chain_id) = chain_id_str.parse::<u64>(){
        let mut tokens = Vec::new();
        if let Some(token_configs) = &network_config.tokens {
          for (symbol, token_config) in token_configs {
            tokens.push(TokenInfo {
              symbol: symbol.clone(),
              address: token_config.address.clone(),
              decimals: token_config.decimals,
            });
          }
        }

        networks.insert(chain_id, Network{
          chain_id,
          name: network_config.name.clone(),
          rpc_url: network_config.rpc_url.clone(),
          tokens,
        });
      }
    }
  }

  networks
}

fn format_balance_smart(balance: f64, symbol: &str) -> String {
    let formatted = if balance >= 1_000_000.0 {
        format!("{:.2}M", balance / 1_000_000.0)
    } else if balance >= 1_000.0 {
        format!("{:.2}K", balance / 1_000.0)
    } else if balance >= 1.0 {
        format!("{:.2}", balance)
    } else if balance > 0.0 {
        format!("{:.6}", balance).trim_end_matches('0').trim_end_matches('.').to_string()
    } else {
        "0".to_string()
    };

    format!("{} {}", formatted, symbol)
}

async fn fetch_balance(
  client: &Client,
  address: &str,
  network: &Network
) -> Result<(u64, f64, String), Box<dyn Error + Send + Sync >> {
  let request_data = JsonRpcRequest {
    jsonrpc: "2.0".to_string(),
    method: "eth_getBalance".to_string(),
    params: vec![address.to_string(), "latest".to_string()],
    id: 1,
  };

  let response = client.post(&network.rpc_url)
    .json(&request_data)
    .send()
    .await?;

  let status = response.status();

  if !response.status().is_success() {
    let error_text = response.text().await?;
    return Err(format!("HTTP error {}: {}", status, error_text).into());
  }

  let response_text = response.text().await?;
  let response_body: JsonRpcResponse = serde_json::from_str(&response_text)?;

  if let Some(hex_str) = response_body.result.strip_prefix("0x"){
    if let Ok(balance) = u128::from_str_radix(hex_str, 16){
      let eth_balance = balance as f64 / 1_000_000_000_000_000_000.0;
      return Ok((network.chain_id, eth_balance, network.name.clone()));
    }
  }
  Err(format!("Failed to parse balance for network {}", network.name).into())
}

async fn fetch_token_balance(
  client: &Client,
  address: &str,
  token: &TokenInfo,
  network: &Network,
) -> Result<TokenBalance, Box<dyn Error + Send + Sync>> {
  let clean_address = address.strip_prefix("0x").unwrap_or(address).to_lowercase();
  let data = format!("0x70a08231000000000000000000000000{}", clean_address);

  let request_data = serde_json::json!({
    "jsonrpc": "2.0",
    "method": "eth_call",
    "params": [
      {
        "to": token.address,
        "data": data
      },
      "latest"
    ],
    "id": 1
  });

  let response = client.post(&network.rpc_url)
    .json(&request_data)
    .send()
    .await?;

  let status = response.status();

  if !status.is_success(){
    let error_text = response.text().await?;
    return Err(format!("HTTP error {}: {}", status, error_text).into());
  }

  let response_text = response.text().await?;

  let response_body: JsonRpcResponse = serde_json::from_str(&response_text)?;

  if let Some(hex_str) = response_body.result.strip_prefix("0x") {
    if let Ok(raw_balance) = u128::from_str_radix(hex_str, 16){
      let divisor = 10_u128.pow(token.decimals as u32) as f64;
      let balance = raw_balance as f64 / divisor;

      return Ok(TokenBalance {
        network_name: network.name.clone(),
        symbol: token.symbol.clone(),
        balance,
      });
    }
  }

  Err(format!("Failed to parse balance for token {} on network {}", token.symbol, network.name).into())
}

fn get_eth_logo() -> &'static str {
r#"------------------------------
--------------4%--------------
-------------44HH-------------
------------444HHH------------
-----------4444HHHH-----------
---------~44444HHHHH~---------
--------4444444HHHHHHW--------
-------4444HHHHWWWWHHHH-------
------KHHHHHHHHWWWWWWWWW------
---------HHHHHHWWWWWW---------
-------44---HHHWWW---HH-------
--------~444?----4HHH~--------
----------44444HHHHH----------
-----------L444HHHq-----------
-------------44HH-------------
--------------4H--------------
------------------------------"#
}

async fn fetch_all_balances(
  address: &str,
  networks: HashMap<u64, Network>
) -> Result<Vec<BalanceResult>, Box<dyn Error>> {
  let client = Client::new();

  let mut tasks: Vec<JoinHandle<Result<BalanceResult, Box<dyn Error + Send + Sync>>>> = Vec::new();

  for (_, network) in &networks {
    let client_clone = client.clone();
    let address_clone = address.to_string();
    let network_clone = network.clone();

    let task = tokio::spawn(async move {
      let (_, balance, name) = fetch_balance(&client_clone, &address_clone, &network_clone).await?;
      Ok(BalanceResult::Native(balance, name))
    });

    tasks.push(task);

    for token in &network.tokens {
      let client_clone = client.clone();
      let address_clone = address.to_string();
      let token_clone = token.clone();
      let network_clone = network.clone();

      let task = tokio::spawn(async move {
        let token_balance = fetch_token_balance(&client_clone, &address_clone, &token_clone, &network_clone).await?;
        Ok(BalanceResult::Token(token_balance))
      });

      tasks.push(task);
    }
  }

  let results = join_all(tasks).await;

  let mut balances = Vec::new();
  for result in results {
    match result {
      Ok(Ok(balance_result)) => {
        balances.push(balance_result);
      },
      Ok(Err(e)) => {
        eprintln!("Error fetching balance: {}", e);
      },
      Err(e) => {
        eprintln!("Task error: {}", e);
      }
    }
  }
  Ok(balances)
}

async fn resolve_address_or_ens(
  input: &str,
  networks: &HashMap<u64, Network>
) -> Result<String, Box<dyn Error>> {
  if input.to_lowercase().ends_with(".eth") {
    let mainnet = networks.get(&1);

    let rpc_url = match mainnet {
      Some(network) => &network.rpc_url,
      None => return Err("Ethereum mainnet configuration not found for ENS resolution".into())
    };

    let provider = ProviderBuilder::new()
      .on_http(rpc_url.parse().unwrap());

    let reader = CCIPReader::new(provider.boxed());

    let resolution_result = match reader.resolve_name(input).await {
      Ok(result) => result,
      Err(e) => return Err(format!("Failed to resolve address for {}: {}", input, e).into())
    };

    let eth_address = format!("{}", resolution_result.addr.value);

    return Ok(eth_address);
  } else {
    if !input.starts_with("0x") || input.len() != 42 {
      return Err("Invalid Ethereum address format".into());
    }

    Ok(input.to_string())
  }
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
    let matches = Command::new("wallet-fetch")
        .version("0.0.1")
        .author("Steve Simkins")
        .about("Neofetch but for your wallet")
        .arg(
            Arg::new("address")
                .help("Address to fetch info for ")
                .index(1)
                .required(false)
        )
        .get_matches();

    let spinner = ProgressBar::new_spinner();
    spinner.set_style(
        ProgressStyle::default_spinner()
            .tick_chars("⠁⠂⠄⡀⢀⠠⠐⠈ ")
            .template("{spinner:.green} {msg}")
            .unwrap()
    );
    spinner.set_message("Fetching balances...");
    spinner.enable_steady_tick(std::time::Duration::from_millis(100));

    let config = read_config()?;

    let input = match matches.get_one::<String>("address"){
      Some(addr) if !addr.is_empty() => addr.to_string(),
      _ => match &config.address {
        Some(addr) if !addr.is_empty() => addr.clone(),
        _ => {
          eprintln!("Error: No address provided. Either pass it as an argument or set it in the config file");
          return Err("No address provided".into());
        }
      },
    };

    let networks = collect_rpc_urls(&config);

    if networks.is_empty(){
      eprintln!("RPC URLs are not defined");
    }

    let address = resolve_address_or_ens(&input, &networks).await?;


    let balances = fetch_all_balances(&address, networks).await?;

    spinner.finish_and_clear();

    if balances.is_empty(){
      println!("No balances found for address {}", address);
    } else {
      let mut network_balances: HashMap<String, Vec<String>> = HashMap::new();

      for balance in balances {
        match balance {
          BalanceResult::Native(eth_balance, network_name) => {
            let balance_str = format_balance_smart(eth_balance, "ETH");
            network_balances.entry(network_name).or_default().push(balance_str);
          },
          BalanceResult::Token(token_balance) => {
            let balance_str = format_balance_smart(token_balance.balance, &token_balance.symbol);
            network_balances.entry(token_balance.network_name).or_default().push(balance_str);
          }
        }
      }

      let logo = get_eth_logo();
      let logo_lines: Vec<&str> = logo.lines().collect();
      let logo_height = logo_lines.len();
      let logo_width = logo_lines.iter().map(|line| line.len()).max().unwrap_or(0);


      println!();
      if !input.starts_with("0x"){
        let ens_display = format!("{}", input.bright_green());
        println!("{}", format!("ENS: {}", ens_display).bright_cyan());
      }
      let address_display = format!("{}", address.bright_green());
      println!("{}", format!("Wallet: {}", address_display).bright_cyan());
      println!("{}", "=".repeat(50).bright_blue());

      let mut info_lines = Vec::new();

      for (network, balances) in network_balances {
        if !info_lines.is_empty() {
          info_lines.push("".to_string());
        }
        info_lines.push(format!("{}: ", network.bright_yellow()));
        for balance in balances {
          info_lines.push(format!("  {} {}", "•".bright_green(), balance.bright_white()));
        }
      }

      let display_lines = std::cmp::max(logo_height, info_lines.len());

      for i in 0..display_lines {
        let logo_line = if i < logo_height {
          logo_lines[i]
        } else {
          &" ".repeat(logo_width)
        };
        let info_line = if i < info_lines.len() { &info_lines[i] } else { "" };

        println!("{}    {}", logo_line.bright_cyan(), info_line);
      }

      println!();
    }

    Ok(())
}