use anyhow::{Context, Result, anyhow};
use clap::Parser;
use std::path::PathBuf;
use std::process::Command;
use std::time::{Duration, Instant};

#[derive(Debug, Parser)]
#[command(name = "dosh-bench")]
struct Args {
    #[arg(long, default_value = "local")]
    server: String,
    #[arg(long, default_value = "default")]
    session: String,
    #[arg(long, default_value_t = 22)]
    ssh_port: u16,
    #[arg(long, default_value_t = 50000)]
    dosh_port: u16,
    #[arg(long)]
    dosh_host: Option<String>,
    #[arg(long, default_value_t = 3)]
    iterations: usize,
    #[arg(long)]
    local_auth: bool,
    #[arg(long)]
    client: Option<PathBuf>,
    #[arg(long, default_value = "~/.local/bin/dosh-auth")]
    ssh_auth_command: String,
    #[arg(long)]
    ssh_key: Option<PathBuf>,
    #[arg(long)]
    ssh_known_hosts: Option<PathBuf>,
    #[arg(long)]
    ssh_control_path: Option<PathBuf>,
    #[arg(long)]
    controlmaster: bool,
    #[arg(long)]
    no_cache: bool,
    #[arg(long)]
    assert_ssh_plus_ms: Option<f64>,
}

fn main() -> Result<()> {
    let args = Args::parse();
    let client = args.client.clone().unwrap_or_else(default_client_path);
    let mut ssh_times = Vec::new();
    let mut dosh_times = Vec::new();
    let generated_control_path = if args.controlmaster {
        Some(std::env::temp_dir().join(format!("dosh-bench-control-{}", std::process::id())))
    } else {
        None
    };
    let control_path = generated_control_path
        .as_ref()
        .or(args.ssh_control_path.as_ref());
    let _controlmaster = if let Some(path) = generated_control_path.as_ref() {
        Some(ControlMaster::start(&args, path.clone())?)
    } else {
        None
    };

    for _ in 0..args.iterations.max(1) {
        if !args.local_auth {
            let mut ssh = Command::new("ssh");
            add_ssh_options(&mut ssh, &args, control_path);
            ssh.arg(&args.server).arg("true");
            ssh_times.push(time_command(&mut ssh)?);
        }

        let mut cmd = Command::new(&client);
        cmd.arg("--attach-only")
            .arg("--session")
            .arg(&args.session)
            .arg("--dosh-port")
            .arg(args.dosh_port.to_string());
        if let Some(host) = &args.dosh_host {
            cmd.arg("--dosh-host").arg(host);
        }
        if args.local_auth {
            cmd.arg("--local-auth").arg(&args.server);
        } else {
            cmd.arg("--ssh-port")
                .arg(args.ssh_port.to_string())
                .arg("--ssh-auth-command")
                .arg(&args.ssh_auth_command);
            if args.no_cache {
                cmd.arg("--no-cache");
            }
            if let Some(key) = &args.ssh_key {
                cmd.arg("--ssh-key").arg(key);
            }
            if let Some(known_hosts) = &args.ssh_known_hosts {
                cmd.arg("--ssh-known-hosts").arg(known_hosts);
            }
            if let Some(control_path) = control_path {
                cmd.arg("--ssh-control-path").arg(control_path);
            }
            cmd.arg(&args.server);
        }
        dosh_times.push(time_command(&mut cmd)?);
    }

    if !ssh_times.is_empty() {
        println!(
            "ssh_true_ms avg={:.2} samples={:?}",
            avg_ms(&ssh_times),
            ssh_times
        );
    }
    println!(
        "dosh_attach_ms avg={:.2} samples={:?}",
        avg_ms(&dosh_times),
        dosh_times
    );
    if let Some(margin) = args.assert_ssh_plus_ms {
        if ssh_times.is_empty() {
            return Err(anyhow!(
                "--assert-ssh-plus-ms requires non-local SSH benchmark"
            ));
        }
        let ssh_avg = avg_ms(&ssh_times);
        let dosh_avg = avg_ms(&dosh_times);
        if dosh_avg > ssh_avg + margin {
            return Err(anyhow!(
                "dosh attach avg {dosh_avg:.2}ms exceeded ssh avg {ssh_avg:.2}ms + {margin:.2}ms"
            ));
        }
        println!("gate ok: dosh avg {dosh_avg:.2}ms <= ssh avg {ssh_avg:.2}ms + {margin:.2}ms");
    }
    Ok(())
}

fn add_ssh_options(cmd: &mut Command, args: &Args, control_path: Option<&PathBuf>) {
    cmd.arg("-p").arg(args.ssh_port.to_string()).arg("-T");
    if let Some(key) = &args.ssh_key {
        cmd.arg("-i").arg(key);
    }
    if let Some(known_hosts) = &args.ssh_known_hosts {
        cmd.arg("-o")
            .arg(format!("UserKnownHostsFile={}", known_hosts.display()));
    }
    if let Some(control_path) = control_path {
        cmd.arg("-S").arg(control_path);
    }
}

struct ControlMaster {
    server: String,
    ssh_port: u16,
    ssh_key: Option<PathBuf>,
    ssh_known_hosts: Option<PathBuf>,
    control_path: PathBuf,
}

impl ControlMaster {
    fn start(args: &Args, control_path: PathBuf) -> Result<Self> {
        let mut cmd = Command::new("ssh");
        cmd.arg("-p")
            .arg(args.ssh_port.to_string())
            .arg("-T")
            .arg("-M")
            .arg("-S")
            .arg(&control_path)
            .arg("-f")
            .arg("-N")
            .arg("-o")
            .arg("ControlPersist=60")
            .arg("-o")
            .arg("ExitOnForwardFailure=yes");
        if let Some(key) = &args.ssh_key {
            cmd.arg("-i").arg(key);
        }
        if let Some(known_hosts) = &args.ssh_known_hosts {
            cmd.arg("-o")
                .arg(format!("UserKnownHostsFile={}", known_hosts.display()));
        }
        cmd.arg(&args.server);
        time_command(&mut cmd).context("start SSH ControlMaster")?;
        Ok(Self {
            server: args.server.clone(),
            ssh_port: args.ssh_port,
            ssh_key: args.ssh_key.clone(),
            ssh_known_hosts: args.ssh_known_hosts.clone(),
            control_path,
        })
    }
}

impl Drop for ControlMaster {
    fn drop(&mut self) {
        let mut cmd = Command::new("ssh");
        cmd.arg("-p")
            .arg(self.ssh_port.to_string())
            .arg("-T")
            .arg("-S")
            .arg(&self.control_path)
            .arg("-O")
            .arg("exit");
        if let Some(key) = &self.ssh_key {
            cmd.arg("-i").arg(key);
        }
        if let Some(known_hosts) = &self.ssh_known_hosts {
            cmd.arg("-o")
                .arg(format!("UserKnownHostsFile={}", known_hosts.display()));
        }
        let _ = cmd.arg(&self.server).output();
    }
}

fn time_command(cmd: &mut Command) -> Result<Duration> {
    let start = Instant::now();
    let output = cmd.output().with_context(|| format!("run {:?}", cmd))?;
    if !output.status.success() {
        return Err(anyhow!(
            "command failed {:?}\nstdout:\n{}\nstderr:\n{}",
            cmd,
            String::from_utf8_lossy(&output.stdout),
            String::from_utf8_lossy(&output.stderr)
        ));
    }
    Ok(start.elapsed())
}

fn avg_ms(samples: &[Duration]) -> f64 {
    let total: f64 = samples.iter().map(|d| d.as_secs_f64() * 1000.0).sum();
    total / samples.len() as f64
}

fn default_client_path() -> PathBuf {
    std::env::current_exe()
        .ok()
        .and_then(|path| path.parent().map(|parent| parent.join("dosh-client")))
        .unwrap_or_else(|| PathBuf::from("dosh-client"))
}