use log::{info, warn};
use std::{
    fs, io::Write, os::unix::fs::PermissionsExt, process::Command, process::Stdio, sync::Arc,
    thread,
};

use super::retry_or_ignore;
use dkl::bootstrap::{Config, SSHServer};

pub async fn start(cfg: &Config) {
    retry_or_ignore(async || {
        info!("ssh: writing authorized keys");

        let ssh_dir = "/root/.ssh";
        let authorized_keys = format!("{ssh_dir}/authorized_keys");

        fs::create_dir_all(ssh_dir)?;
        fs::set_permissions(ssh_dir, fs::Permissions::from_mode(0o700))?;

        let mut ak = Vec::new();

        for auth in &cfg.auths {
            let Some(ref key) = auth.ssh_key else {
                continue;
            };
            writeln!(ak, "{key} {}", auth.name)?;
        }

        fs::write(authorized_keys, ak)?;
        Ok(())
    })
    .await;

    let cfg = Arc::new(cfg.ssh.clone());
    retry_or_ignore(async move || {
        // don't pre-start sshd as it should rarely be useful at this stage, use inetd-style.
        let listen_addr = cfg.listen.clone();
        info!("ssh: starting listener on {listen_addr}");

        let listener = tokio::net::TcpListener::bind(listen_addr).await?;

        let cfg = cfg.clone();
        let listener = listener.into_std().unwrap();
        thread::spawn(move || handle_ssh_connections(listener, &cfg));

        Ok(())
    })
    .await;
}

fn handle_ssh_connections(listener: std::net::TcpListener, cfg: &SSHServer) {
    let mut sshd_args = Vec::new();

    sshd_args.extend(["-i", "-E", "/var/log/sshd.log"]);

    let mut options = Vec::new();

    if let Some(ref user_ca) = cfg.user_ca {
        options.push(format!("TrustedUserCAKeys={user_ca}"));
    }

    for opt in &options {
        sshd_args.extend(["-o", opt]);
    }

    let mut keygen_done = false;

    loop {
        let (stream, remote) = match listener.accept() {
            Ok(v) => v,
            Err(e) => {
                warn!("ssh: listener stopped: {e}");
                return;
            }
        };

        if !keygen_done {
            // make sure we have ssh host keys even if not provided
            if (Command::new("ssh-keygen").arg("-A").status())
                .inspect_err(|e| warn!("ssh-keygen failed: {e}"))
                .is_ok_and(|s| s.success())
            {
                keygen_done = true
            }
        }

        use std::os::unix::io::{AsRawFd, FromRawFd};
        let fd = stream.as_raw_fd();

        let mut cmd = Command::new("/usr/sbin/sshd");
        cmd.args(&sshd_args);

        cmd.stdin(unsafe { Stdio::from_raw_fd(fd) });
        cmd.stdout(unsafe { Stdio::from_raw_fd(fd) });
        cmd.stderr(Stdio::null());

        let Ok(mut child) =
            (cmd.spawn()).inspect_err(|e| warn!("ssh: failed to start server: {e}"))
        else {
            continue;
        };

        let pid = child.id();
        info!("ssh: new connection from {remote}, sshd PID {pid}");

        tokio::spawn(async move {
            let _ = child.wait();
        });
    }
}
