use anyhow::{Context, Result, anyhow};
use clap::Parser;
use portable_pty::{CommandBuilder, NativePtySystem, PtySize, PtySystem};
use std::ffi::OsStr;
use std::io::Read;
use std::path::PathBuf;
use std::process::Command;
use std::sync::mpsc;
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)]
    warm_cache: bool,
    #[arg(long)]
    skip_ssh_baseline: bool,
    #[arg(long)]
    include_mosh: bool,
    #[arg(long, default_value = "mosh")]
    mosh: PathBuf,
    #[arg(long, default_value = "mosh-server")]
    mosh_server_command: String,
    #[arg(long)]
    mosh_port: Option<String>,
    #[arg(long)]
    assert_ssh_plus_ms: Option<f64>,
    #[arg(long)]
    assert_mosh_minus_ms: Option<f64>,
    #[arg(long)]
    assert_dosh_max_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 mut mosh_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
    };
    if args.no_cache && args.warm_cache {
        return Err(anyhow!("--warm-cache cannot be used with --no-cache"));
    }
    let dosh_label = if args.warm_cache {
        let _ = time_dosh_attach(&client, &args, control_path)?;
        "dosh_cached_attach_ms"
    } else {
        "dosh_attach_ms"
    };

    for _ in 0..args.iterations.max(1) {
        if !args.local_auth && !args.skip_ssh_baseline {
            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)?);
        }

        dosh_times.push(time_dosh_attach(&client, &args, control_path)?);

        if args.include_mosh {
            mosh_times.push(time_mosh_in_pty(&args)?);
        }
    }

    if !ssh_times.is_empty() {
        println!(
            "ssh_true_ms avg={:.2} samples={:?}",
            avg_ms(&ssh_times),
            ssh_times
        );
    }
    println!(
        "{dosh_label} avg={:.2} samples={:?}",
        avg_ms(&dosh_times),
        dosh_times
    );
    if !mosh_times.is_empty() {
        println!(
            "mosh_start_true_ms avg={:.2} samples={:?}",
            avg_ms(&mosh_times),
            mosh_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");
    }
    if let Some(margin) = args.assert_mosh_minus_ms {
        if mosh_times.is_empty() {
            return Err(anyhow!(
                "--assert-mosh-minus-ms requires --include-mosh benchmark"
            ));
        }
        let dosh_avg = avg_ms(&dosh_times);
        let mosh_avg = avg_ms(&mosh_times);
        if dosh_avg + margin > mosh_avg {
            return Err(anyhow!(
                "dosh attach avg {dosh_avg:.2}ms was not at least {margin:.2}ms faster than mosh avg {mosh_avg:.2}ms"
            ));
        }
        println!("gate ok: dosh avg {dosh_avg:.2}ms + {margin:.2}ms <= mosh avg {mosh_avg:.2}ms");
    }
    if let Some(max_ms) = args.assert_dosh_max_ms {
        let dosh_avg = avg_ms(&dosh_times);
        if dosh_avg > max_ms {
            return Err(anyhow!(
                "{dosh_label} avg {dosh_avg:.2}ms exceeded max {max_ms:.2}ms"
            ));
        }
        println!("gate ok: {dosh_label} avg {dosh_avg:.2}ms <= {max_ms:.2}ms");
    }
    Ok(())
}

fn time_dosh_attach(
    client: &PathBuf,
    args: &Args,
    control_path: Option<&PathBuf>,
) -> Result<Duration> {
    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);
    }
    time_command(&mut cmd)
}

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);
    }
}

fn mosh_ssh_command(args: &Args) -> String {
    let mut parts = vec![
        "ssh".to_string(),
        "-p".to_string(),
        args.ssh_port.to_string(),
    ];
    if let Some(key) = &args.ssh_key {
        parts.push("-i".to_string());
        parts.push(shell_word(key.as_os_str()));
    }
    if let Some(known_hosts) = &args.ssh_known_hosts {
        parts.push("-o".to_string());
        parts.push(shell_word(OsStr::new(&format!(
            "UserKnownHostsFile={}",
            known_hosts.display()
        ))));
    }
    parts.join(" ")
}

fn time_mosh_in_pty(args: &Args) -> Result<Duration> {
    let pty_system = NativePtySystem::default();
    let pair = pty_system
        .openpty(PtySize {
            rows: 24,
            cols: 80,
            pixel_width: 0,
            pixel_height: 0,
        })
        .context("open pty for mosh benchmark")?;
    let mut command_line = vec![
        shell_word(args.mosh.as_os_str()),
        shell_word(OsStr::new(&format!("--ssh={}", mosh_ssh_command(args)))),
        shell_word(OsStr::new(&format!(
            "--server={}",
            args.mosh_server_command
        ))),
        "--no-init".to_string(),
    ];
    if let Some(port) = &args.mosh_port {
        command_line.push("-p".to_string());
        command_line.push(shell_word(OsStr::new(port)));
    }
    command_line.push(shell_word(OsStr::new(&args.server)));
    command_line.push("--".to_string());
    command_line.push("true".to_string());

    let mut cmd = CommandBuilder::new("/bin/sh");
    cmd.arg("-c");
    cmd.arg(command_line.join(" "));

    let start = Instant::now();
    let mut child = pair
        .slave
        .spawn_command(cmd)
        .context("spawn mosh benchmark in pty")?;
    drop(pair.slave);

    let mut reader = pair.master.try_clone_reader().context("clone pty reader")?;
    let (tx, rx) = mpsc::channel();
    std::thread::spawn(move || {
        let mut output = Vec::new();
        let _ = reader.read_to_end(&mut output);
        let _ = tx.send(output);
    });

    {
        let _writer = pair.master.take_writer().context("take pty writer")?;
    }

    let status = child.wait().context("wait for mosh benchmark")?;
    let elapsed = start.elapsed();
    drop(pair.master);
    let output = rx.recv().unwrap_or_default();
    if !status.success() {
        return Err(anyhow!(
            "mosh command failed with status {status:?}\noutput:\n{}",
            String::from_utf8_lossy(&output)
        ));
    }
    Ok(elapsed)
}

fn shell_word(value: &OsStr) -> String {
    let value = value.to_string_lossy();
    format!("'{}'", value.replace('\'', "'\\''"))
}

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"))
}