use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::time::Duration;

use anyhow::{Result, bail};

use crate::instance_store;
use crate::process;

/// SSH client binary. On macOS the nix-store `ssh` is blocked from the
/// guest's vmnet subnet by Local Network Privacy (connect → EHOSTUNREACH);
/// Apple's `/usr/bin/ssh` carries the local-network entitlement, so use the
/// system tools there. On Linux, resolve from PATH as usual.
#[cfg(target_os = "macos")]
pub const SSH_PROGRAM: &str = "/usr/bin/ssh";
#[cfg(not(target_os = "macos"))]
pub const SSH_PROGRAM: &str = "ssh";

/// `ssh-keyscan` binary — same Local Network Privacy constraint as [`SSH_PROGRAM`].
#[cfg(target_os = "macos")]
pub const SSH_KEYSCAN_PROGRAM: &str = "/usr/bin/ssh-keyscan";
#[cfg(not(target_os = "macos"))]
pub const SSH_KEYSCAN_PROGRAM: &str = "ssh-keyscan";

/// How often `wait_for_ssh` retries. Kept well under a second so the ready
/// time epi reports tracks the guest's actual boot rather than rounding up to
/// the next poll.
const SSH_POLL_INTERVAL: Duration = Duration::from_millis(200);

pub fn user() -> String {
    std::env::var("USER").unwrap_or_else(|_| "user".to_string())
}

/// Returns the path to the SSH config file for an instance.
pub fn config_path(instance: &str) -> PathBuf {
    instance_store::instance_path(instance, "ssh_config")
}

/// Generate an SSH config file for an instance and write it to the given path.
///
/// When `known_hosts` is `None`, host key checking is disabled (for initial boot).
/// When `Some(path)`, strict host key checking is enabled using the given known_hosts file.
pub fn generate_config(
    output: &Path,
    instance: &str,
    ssh: SocketAddr,
    username: &str,
    ssh_key_path: &Path,
    known_hosts: Option<&Path>,
    extra_config: &[String],
) -> Result<()> {
    let (strict_checking, known_hosts_file) = match known_hosts {
        Some(path) => ("yes", path.display().to_string()),
        None => ("no", "/dev/null".to_string()),
    };
    let mut contents = format!(
        "Host {instance}\n\
         \x20   HostName {host}\n\
         \x20   Port {port}\n\
         \x20   User {username}\n\
         \x20   IdentityFile {key}\n\
         \x20   IdentitiesOnly yes\n\
         \x20   StrictHostKeyChecking {strict_checking}\n\
         \x20   UserKnownHostsFile {known_hosts_file}\n\
         \x20   LogLevel ERROR\n",
        host = ssh.ip(),
        port = ssh.port(),
        key = ssh_key_path.display(),
    );
    for line in extra_config {
        contents.push_str(&format!("    {line}\n"));
    }
    std::fs::write(output, &contents)?;
    Ok(())
}

/// Returns the path to the known_hosts file for an instance.
pub fn known_hosts_path(instance: &str) -> PathBuf {
    instance_store::instance_path(instance, "known_hosts")
}

/// Record the guest's SSH host key using ssh-keyscan.
///
/// Returns `true` if the key was successfully recorded, `false` if keyscan failed.
/// Does not return an error on keyscan failure — the caller can fall back to the
/// untrusted config.
pub fn record_host_key(ssh: SocketAddr, known_hosts: &Path) -> Result<bool> {
    let port_str = ssh.port().to_string();
    let host_str = ssh.ip().to_string();
    let out = process::run(SSH_KEYSCAN_PROGRAM, &["-p", &port_str, &host_str])?;

    if !out.success() || out.stdout.is_empty() {
        return Ok(false);
    }

    std::fs::write(known_hosts, &out.stdout)?;
    Ok(true)
}

/// Record the guest host key and rewrite the SSH config to enable strict checking.
///
/// If keyscan fails, logs a warning and leaves the untrusted config in place.
pub fn trust_host_key(
    instance: &str,
    ssh: SocketAddr,
    username: &str,
    ssh_key_path: &Path,
    extra_config: &[String],
) -> Result<()> {
    let known_hosts = known_hosts_path(instance);
    if record_host_key(ssh, &known_hosts)? {
        generate_config(
            &config_path(instance),
            instance,
            ssh,
            username,
            ssh_key_path,
            Some(&known_hosts),
            extra_config,
        )?;
    } else {
        eprintln!(
            "warning: could not record host key for {instance}, keeping untrusted SSH config"
        );
    }
    Ok(())
}

/// Copy a nix store closure to a running instance via SSH.
pub fn nix_copy_closure(instance: &str, store_path: &str) -> Result<()> {
    let config = config_path(instance);
    // nix invokes ssh from PATH; NIX_SSHOPTS only carries options. On macOS
    // point it at the entitled system ssh via an absolute -o is not possible,
    // so prepend /usr/bin to PATH so `nix copy` picks up /usr/bin/ssh.
    let ssh_opts = format!("-F {}", config.display());
    let target = format!("ssh://{instance}");
    let mut env: Vec<(&str, &str)> = vec![("NIX_SSHOPTS", &ssh_opts)];
    let patched_path;
    if cfg!(target_os = "macos") {
        let existing = std::env::var("PATH").unwrap_or_default();
        patched_path = format!("/usr/bin:{existing}");
        env.push(("PATH", &patched_path));
    }
    let out = process::run_with_env("nix", &["copy", "--to", &target, store_path], &env)?;
    if !out.success() {
        bail!("nix copy failed (exit {}): {}", out.status, out.stderr);
    }
    Ok(())
}

/// Run a command on a running instance via SSH.
pub fn run_on_guest(instance: &str, command: &str) -> Result<()> {
    let config = config_path(instance);
    let config_str = config.to_string_lossy();
    let out = process::run(SSH_PROGRAM, &["-F", &config_str, instance, "--", command])?;
    if !out.success() {
        bail!(
            "remote command failed (exit {}): {}",
            out.status,
            out.stderr
        );
    }
    Ok(())
}

/// Poll until SSH is reachable via the config file, or timeout.
pub fn wait_for_ssh(config: &Path, instance: &str, timeout_seconds: u64) -> Result<()> {
    let start = std::time::Instant::now();
    let timeout = Duration::from_secs(timeout_seconds);
    let config_str = config.to_string_lossy();

    loop {
        if start.elapsed() >= timeout {
            bail!("SSH not reachable after {timeout_seconds}s — instance may still be booting");
        }

        let out = process::run(
            SSH_PROGRAM,
            &[
                "-F",
                &config_str,
                "-o",
                "BatchMode=yes",
                "-o",
                "ConnectTimeout=5",
                instance,
                "true",
            ],
        )?;

        if out.success() {
            return Ok(());
        }

        std::thread::sleep(SSH_POLL_INTERVAL);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn ssh_poll_interval_is_finer_than_boot_time_differences() {
        // A coarse interval quantises the ready time epi reports, hiding
        // guest boot improvements smaller than one poll.
        assert!(SSH_POLL_INTERVAL <= Duration::from_millis(250));
    }

    #[test]
    fn ssh_programs_are_system_tools_on_macos() {
        if cfg!(target_os = "macos") {
            assert_eq!(SSH_PROGRAM, "/usr/bin/ssh");
            assert_eq!(SSH_KEYSCAN_PROGRAM, "/usr/bin/ssh-keyscan");
        } else {
            assert_eq!(SSH_PROGRAM, "ssh");
            assert_eq!(SSH_KEYSCAN_PROGRAM, "ssh-keyscan");
        }
    }

    #[test]
    fn test_generate_config_untrusted() {
        let dir = tempfile::tempdir().unwrap();
        let config = dir.path().join("ssh_config");
        let key_path = Path::new("/home/adam/.epi/state/myvm/id_ed25519");

        generate_config(
            &config,
            "myvm",
            "127.0.0.1:12345".parse().unwrap(),
            "adam",
            key_path,
            None,
            &[],
        )
        .unwrap();

        let contents = std::fs::read_to_string(&config).unwrap();
        assert!(contents.starts_with("Host myvm\n"));
        assert!(contents.contains("HostName 127.0.0.1"));
        assert!(contents.contains("Port 12345"));
        assert!(contents.contains("User adam"));
        assert!(contents.contains(&format!("IdentityFile {}", key_path.display())));
        assert!(contents.contains("IdentitiesOnly yes"));
        assert!(contents.contains("StrictHostKeyChecking no"));
        assert!(contents.contains("UserKnownHostsFile /dev/null"));
        assert!(contents.contains("LogLevel ERROR"));

        // Verify it's valid SSH config format - all options indented under Host
        let lines: Vec<&str> = contents.lines().collect();
        assert!(!lines[0].starts_with(' '));
        for line in &lines[1..] {
            assert!(line.starts_with("    "));
        }
    }

    #[test]
    fn test_generate_config_never_contains_remote_command() {
        let dir = tempfile::tempdir().unwrap();
        let config = dir.path().join("ssh_config");
        let key_path = Path::new("/home/adam/.epi/state/myvm/id_ed25519");

        generate_config(
            &config,
            "myvm",
            "127.0.0.1:12345".parse().unwrap(),
            "adam",
            key_path,
            None,
            &[],
        )
        .unwrap();

        let contents = std::fs::read_to_string(&config).unwrap();
        assert!(!contents.contains("RemoteCommand"));
        assert!(!contents.contains("RequestTTY"));
    }

    #[test]
    fn test_generate_config_with_extra_lines() {
        let dir = tempfile::tempdir().unwrap();
        let config = dir.path().join("ssh_config");
        let key_path = Path::new("/home/adam/.epi/state/myvm/id_ed25519");

        let extra = vec![
            "LocalForward /tmp/local.sock /tmp/remote.sock".to_string(),
            "ForwardAgent yes".to_string(),
        ];
        generate_config(
            &config,
            "myvm",
            "127.0.0.1:12345".parse().unwrap(),
            "adam",
            key_path,
            None,
            &extra,
        )
        .unwrap();

        let contents = std::fs::read_to_string(&config).unwrap();
        assert!(contents.contains("    LocalForward /tmp/local.sock /tmp/remote.sock\n"));
        assert!(contents.contains("    ForwardAgent yes\n"));

        // Extra lines should come after the standard config
        let local_fwd_pos = contents.find("LocalForward").unwrap();
        let log_level_pos = contents.find("LogLevel").unwrap();
        assert!(local_fwd_pos > log_level_pos);
    }

    #[test]
    fn test_generate_config_with_empty_extra_lines() {
        let dir = tempfile::tempdir().unwrap();
        let config = dir.path().join("ssh_config");
        let key_path = Path::new("/home/adam/.epi/state/myvm/id_ed25519");

        generate_config(
            &config,
            "myvm",
            "127.0.0.1:12345".parse().unwrap(),
            "adam",
            key_path,
            None,
            &[],
        )
        .unwrap();

        let contents = std::fs::read_to_string(&config).unwrap();
        // Should be identical to the no-extra-config case
        assert!(contents.ends_with("LogLevel ERROR\n"));
    }

    #[test]
    fn test_generate_config_writes_supplied_host() {
        let dir = tempfile::tempdir().unwrap();
        let config = dir.path().join("ssh_config");
        let key_path = Path::new("/home/adam/.epi/state/myvm/id_ed25519");

        generate_config(
            &config,
            "myvm",
            "192.168.64.42:22".parse().unwrap(),
            "adam",
            key_path,
            None,
            &[],
        )
        .unwrap();

        let contents = std::fs::read_to_string(&config).unwrap();
        assert!(contents.contains("HostName 192.168.64.42"));
        assert!(contents.contains("Port 22"));
        assert!(!contents.contains("HostName 127.0.0.1"));
    }

    #[test]
    fn test_generate_config_writes_ipv6_host() {
        let dir = tempfile::tempdir().unwrap();
        let config = dir.path().join("ssh_config");
        let key_path = Path::new("/home/adam/.epi/state/myvm/id_ed25519");

        generate_config(
            &config,
            "myvm",
            "[::1]:22".parse().unwrap(),
            "adam",
            key_path,
            None,
            &[],
        )
        .unwrap();

        let contents = std::fs::read_to_string(&config).unwrap();
        assert!(contents.contains("HostName ::1"));
        assert!(contents.contains("Port 22"));
    }

    #[test]
    fn test_generate_config_trusted() {
        let dir = tempfile::tempdir().unwrap();
        let config = dir.path().join("ssh_config");
        let key_path = Path::new("/home/adam/.epi/state/myvm/id_ed25519");
        let known_hosts = dir.path().join("known_hosts");

        generate_config(
            &config,
            "myvm",
            "127.0.0.1:12345".parse().unwrap(),
            "adam",
            key_path,
            Some(&known_hosts),
            &[],
        )
        .unwrap();

        let contents = std::fs::read_to_string(&config).unwrap();
        assert!(contents.contains("StrictHostKeyChecking yes"));
        assert!(contents.contains(&format!("UserKnownHostsFile {}", known_hosts.display())));
        // Should NOT contain the untrusted values
        assert!(!contents.contains("StrictHostKeyChecking no"));
        assert!(!contents.contains("UserKnownHostsFile /dev/null"));
    }
}