//! Antigravity OpenAI Proxy — Rust edition v3 (stealth hardened).
//!
//! Single-binary replacement for server.py. BoringSSL TLS impersonation,
//! byte-exact protobuf encoding, Chrome header fingerprinting, cascade
//! session management, warmup + heartbeat lifecycle mimicry.

mod api;
mod backend;
mod constants;
mod mitm;
mod proto;
mod quota;
mod session;
mod standalone;
mod warmup;

use api::AppState;
use backend::Backend;
use clap::Parser;
use session::SessionManager;
use std::sync::Arc;
use tracing::{info, warn};

use mitm::store::MitmStore;

#[derive(Parser)]
#[command(name = "antigravity-proxy", about = "Antigravity OpenAI Proxy (stealth)")]
struct Cli {
    /// Port to listen on
    #[arg(long, default_value_t = 8741)]
    port: u16,

    /// Enable info-level logging (-v)
    #[arg(short, long)]
    verbose: bool,

    /// Enable debug-level logging (-d)
    #[arg(short, long)]
    debug: bool,

    /// Disable the MITM proxy (no API interception)
    #[arg(long)]
    no_mitm: bool,

    /// MITM proxy port (default: 8742, matches wrapper script)
    #[arg(long, default_value_t = 8742)]
    mitm_port: u16,

    /// Disable standalone LS — attach to the real running LS instead
    #[arg(long)]
    no_standalone: bool,
}

#[tokio::main]
async fn main() {
    // Ignore SIGPIPE — prevents instant death when piped through tee/grep
    #[cfg(unix)]
    {
        use tokio::signal::unix::{signal, SignalKind};
        let mut sigpipe = signal(SignalKind::pipe()).expect("failed to install SIGPIPE handler");
        tokio::spawn(async move {
            loop {
                sigpipe.recv().await;
                // Silently ignore SIGPIPE
            }
        });
    }

    // Install rustls CryptoProvider early — prevents panic under concurrent load
    let _ = rustls::crypto::ring::default_provider().install_default();

    let cli = Cli::parse();

    // Flag > env var > default (warn)
    let log_level = if cli.debug {
        "debug"
    } else if cli.verbose {
        "info"
    } else {
        // Fall back to RUST_LOG env, or warn-only
        ""
    };

    let filter = if log_level.is_empty() {
        tracing_subscriber::EnvFilter::try_from_default_env()
            .unwrap_or_else(|_| "warn".into())
    } else {
        tracing_subscriber::EnvFilter::new(log_level)
    };

    tracing_subscriber::fmt()
        .with_env_filter(filter)
        .init();

    // ── Step 1: Bind main port FIRST (fail fast, before spawning anything) ────
    let addr = format!("127.0.0.1:{}", cli.port);
    let listener = match tokio::net::TcpListener::bind(&addr).await {
        Ok(l) => l,
        Err(e) => {
            eprintln!("Fatal: cannot bind to {addr}: {e}");
            eprintln!("Hint: kill $(lsof -ti:{}) 2>/dev/null", cli.port);
            std::process::exit(1);
        }
    };

    // ── Step 2: Backend discovery (or standalone LS spawn) ─────────────────────
    let standalone_ls = if !cli.no_standalone {
        // Standalone mode: discover main LS config, spawn our own
        let main_config = match standalone::discover_main_ls_config() {
            Ok(c) => c,
            Err(e) => {
                eprintln!("Fatal: {e}");
                std::process::exit(1);
            }
        };
        // Build MITM config if MITM is enabled
        let mitm_cfg = if !cli.no_mitm {
            let ca_path = dirs_data_dir()
                .join("mitm-ca.pem")
                .to_string_lossy()
                .to_string();
            Some(standalone::StandaloneMitmConfig {
                proxy_addr: format!("http://127.0.0.1:{}", cli.mitm_port),
                ca_cert_path: ca_path,
            })
        } else {
            None
        };

        let ls = match standalone::StandaloneLS::spawn(&main_config, mitm_cfg.as_ref()) {
            Ok(ls) => ls,
            Err(e) => {
                eprintln!("Fatal: failed to spawn standalone LS: {e}");
                std::process::exit(1);
            }
        };
        // Wait for it to be ready
        let rt_ls_port = ls.port;
        let rt_ls_csrf = ls.csrf.clone();
        tokio::task::block_in_place(|| {
            tokio::runtime::Handle::current().block_on(async {
                if let Err(e) = ls.wait_ready(10).await {
                    eprintln!("Fatal: {e}");
                    std::process::exit(1);
                }
            });
        });
        info!(port = rt_ls_port, "Standalone LS ready");
        Some((ls, rt_ls_port, rt_ls_csrf))
    } else {
        None
    };

    let backend = Arc::new(if let Some((_, port, ref csrf)) = standalone_ls {
        // Build backend pointing at standalone LS
        let oauth = std::env::var("ANTIGRAVITY_OAUTH_TOKEN")
            .ok()
            .filter(|s| !s.is_empty())
            .or_else(|| {
                let home = std::env::var("HOME").unwrap_or_default();
                let path = format!("{home}/.config/antigravity-proxy-token");
                std::fs::read_to_string(&path)
                    .ok()
                    .map(|s| s.trim().to_string())
                    .filter(|s| !s.is_empty())
            })
            .unwrap_or_default();
        match Backend::new_with_config(port, csrf.clone(), oauth) {
            Ok(b) => b,
            Err(e) => {
                eprintln!("Fatal: {e}");
                std::process::exit(1);
            }
        }
    } else {
        // Normal mode: discover existing LS
        match Backend::new() {
            Ok(b) => b,
            Err(e) => {
                eprintln!("Fatal: {e}");
                std::process::exit(1);
            }
        }
    });

    let (pid, https_port, csrf, token) = backend.info().await;

    // ── Step 3: MITM proxy (after port is secured) ────────────────────────────
    let mitm_store = MitmStore::new();
    let (mitm_port_actual, mitm_handle) = if !cli.no_mitm {
        let data_dir = dirs_data_dir();
        match mitm::ca::MitmCa::load_or_generate(&data_dir) {
            Ok(ca) => {
                let ca = Arc::new(ca);
                let ca_pem = ca.ca_pem_path.display().to_string();
                let config = mitm::proxy::MitmConfig {
                    port: cli.mitm_port,
                    modify_requests: true,
                };
                match mitm::proxy::run(ca, mitm_store.clone(), config).await {
                    Ok((port, handle)) => {
                        info!(port, ca = %ca_pem, "MITM proxy started");
                        // Write actual port to file for wrapper script discovery
                        let port_file = data_dir.join("mitm-port");
                        if let Err(e) = std::fs::write(&port_file, port.to_string()) {
                            warn!("Failed to write MITM port file: {e}");
                        }
                        (Some((port, ca_pem)), Some(handle))
                    }
                    Err(e) => {
                        warn!("MITM proxy failed to start: {e}");
                        (None, None)
                    }
                }
            }
            Err(e) => {
                warn!("MITM CA generation failed: {e}");
                (None, None)
            }
        }
    } else {
        info!("MITM proxy disabled (--no-mitm)");
        (None, None)
    };

    // ── Step 4: Warmup + heartbeat ────────────────────────────────────────────
    warmup::warmup_sequence(&backend).await;
    let heartbeat_handle = warmup::start_heartbeat(Arc::clone(&backend));

    // ── Step 4b: Quota monitor ────────────────────────────────────────────────
    let quota_store = quota::QuotaStore::new();
    quota_store.clone().start_polling(Arc::clone(&backend));
    info!("Quota monitor started (polling every 60s)");

    let state = Arc::new(AppState {
        backend,
        sessions: SessionManager::new(),
        mitm_store,
        quota_store,
    });

    // Periodic backend refresh — keeps LS connection details fresh
    // (skip in standalone mode — the port is fixed and discover() would overwrite it)
    let is_standalone = !cli.no_standalone;
    let refresh_backend = Arc::clone(&state.backend);
    let refresh_handle = tokio::spawn(async move {
        if is_standalone {
            // In standalone mode, the backend config is fixed — no refresh needed
            std::future::pending::<()>().await;
            return;
        }
        loop {
            tokio::time::sleep(tokio::time::Duration::from_secs(60)).await;
            if let Err(e) = refresh_backend.refresh().await {
                warn!("Periodic refresh failed: {e}");
            }
        }
    });

    // ── Step 5: Start serving ─────────────────────────────────────────────────
    let app = api::router(state.clone());

    print_banner(cli.port, &pid, &https_port, &csrf, &token, &mitm_port_actual, is_standalone);
    info!("Listening on http://{addr}");

    axum::serve(listener, app)
        .with_graceful_shutdown(shutdown_signal())
        .await
        .expect("server error");

    // ── Cleanup: abort all background tasks ───────────────────────────────────
    heartbeat_handle.abort();
    refresh_handle.abort();
    if let Some(h) = mitm_handle {
        h.abort();
    }
    // Kill standalone LS if we spawned one
    if let Some((mut ls, _, _)) = standalone_ls {
        ls.kill();
    }
    // Remove stale MITM port file
    let _ = std::fs::remove_file(dirs_data_dir().join("mitm-port"));
    eprintln!("  \x1b[1;32m✓ Server shutdown complete\x1b[0m\n");
    info!("Server shutdown complete");
}

/// Wait for SIGINT (Ctrl+C) or SIGTERM for graceful shutdown.
async fn shutdown_signal() {
    let ctrl_c = async {
        tokio::signal::ctrl_c()
            .await
            .expect("failed to install Ctrl+C handler");
    };

    #[cfg(unix)]
    let terminate = async {
        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
            .expect("failed to install SIGTERM handler")
            .recv()
            .await;
    };

    #[cfg(not(unix))]
    let terminate = std::future::pending::<()>();

    tokio::select! {
        _ = ctrl_c => {
            eprintln!("\n  \x1b[1;33m⚡ Shutting down gracefully...\x1b[0m");
            info!("Received SIGINT, shutting down...");
        },
        _ = terminate => {
            eprintln!("\n  \x1b[1;33m⚡ Received SIGTERM, shutting down...\x1b[0m");
            info!("Received SIGTERM, shutting down...");
        },
    }
}

fn print_banner(port: u16, pid: &str, https_port: &str, csrf: &str, token: &str, mitm: &Option<(u16, String)>, is_standalone: bool) {
    let chrome_major = &*constants::CHROME_MAJOR;
    let ver = crate::constants::antigravity_version();

    println!();
    println!("  \x1b[1;35m>> antigravity-proxy\x1b[0m \x1b[2mv{ver}\x1b[0m");
    println!("  \x1b[2m────────────────────────────────────────────────\x1b[0m");
    println!();
    println!("  \x1b[1mcore\x1b[0m");
    println!("  \x1b[36m  tls\x1b[0m        BoringSSL (Chrome {chrome_major})");
    println!("  \x1b[36m  listen\x1b[0m     http://127.0.0.1:{port}");
    println!("  \x1b[36m  ls pid\x1b[0m     {pid}");
    println!("  \x1b[36m  https\x1b[0m      :{https_port}");
    println!("  \x1b[36m  csrf\x1b[0m       {csrf}");
    println!("  \x1b[36m  oauth\x1b[0m      {token}");
    println!();

    // MITM section
    if let Some((mitm_port, ca_path)) = mitm {
        println!("  \x1b[1mmitm\x1b[0m");
        println!("  \x1b[36m  proxy\x1b[0m      127.0.0.1:{mitm_port}");
        println!("  \x1b[36m  ca cert\x1b[0m    {ca_path}");

        // Check if wrapper is installed
        if is_standalone {
            println!("  \x1b[36m  wrapper\x1b[0m    \x1b[32miptables (standalone)\x1b[0m");
        } else if check_wrapper_installed() {
            println!("  \x1b[36m  wrapper\x1b[0m    \x1b[32minstalled\x1b[0m");
        } else {
            println!("  \x1b[36m  wrapper\x1b[0m    \x1b[33mnot installed\x1b[0m");
        }
        println!();
    } else {
        println!("  \x1b[1mmitm\x1b[0m        \x1b[33mdisabled\x1b[0m");
        println!();
    }

    // Routes
    println!("  \x1b[1mroutes\x1b[0m");
    println!("  \x1b[33m  POST\x1b[0m  /v1/responses");
    println!("  \x1b[33m  POST\x1b[0m  /v1/chat/completions");
    println!("  \x1b[33m  POST\x1b[0m  /v1/gemini");
    println!("  \x1b[32m  GET \x1b[0m  /v1/models");
    println!("  \x1b[32m  GET \x1b[0m  /v1/sessions");
    println!("  \x1b[31m  DEL \x1b[0m  /v1/sessions/:id");
    println!("  \x1b[33m  POST\x1b[0m  /v1/token");
    println!("  \x1b[32m  GET \x1b[0m  /v1/usage");
    println!("  \x1b[32m  GET \x1b[0m  /v1/quota");
    println!("  \x1b[32m  GET \x1b[0m  /health");
    println!();

    // Status line
    let mitm_tag = if mitm.is_some() { "\x1b[32mmitm\x1b[0m" } else { "\x1b[31mmitm\x1b[0m" };
    println!("  \x1b[2mstealth:\x1b[0m \x1b[32mwarmup\x1b[0m \x1b[32mheartbeat\x1b[0m \x1b[32mjitter\x1b[0m {mitm_tag}");
    println!();

    // Setup hints
    if let Some((mitm_port, ca_path)) = mitm {
        if is_standalone {
            // Standalone mode uses iptables UID isolation — no wrapper needed
        } else if !check_wrapper_installed() {
            println!("  \x1b[1;33m[!]\x1b[0m mitm wrapper not installed");
            println!("      \x1b[2mrun:\x1b[0m  ./scripts/mitm-wrapper.sh install");
            println!("      \x1b[2mor:\x1b[0m   HTTPS_PROXY=http://127.0.0.1:{mitm_port}");
            println!("             NODE_EXTRA_CA_CERTS={ca_path}");
            println!();
        }
    }

    if token == "NOT SET" {
        println!("  \x1b[1;33m[!]\x1b[0m no oauth token");
        println!("      export ANTIGRAVITY_OAUTH_TOKEN=ya29.xxx");
        println!("      curl -X POST http://127.0.0.1:{port}/v1/token -d '{{\"token\":\"ya29.xxx\"}}'");
        println!("      echo 'ya29.xxx' > ~/.config/antigravity-proxy-token");
        println!();
    }
}

/// Check if the MITM wrapper is installed by looking for the .real backup file
/// next to the LS binary. Uses /proc to find the real LS path dynamically.
fn check_wrapper_installed() -> bool {
    // Find the LS binary path from known PID or by scanning /proc
    if let Some(ls_path) = find_ls_binary_path() {
        let real_path = format!("{ls_path}.real");
        return std::path::Path::new(&real_path).exists();
    }
    false
}

/// Find the LS binary path by reading /proc/<pid>/exe for known language server processes.
fn find_ls_binary_path() -> Option<String> {
    // Try all running processes, look for ones that look like the LS
    let proc = std::path::Path::new("/proc");
    if !proc.exists() {
        return None;
    }

    if let Ok(entries) = std::fs::read_dir(proc) {
        for entry in entries.flatten() {
            let name = entry.file_name();
            let name_str = name.to_string_lossy();
            // Only look at numeric dirs (PIDs)
            if !name_str.chars().all(|c| c.is_ascii_digit()) {
                continue;
            }
            let exe_link = entry.path().join("exe");
            if let Ok(target) = std::fs::read_link(&exe_link) {
                let target_str = target.to_string_lossy();
                // Strip " (deleted)" suffix from unlinked binaries
                let target_clean = target_str.trim_end_matches(" (deleted)");
                // Match any binary that looks like the Antigravity LS
                if target_clean.contains("language_server_linux")
                    || target_clean.contains("antigravity-language-server")
                {
                    // Strip .real suffix — if the wrapper exec'd the backup, we want the base name
                    let path = target_clean.trim_end_matches(".real");
                    return Some(path.to_string());
                }
            }
        }
    }
    None
}

/// Get the data directory for storing MITM CA cert/key.
fn dirs_data_dir() -> std::path::PathBuf {
    let home = std::env::var("HOME").unwrap_or_else(|_| "/tmp".to_string());
    std::path::PathBuf::from(home).join(".config").join("antigravity-proxy")
}