//! OSTP Client CLI - Stealth VPN Client for Linux
//!
//! Usage:
//!   sudo ostp-client-linux connect --server 1.2.3.4:443 --psk <hex-key>
//!   sudo ostp-client-linux setup
//!   ostp-client-linux profiles

use anyhow::{Context, Result};
use clap::{Parser, Subcommand};
use console::{style, Emoji};
use dialoguer::{Confirm, Input, Select};
use ostp::{ClientConfig, OstpClient};
use std::net::SocketAddr;
use std::path::PathBuf;
use tracing_subscriber::{fmt, prelude::*, EnvFilter};

static ROCKET: Emoji<'_, '_> = Emoji("🚀 ", "");
static LOCK: Emoji<'_, '_> = Emoji("🔒 ", "");
static CHECK: Emoji<'_, '_> = Emoji("✅ ", "[OK] ");
static WARN: Emoji<'_, '_> = Emoji("⚠️  ", "[!] ");
static GLOBE: Emoji<'_, '_> = Emoji("🌍 ", "");
static CROSS: Emoji<'_, '_> = Emoji("❌ ", "[X] ");

#[derive(Parser)]
#[command(name = "ostp-client-linux")]
#[command(author = "ospab.team")]
#[command(version)]
#[command(about = "OSTP Stealth VPN Client for Linux", long_about = None)]
struct Cli {
    #[command(subcommand)]
    command: Commands,

    /// Log level (trace, debug, info, warn, error)
    #[arg(long, default_value = "info", global = true)]
    log_level: String,
}

#[derive(Subcommand)]
enum Commands {
    /// Connect to an OSTP server
    Connect {
        /// Server address (e.g., vpn.example.com:443)
        #[arg(short, long)]
        server: Option<SocketAddr>,

        /// Pre-shared key in hex format (64 chars)
        #[arg(short, long, env = "OSTP_PSK")]
        psk: Option<String>,

        /// Country code for SNI mimicry (RU, US, DE, NO, CN)
        #[arg(short = 'c', long)]
        country: Option<String>,

        /// Use saved profile
        #[arg(long, conflicts_with_all = ["server", "psk"])]
        profile: Option<String>,

        /// Run in background (daemon mode)
        #[arg(short = 'd', long)]
        daemon: bool,
    },
    /// Disconnect from VPN
    Disconnect,
    /// Show connection status
    Status,
    /// Interactive setup wizard
    Setup,
    /// Test connection to server (handshake only)
    Test {
        /// Server address
        #[arg(short, long)]
        server: SocketAddr,

        /// Pre-shared key in hex format
        #[arg(short, long)]
        psk: String,
    },
    /// Manage connection profiles
    Profiles {
        #[command(subcommand)]
        action: Option<ProfileAction>,
    },
}

#[derive(Subcommand)]
enum ProfileAction {
    /// List all profiles
    List,
    /// Add new profile
    Add {
        /// Profile name
        #[arg(short, long)]
        name: String,
        /// Server address
        #[arg(short, long)]
        server: SocketAddr,
        /// Pre-shared key
        #[arg(short, long)]
        psk: String,
        /// Country code
        #[arg(short, long, default_value = "US")]
        country: String,
    },
    /// Remove profile
    Remove {
        /// Profile name
        name: String,
    },
    /// Set default profile
    SetDefault {
        /// Profile name
        name: String,
    },
}

#[derive(serde::Deserialize, serde::Serialize, Default)]
struct ProfilesFile {
    profiles: Vec<Profile>,
    default_profile: Option<String>,
}

#[derive(serde::Deserialize, serde::Serialize, Clone)]
struct Profile {
    name: String,
    server: String,
    psk: String,
    country: String,
}

fn setup_logging(level: &str) {
    let filter = EnvFilter::try_from_default_env()
        .unwrap_or_else(|_| EnvFilter::new(level));

    tracing_subscriber::registry()
        .with(filter)
        .with(fmt::layer())
        .init();
}

fn get_profiles_path() -> PathBuf {
    let home = std::env::var("HOME").unwrap_or_else(|_| "/root".to_string());
    PathBuf::from(home).join(".config/ostp/profiles.json")
}

fn load_profiles() -> Result<ProfilesFile> {
    let path = get_profiles_path();
    if !path.exists() {
        return Ok(ProfilesFile::default());
    }
    let content = std::fs::read_to_string(&path)
        .context("Failed to read profiles file")?;
    let profiles: ProfilesFile = serde_json::from_str(&content)
        .context("Failed to parse profiles file")?;
    Ok(profiles)
}

fn save_profiles(profiles: &ProfilesFile) -> Result<()> {
    let path = get_profiles_path();
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    let content = serde_json::to_string_pretty(profiles)?;
    std::fs::write(&path, content)?;
    Ok(())
}

fn check_root() -> Result<()> {
    unsafe {
        if libc::geteuid() != 0 {
            anyhow::bail!("This command requires root privileges. Please run with sudo.");
        }
    }
    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    let cli = Cli::parse();
    setup_logging(&cli.log_level);

    match cli.command {
        Commands::Connect {
            server,
            psk,
            country,
            profile,
            daemon,
        } => {
            check_root()?;

            let (server_addr, psk_hex, country_code) = if let Some(profile_name) = profile {
                // Load from profile
                let profiles = load_profiles()?;
                let profile = profiles
                    .profiles
                    .iter()
                    .find(|p| p.name == profile_name)
                    .ok_or_else(|| anyhow::anyhow!("Profile '{}' not found", profile_name))?;

                let addr: SocketAddr = profile.server.parse()
                    .context("Invalid server address in profile")?;

                (addr, profile.psk.clone(), profile.country.clone())
            } else {
                // Use command-line args
                let srv = server.ok_or_else(|| anyhow::anyhow!("--server is required"))?;
                let key = psk.ok_or_else(|| anyhow::anyhow!("--psk is required"))?;
                let cc = country.unwrap_or_else(|| "US".to_string());

                (srv, key, cc)
            };

            connect(server_addr, &psk_hex, &country_code, daemon).await
        }
        Commands::Disconnect => {
            check_root()?;
            disconnect().await
        }
        Commands::Status => {
            show_status().await
        }
        Commands::Setup => {
            run_setup_wizard().await
        }
        Commands::Test { server, psk } => {
            test_connection(server, &psk).await
        }
        Commands::Profiles { action } => {
            handle_profiles(action).await
        }
    }
}

async fn connect(
    server: SocketAddr,
    psk: &str,
    country: &str,
    daemon: bool,
) -> Result<()> {
    println!("{} {}Connecting to OSTP server...", ROCKET, style("").bold());
    println!("   Server:  {}", style(server).cyan());
    println!("   Country: {}", style(country).yellow());
    println!();

    // Validate PSK format
    if psk.len() != 64 {
        anyhow::bail!("PSK must be 64 hex characters (32 bytes)");
    }
    let psk_bytes = hex::decode(psk).context("Invalid PSK hex format")?;

    // Run anti-analysis checks (production mode only)
    #[cfg(not(debug_assertions))]
    {
        if ostp_guard::anti_vm::is_virtual_machine() {
            println!("{} {}Security check failed", CROSS, style("").red());
            return Ok(());
        }
    }

    // Create OSTP client config
    let config = ClientConfig {
        server_addr: server,
        psk: psk_bytes.try_into().unwrap(),
        country_code: country.to_string(),
    };

    // Connect to server
    println!("{} {}Establishing secure tunnel...", LOCK, style("").bold());
    let mut client = OstpClient::new(config);
    let _stream = client.connect().await
        .context("Failed to connect to server")?;

    println!("{} {}Connected successfully!", CHECK, style("").green());
    println!();

    // Create TUN interface
    println!("{} {}Creating TUN interface...", GLOBE, style("").bold());
    let tun = create_tun_interface()?;
    
    println!("{} {}Configuring routes...", GLOBE, style("").bold());
    configure_routing(&tun)?;

    println!();
    println!("{} {}VPN connection established!", CHECK, style("").green().bold());
    println!();
    println!("   Interface: {}", style("ostp0").cyan());
    println!("   Gateway:   {}", style("10.X.0.1").cyan());
    println!();

    if daemon {
        println!("Running in background mode. Press Ctrl+C to stop.");
        // TODO: Daemonize process
    } else {
        println!("Press Ctrl+C to disconnect...");
    }

    // Main relay loop
    tokio::select! {
        result = relay_traffic(&mut client, &tun) => {
            result?;
        }
        _ = tokio::signal::ctrl_c() => {
            println!();
            println!("{} {}Disconnecting...", WARN, style("").yellow());
        }
    }

    cleanup_routing(&tun)?;
    println!("{} {}Disconnected", CHECK, style("").green());

    Ok(())
}

async fn disconnect() -> Result<()> {
    println!("{} {}Disconnecting from VPN...", WARN, style("").yellow());
    
    // Find and kill running ostp-client-linux process
    let output = std::process::Command::new("pgrep")
        .arg("-f")
        .arg("ostp-client-linux connect")
        .output()?;

    if output.status.success() {
        let pids = String::from_utf8_lossy(&output.stdout);
        for pid in pids.lines() {
            if let Ok(pid_num) = pid.trim().parse::<i32>() {
                // Don't kill ourselves
                if pid_num != std::process::id() as i32 {
                    println!("   Killing process {}", pid_num);
                    unsafe {
                        libc::kill(pid_num, libc::SIGTERM);
                    }
                }
            }
        }
        println!("{} {}Disconnected", CHECK, style("").green());
    } else {
        println!("{} {}No active VPN connection found", WARN, style("").yellow());
    }

    Ok(())
}

async fn show_status() -> Result<()> {
    println!("{} {}VPN Connection Status", GLOBE, style("").bold());
    println!();

    // Check if client is running
    let output = std::process::Command::new("pgrep")
        .arg("-f")
        .arg("ostp-client-linux connect")
        .output()?;

    if output.status.success() && !output.stdout.is_empty() {
        println!("   Status: {}", style("Connected").green().bold());
        
        // Try to get interface stats
        if let Ok(stats) = get_interface_stats("ostp0") {
            println!();
            println!("   Interface:  ostp0");
            println!("   RX Bytes:   {} MB", stats.rx_bytes / 1_000_000);
            println!("   TX Bytes:   {} MB", stats.tx_bytes / 1_000_000);
            println!("   RX Packets: {}", stats.rx_packets);
            println!("   TX Packets: {}", stats.tx_packets);
        }
    } else {
        println!("   Status: {}", style("Disconnected").red().bold());
    }

    println!();
    Ok(())
}

async fn test_connection(server: SocketAddr, psk: &str) -> Result<()> {
    println!("{} {}Testing connection to {}...", ROCKET, style("").bold(), server);
    println!();

    let psk_bytes = hex::decode(psk).context("Invalid PSK hex format")?;
    let config = ClientConfig {
        server_addr: server,
        psk: psk_bytes.try_into().unwrap(),
        country_code: "US".to_string(),
    };

    let mut client = OstpClient::new(config);
    match client.connect().await {
        Ok(_) => {
            println!("{} {}Connection test successful!", CHECK, style("").green());
            println!("   Server is reachable");
            println!("   PSK is valid");
            println!("   Handshake completed");
        }
        Err(e) => {
            println!("{} {}Connection test failed", CROSS, style("").red());
            println!("   Error: {}", e);
        }
    }

    println!();
    Ok(())
}

async fn run_setup_wizard() -> Result<()> {
    println!("{}", style("═══════════════════════════════════════").cyan());
    println!("{} {}", ROCKET, style("OSTP Client Setup Wizard").cyan().bold());
    println!("{}", style("═══════════════════════════════════════").cyan());
    println!();

    let server_input: String = Input::new()
        .with_prompt("Server address (e.g., vpn.example.com:443)")
        .interact()?;

    let server: SocketAddr = server_input.parse()
        .context("Invalid server address format")?;

    let psk: String = Input::new()
        .with_prompt("Pre-shared key (64 hex characters)")
        .interact()?;

    if psk.len() != 64 {
        anyhow::bail!("PSK must be exactly 64 hex characters");
    }
    hex::decode(&psk).context("PSK must be valid hex")?;

    let countries = vec!["US", "RU", "DE", "NO", "CN", "FR", "GB", "JP"];
    let country_idx = Select::new()
        .with_prompt("Select country for SNI mimicry")
        .items(&countries)
        .default(0)
        .interact()?;
    let country = countries[country_idx].to_string();

    let profile_name: String = Input::new()
        .with_prompt("Profile name (e.g., 'US-West')")
        .default("default".to_string())
        .interact()?;

    // Save profile
    let mut profiles = load_profiles()?;
    
    // Remove existing profile with same name
    profiles.profiles.retain(|p| p.name != profile_name);
    
    profiles.profiles.push(Profile {
        name: profile_name.clone(),
        server: server.to_string(),
        psk: psk.clone(),
        country: country.clone(),
    });

    let set_default = Confirm::new()
        .with_prompt("Set as default profile?")
        .default(profiles.default_profile.is_none())
        .interact()?;

    if set_default {
        profiles.default_profile = Some(profile_name.clone());
    }

    save_profiles(&profiles)?;

    println!();
    println!("{} {}Profile '{}' saved!", CHECK, style("").green(), profile_name);
    println!();
    println!("To connect, run:");
    println!("   {}", style(format!("sudo ostp-client-linux connect --profile {}", profile_name)).cyan());
    println!();

    Ok(())
}

async fn handle_profiles(action: Option<ProfileAction>) -> Result<()> {
    match action.unwrap_or(ProfileAction::List) {
        ProfileAction::List => {
            let profiles = load_profiles()?;
            
            if profiles.profiles.is_empty() {
                println!("{} {}No profiles configured", WARN, style("").yellow());
                println!();
                println!("Run {} to create a profile", style("ostp-client-linux setup").cyan());
                return Ok(());
            }

            println!("{} {}Saved Profiles", GLOBE, style("").bold());
            println!("{}", style("─".repeat(60)).dim());
            
            for profile in &profiles.profiles {
                let is_default = profiles.default_profile.as_ref() == Some(&profile.name);
                let marker = if is_default { " (default)" } else { "" };
                
                println!();
                println!("   {}{}", style(&profile.name).cyan().bold(), style(marker).green());
                println!("   Server:  {}", profile.server);
                println!("   Country: {}", profile.country);
            }
            
            println!();
        }
        ProfileAction::Add { name, server, psk, country } => {
            let mut profiles = load_profiles()?;
            
            profiles.profiles.retain(|p| p.name != name);
            profiles.profiles.push(Profile {
                name: name.clone(),
                server: server.to_string(),
                psk,
                country,
            });
            
            save_profiles(&profiles)?;
            println!("{} {}Profile '{}' added", CHECK, style("").green(), name);
        }
        ProfileAction::Remove { name } => {
            let mut profiles = load_profiles()?;
            let initial_len = profiles.profiles.len();
            
            profiles.profiles.retain(|p| p.name != name);
            
            if profiles.profiles.len() == initial_len {
                println!("{} {}Profile '{}' not found", CROSS, style("").red(), name);
            } else {
                if profiles.default_profile.as_ref() == Some(&name) {
                    profiles.default_profile = None;
                }
                save_profiles(&profiles)?;
                println!("{} {}Profile '{}' removed", CHECK, style("").green(), name);
            }
        }
        ProfileAction::SetDefault { name } => {
            let mut profiles = load_profiles()?;
            
            if profiles.profiles.iter().any(|p| p.name == name) {
                profiles.default_profile = Some(name.clone());
                save_profiles(&profiles)?;
                println!("{} {}'{}' set as default profile", CHECK, style("").green(), name);
            } else {
                println!("{} {}Profile '{}' not found", CROSS, style("").red(), name);
            }
        }
    }
    
    Ok(())
}

// TUN interface management
fn create_tun_interface() -> Result<TunInterface> {
    // Placeholder - will use osn crate
    tracing::warn!("TUN interface creation not yet fully implemented - using stub");
    println!("   {}[STUB] TUN interface (ostp0)", style("⚠ ").yellow());
    Ok(TunInterface {
        name: "ostp0".to_string(),
    })
}

fn configure_routing(_tun: &TunInterface) -> Result<()> {
    // Configure routes via ip command
    tracing::warn!("Route configuration not yet fully implemented - using stub");
    println!("   {}[STUB] Routes configured", style("⚠ ").yellow());
    Ok(())
}

fn cleanup_routing(_tun: &TunInterface) -> Result<()> {
    tracing::debug!("Route cleanup - stub");
    Ok(())
}

async fn relay_traffic(_client: &mut OstpClient, _tun: &TunInterface) -> Result<()> {
    // Main packet relay loop
    tracing::warn!("Traffic relay not yet fully implemented - connection will stay idle");
    println!("   {}[STUB] Traffic relay active (no actual forwarding yet)", style("⚠ ").yellow());
    
    // Keep connection alive - wait indefinitely
    // In real implementation, this would be the packet forwarding loop
    tokio::time::sleep(tokio::time::Duration::from_secs(u64::MAX)).await;
    Ok(())
}

struct TunInterface {
    name: String,
}

struct InterfaceStats {
    rx_bytes: u64,
    tx_bytes: u64,
    rx_packets: u64,
    tx_packets: u64,
}

fn get_interface_stats(iface: &str) -> Result<InterfaceStats> {
    let rx_bytes = std::fs::read_to_string(format!("/sys/class/net/{}/statistics/rx_bytes", iface))?
        .trim()
        .parse()?;
    let tx_bytes = std::fs::read_to_string(format!("/sys/class/net/{}/statistics/tx_bytes", iface))?
        .trim()
        .parse()?;
    let rx_packets = std::fs::read_to_string(format!("/sys/class/net/{}/statistics/rx_packets", iface))?
        .trim()
        .parse()?;
    let tx_packets = std::fs::read_to_string(format!("/sys/class/net/{}/statistics/tx_packets", iface))?
        .trim()
        .parse()?;

    Ok(InterfaceStats {
        rx_bytes,
        tx_bytes,
        rx_packets,
        tx_packets,
    })
}
