//! Host-owned SSH transport for remote worker execution and bootstrap execution.
//!
//! Enforces:
//! 1. Genuine SSH transport with authentic known-hosts, keys, and target login shell.
//! 2. Separation of SSH stderr into a background-drained bounded buffer so verbose
//!    banners, warnings, or connection errors cannot deadlock the channel.
//! 3. Host ownership of all transport process creation (no unmanaged child SSH loops).
//! 4. Bounded execution primitives for remote bootstrap and smoke verification.

use std::{path::PathBuf, process::Stdio, sync::Arc, time::Duration};
use tokio::{
    io::AsyncReadExt,
    process::{Child, ChildStderr, ChildStdin, ChildStdout, Command},
    sync::{watch, Mutex},
};

pub const MAX_STDERR_TAIL_BYTES: usize = 65_536;

/// SSH host key checking policy.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum StrictHostKeyChecking {
    /// Strict checking against known_hosts: connection fails if host key is missing or changed.
    Yes,
    /// Automatically records new host keys on first connect, but strictly rejects changed keys.
    AcceptNew,
    /// Insecure: disables host key verification completely (for testing only).
    No,
}

/// SSH configuration for connecting to a target.
#[derive(Debug, Clone)]
pub struct SshConfig {
    pub host: String,
    pub port: Option<u16>,
    pub user: Option<String>,
    pub identity_file: Option<PathBuf>,
    pub known_hosts_path: Option<PathBuf>,
    pub strict_host_key_checking: StrictHostKeyChecking,
    pub extra_args: Vec<String>,
}

impl SshConfig {
    pub fn new(host: impl Into<String>) -> Self {
        Self {
            host: host.into(),
            port: None,
            user: None,
            identity_file: None,
            known_hosts_path: None,
            strict_host_key_checking: StrictHostKeyChecking::AcceptNew,
            extra_args: Vec::new(),
        }
    }

    pub fn with_user(mut self, user: impl Into<String>) -> Self {
        self.user = Some(user.into());
        self
    }

    pub fn with_port(mut self, port: u16) -> Self {
        self.port = Some(port);
        self
    }

    pub fn with_identity_file(mut self, path: impl Into<PathBuf>) -> Self {
        self.identity_file = Some(path.into());
        self
    }

    pub fn with_known_hosts(mut self, path: impl Into<PathBuf>) -> Self {
        self.known_hosts_path = Some(path.into());
        self
    }

    pub fn with_strict_host_key_checking(mut self, policy: StrictHostKeyChecking) -> Self {
        self.strict_host_key_checking = policy;
        self
    }
}

/// An active SSH child process with isolated stdin/stdout streams and a separately-drained stderr tail.
pub struct SshChild {
    pub stdin: ChildStdin,
    pub stdout: ChildStdout,
    pub process: SshProcess,
}

/// The retained process and diagnostic handle after stdin/stdout are handed to the multiplexer.
pub struct SshProcess {
    pub stderr_tail: Arc<Mutex<Vec<u8>>>,
    child: Child,
    stderr_drained: watch::Receiver<bool>,
}

impl SshProcess {
    /// Kills the SSH transport child process.
    pub async fn kill(&mut self) -> Result<(), std::io::Error> {
        self.child.kill().await
    }

    /// Awaits the SSH transport and its diagnostic stream.
    pub async fn wait(&mut self) -> Result<std::process::ExitStatus, std::io::Error> {
        let status = self.child.wait().await;
        let _ = self.stderr_drained.wait_for(|drained| *drained).await;
        status
    }

    /// Retrieves the current drained stderr tail as string.
    pub async fn stderr_string(&self) -> String {
        let guard = self.stderr_tail.lock().await;
        String::from_utf8_lossy(&guard).to_string()
    }
}

impl SshChild {
    pub fn split(self) -> (ChildStdin, ChildStdout, SshProcess) {
        (self.stdin, self.stdout, self.process)
    }

    /// Kills the SSH transport child process.
    pub async fn kill(&mut self) -> Result<(), std::io::Error> {
        self.process.kill().await
    }

    /// Awaits the SSH transport process exit.
    pub async fn wait(&mut self) -> Result<std::process::ExitStatus, std::io::Error> {
        self.process.wait().await
    }

    /// Retrieves the current drained stderr tail as string.
    pub async fn stderr_string(&self) -> String {
        self.process.stderr_string().await
    }
}

/// A remote command whose three streams are handed to the caller UNDRAINED.
///
/// This is the evidence-complete counterpart to [`SshChild`]. `SshChild` drains stderr
/// into a bounded tail in the background, which guarantees the channel cannot deadlock
/// but makes complete stderr unrecoverable. When a caller must retain *complete*
/// output as retrievable evidence, it takes ownership of the raw streams instead.
///
/// # Hazard
///
/// The caller MUST drain `stderr` concurrently with `stdout`. A remote process that
/// fills its stderr pipe while the host reads only stdout will block forever. The
/// bounded background drain that normally prevents this is deliberately absent here,
/// so that responsibility transfers with the streams. The caller is likewise
/// responsible for bounding what it retains and for reaping the child.
pub struct SshStreams {
    pub stdin: ChildStdin,
    pub stdout: ChildStdout,
    pub stderr: ChildStderr,
    pub process: SshStreamingProcess,
}

/// Process control for a streaming remote command.
///
/// Deliberately does NOT expose `stderr_string()`: in streaming mode no background
/// tail is collected, so such an accessor could only ever return an empty buffer and
/// would misrepresent absent evidence as observed silence.
pub struct SshStreamingProcess {
    child: Child,
}

impl SshStreamingProcess {
    pub async fn kill(&mut self) -> Result<(), std::io::Error> {
        self.child.kill().await
    }

    pub async fn wait(&mut self) -> Result<std::process::ExitStatus, std::io::Error> {
        self.child.wait().await
    }
}

/// Result of a bounded remote command execution.
#[derive(Debug, Clone)]
pub struct RemoteExecResult {
    pub exit_code: i32,
    pub stdout: Vec<u8>,
    pub stderr: Vec<u8>,
    pub stdout_truncated: bool,
    pub stderr_truncated: bool,
}

pub struct SshTransport;

impl SshTransport {
    /// Escapes an argument using POSIX shell single-quote wrapping.
    pub fn shell_escape(arg: &str) -> String {
        if arg.is_empty() {
            return "''".to_string();
        }
        if arg
            .bytes()
            .all(|b| b.is_ascii_alphanumeric() || b"-_./:=+@".contains(&b))
        {
            return arg.to_string();
        }
        format!("'{}'", arg.replace('\'', "'\\''"))
    }

    /// Builds the SSH argument list for launching a remote command.
    pub fn build_args(config: &SshConfig, remote_cmd: &str) -> Vec<String> {
        let host_key_check = match config.strict_host_key_checking {
            StrictHostKeyChecking::Yes => "yes",
            StrictHostKeyChecking::AcceptNew => "accept-new",
            StrictHostKeyChecking::No => "no",
        };

        let mut args = vec![
            "-o".to_string(),
            "BatchMode=yes".to_string(),
            "-o".to_string(),
            format!("StrictHostKeyChecking={host_key_check}"),
        ];

        if let Some(port) = config.port {
            args.push("-p".to_string());
            args.push(port.to_string());
        }

        if let Some(ref user) = config.user {
            args.push("-l".to_string());
            args.push(user.clone());
        }

        if let Some(ref key) = config.identity_file {
            args.push("-i".to_string());
            args.push(key.display().to_string());
        }

        if let Some(ref kh) = config.known_hosts_path {
            args.push("-o".to_string());
            args.push(format!("UserKnownHostsFile={}", kh.display()));
        }

        for extra in &config.extra_args {
            args.push(extra.clone());
        }

        args.push(config.host.clone());
        args.push(remote_cmd.to_string());

        args
    }

    /// Spawns an SSH child process executing `remote_cmd` on the target host.
    pub fn spawn(config: &SshConfig, remote_cmd: &str) -> Result<SshChild, std::io::Error> {
        let args = Self::build_args(config, remote_cmd);

        let mut cmd = Command::new("ssh");
        cmd.args(&args);
        cmd.stdin(Stdio::piped());
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());

        let mut child = cmd.spawn()?;
        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| std::io::Error::other("missing child stdin"))?;
        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| std::io::Error::other("missing child stdout"))?;
        let mut stderr = child
            .stderr
            .take()
            .ok_or_else(|| std::io::Error::other("missing child stderr"))?;

        let stderr_tail = Arc::new(Mutex::new(Vec::new()));
        let tail_clone = stderr_tail.clone();
        let (stderr_drained_tx, stderr_drained) = watch::channel(false);

        // Background task to drain SSH stderr separately so banners or error messages
        // never block the pipe or intermingle with stdout.
        tokio::spawn(async move {
            let mut buf = [0u8; 4096];
            loop {
                match stderr.read(&mut buf).await {
                    Ok(0) => break,
                    Ok(n) => {
                        let mut guard = tail_clone.lock().await;
                        guard.extend_from_slice(&buf[..n]);
                        if guard.len() > MAX_STDERR_TAIL_BYTES {
                            let overflow = guard.len() - MAX_STDERR_TAIL_BYTES;
                            guard.drain(0..overflow);
                        }
                    }
                    Err(_) => break,
                }
            }
            let _ = stderr_drained_tx.send(true);
        });

        Ok(SshChild {
            stdin,
            stdout,
            process: SshProcess {
                stderr_tail,
                child,
                stderr_drained,
            },
        })
    }

    /// Executes an explicit argv command on the remote target with bounded output capture and timeout.
    ///
    /// Avoids unescaped shell string interpolation: escapes each element of `argv` safely.
    /// Bounded at `max_output_bytes` to prevent runaway memory allocation.
    /// Spawns an explicit argv command and hands the caller all three UNDRAINED streams.
    ///
    /// Escaping is identical to [`Self::spawn_argv`] and [`Self::exec_argv`]. Unlike those,
    /// no background stderr drain is installed and no output bound is applied, so complete
    /// stdout AND stderr can be retained as retrievable evidence.
    ///
    /// # Hazard
    ///
    /// The caller MUST drain stderr concurrently with stdout or the remote process can block
    /// on a full pipe. See [`SshStreams`].
    pub fn spawn_argv_streaming(
        config: &SshConfig,
        argv: &[&str],
    ) -> Result<SshStreams, std::io::Error> {
        if argv.is_empty() {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "empty argv for remote spawn",
            ));
        }

        let escaped_cmd = argv
            .iter()
            .map(|a| Self::shell_escape(a))
            .collect::<Vec<_>>()
            .join(" ");
        let args = Self::build_args(config, &escaped_cmd);

        let mut cmd = Command::new("ssh");
        cmd.args(&args);
        cmd.stdin(Stdio::piped());
        cmd.stdout(Stdio::piped());
        cmd.stderr(Stdio::piped());

        let mut child = cmd.spawn()?;
        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| std::io::Error::other("missing child stdin"))?;
        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| std::io::Error::other("missing child stdout"))?;
        let stderr = child
            .stderr
            .take()
            .ok_or_else(|| std::io::Error::other("missing child stderr"))?;

        Ok(SshStreams {
            stdin,
            stdout,
            stderr,
            process: SshStreamingProcess { child },
        })
    }

    /// Spawns an explicit argv command on the remote target, leaving stdin/stdout/stderr open.
    ///
    /// Identical escaping to [`Self::exec_argv`]: every element of `argv` is POSIX
    /// single-quote escaped and joined, so there is no shell interpolation of
    /// user-supplied values beyond per-argument escaping. Unlike `exec_argv` this
    /// performs no capture, bounding or timeout - the caller owns the streams and is
    /// responsible for draining, bounding and reaping them.
    pub fn spawn_argv(config: &SshConfig, argv: &[&str]) -> Result<SshChild, std::io::Error> {
        if argv.is_empty() {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "empty argv for remote spawn",
            ));
        }

        let escaped_cmd = argv
            .iter()
            .map(|a| Self::shell_escape(a))
            .collect::<Vec<_>>()
            .join(" ");

        Self::spawn(config, &escaped_cmd)
    }

    pub async fn exec_argv(
        config: &SshConfig,
        argv: &[&str],
        max_output_bytes: usize,
        timeout_duration: Duration,
    ) -> Result<RemoteExecResult, std::io::Error> {
        if argv.is_empty() {
            return Err(std::io::Error::new(
                std::io::ErrorKind::InvalidInput,
                "empty argv for remote exec",
            ));
        }

        let escaped_cmd = argv
            .iter()
            .map(|a| Self::shell_escape(a))
            .collect::<Vec<_>>()
            .join(" ");

        let child = Self::spawn(config, &escaped_cmd)?;

        let (_stdin, mut stdout, mut process) = child.split();
        let stderr_tail = process.stderr_tail.clone();

        let run_fut = async {
            let mut stdout_buf = Vec::new();
            let mut stdout_truncated = false;
            let mut chunk = [0u8; 4096];

            loop {
                let n = stdout.read(&mut chunk).await?;
                if n == 0 {
                    break;
                }
                if stdout_buf.len() < max_output_bytes {
                    let available = max_output_bytes - stdout_buf.len();
                    let to_take = n.min(available);
                    stdout_buf.extend_from_slice(&chunk[..to_take]);
                    if to_take < n {
                        stdout_truncated = true;
                    }
                } else {
                    stdout_truncated = true;
                }
            }

            let status = process.wait().await?;
            let stderr_bytes = {
                let guard = stderr_tail.lock().await;
                guard.clone()
            };
            let stderr_truncated = stderr_bytes.len() >= MAX_STDERR_TAIL_BYTES;

            Ok(RemoteExecResult {
                exit_code: status.code().unwrap_or(-1),
                stdout: stdout_buf,
                stderr: stderr_bytes,
                stdout_truncated,
                stderr_truncated,
            })
        };

        match tokio::time::timeout(timeout_duration, run_fut).await {
            Ok(res) => res,
            Err(_) => {
                let _ = process.kill().await;
                Err(std::io::Error::new(
                    std::io::ErrorKind::TimedOut,
                    "remote exec timed out",
                ))
            }
        }
    }

    /// Convenience execution helper for bootstrap routines returning (ExitStatus, stdout, stderr_string).
    pub async fn execute(
        config: &SshConfig,
        remote_cmd: &str,
    ) -> Result<(std::process::ExitStatus, Vec<u8>, String), std::io::Error> {
        let mut child = Self::spawn(config, remote_cmd)?;
        let mut stdout_bytes = Vec::new();
        let _ = child.stdout.read_to_end(&mut stdout_bytes).await?;
        let status = child.wait().await?;
        let stderr_str = child.stderr_string().await;
        Ok((status, stdout_bytes, stderr_str))
    }
}