use clap::{Parser, ValueEnum};

mod bench;
mod client;
mod relay;
mod server;
mod setup;
mod stats;

#[derive(Clone, Copy, ValueEnum)]
pub enum GeoLevel {
    /// Aggregate by country only
    Country,
    /// Aggregate by city (includes country)
    City,
}

#[derive(Parser)]
#[command(name = "layover", about = "SSH through firewalls")]
enum Cli {
    /// Run the relay server
    Serve {
        /// Address to listen on
        #[arg(long, default_value = "0.0.0.0:443")]
        listen: String,

        /// TLS certificate file (PEM)
        #[arg(long, required_unless_present = "no_tls")]
        cert: Option<String>,

        /// TLS private key file (PEM)
        #[arg(long, required_unless_present = "no_tls")]
        key: Option<String>,

        /// Disable TLS (development only)
        #[arg(long)]
        no_tls: bool,

        /// Log connection lifecycle details (incoming/outgoing linking, resolved IPs, durations)
        #[arg(long, short)]
        verbose: bool,

        /// Path to MaxMind GeoLite2 database file (.mmdb) for geographic stats
        #[arg(long)]
        geoip_db: Option<String>,

        /// Geographic aggregation level for stats
        #[arg(long, default_value = "city")]
        geo_level: GeoLevel,

        /// Path to persist stats history (JSON file, created if missing)
        #[arg(long)]
        stats_file: Option<String>,

        /// Allow relaying to private/loopback IPs (for benchmarking only)
        #[arg(long, hide = true)]
        allow_private: bool,
    },

    /// Connect to a destination through the relay (ProxyCommand mode)
    Connect {
        /// Destination in host:port format
        destination: String,

        /// Relay server address
        #[arg(long, env = "LAYOVER_SERVER", default_value = "layover.sh")]
        server: String,

        /// Disable TLS (development only)
        #[arg(long)]
        no_tls: bool,
    },

    /// Run performance benchmarks (throughput, latency, concurrency)
    Bench {
        /// Number of parallel connections for concurrency tests
        #[arg(long, default_value = "1,10,50,100")]
        concurrency: String,

        /// Bytes to push per connection for throughput tests
        #[arg(long, default_value = "10485760")]
        bytes: u64,

        /// Number of round-trips for latency tests
        #[arg(long, default_value = "1000")]
        latency_rounds: u64,

        /// Save results to this JSON file
        #[arg(long)]
        output: Option<String>,
    },

    /// Run echo server for benchmarking (internal use)
    #[command(hide = true)]
    BenchEcho {
        #[arg(long)]
        listen: String,
    },

    /// Generate SSH config to route connections through Layover
    Setup {
        /// Relay server address
        #[arg(long, env = "LAYOVER_SERVER", default_value = "layover.sh")]
        server: String,

        /// SSH host pattern(s) to route through Layover (default: * with relay exclusion)
        #[arg(long)]
        host: Vec<String>,

        /// Append config to ~/.ssh/config instead of printing
        #[arg(long)]
        install: bool,

        /// Pass --no-tls in the ProxyCommand (development only)
        #[arg(long)]
        no_tls: bool,
    },
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let cli = Cli::parse();
    match cli {
        Cli::Serve {
            listen,
            cert,
            key,
            no_tls,
            verbose,
            geoip_db,
            geo_level,
            stats_file,
            allow_private,
        } => server::run(&listen, cert.as_deref(), key.as_deref(), no_tls, verbose, geoip_db.as_deref(), geo_level, stats_file.as_deref(), allow_private).await,
        Cli::Bench {
            concurrency,
            bytes,
            latency_rounds,
            output,
        } => {
            let levels: Vec<usize> = concurrency.split(',').map(|s| s.trim().parse().expect("invalid concurrency value")).collect();
            bench::run(&levels, bytes, latency_rounds, output.as_deref()).await
        }
        Cli::BenchEcho { listen } => bench::run_echo(&listen).await,
        Cli::Connect {
            destination,
            server,
            no_tls,
        } => client::run(&destination, &server, no_tls).await,
        Cli::Setup {
            server,
            host,
            install,
            no_tls,
        } => setup::run(&server, &host, install, no_tls),
    }
}
