use crate::command::{self, OutKind, Spec};
use crate::protocol::{self, Message, v1};
use nix::unistd::{Group, User};
use std::ffi::OsString;
use std::path::PathBuf;
use std::time::Duration;
use tokio::io::AsyncWriteExt;
use tokio::sync::mpsc::Sender;
use tokio_vsock::{VsockAddr, VsockStream};
use tracing::{info, warn};

const DEFAULT_USER: &str = "spindle-workflow";
const VSOCK_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);

pub async fn run(id: String, req: v1::ExecStart, out: Sender<Message>, host_cid: u32) {
    let send_exit = async |exit_code: i32, error: Option<String>, timed_out: bool| {
        let msg = Message {
            id: id.clone(),
            exec_exit: Some(v1::ExecExit {
                exit_code,
                error: protocol::error_or_empty(error),
                timed_out,
            }),
            ..Default::default()
        };
        let _ = out.send(msg).await;
    };

    if req.argv.is_empty() {
        send_exit(127, Some("missing argv".to_owned()), false).await;
        return;
    }

    let user = if req.user.is_empty() {
        DEFAULT_USER
    } else {
        req.user.as_str()
    };
    let run_as = match resolve_user(user) {
        Ok(run_as) => run_as,
        Err(err) => {
            send_exit(127, Some(err), false).await;
            return;
        }
    };

    let mut env = run_as.login_env();
    let runtime_dir = run_as.runtime_dir();
    match runtime_dir.try_exists() {
        Ok(true) => env.push((
            OsString::from("XDG_RUNTIME_DIR"),
            runtime_dir.into_os_string(),
        )),
        Ok(false) => {}
        Err(err) => warn!(error = %err, "could not stat XDG_RUNTIME_DIR for workflow user"),
    }

    let mut spec = Spec::new(req.argv[0].clone())
        .args(req.argv[1..].iter().cloned())
        .envs(env)
        .envs(parse_env(&req.env))
        .run_as(run_as.uid, run_as.gid);

    if req.stdio_vsock_port == 0 {
        send_exit(
            1,
            Some("exec is missing a stdio vsock port".to_owned()),
            false,
        )
        .await;
        return;
    }
    let addr = VsockAddr::new(host_cid, req.stdio_vsock_port);
    let conn = match tokio::time::timeout(VSOCK_CONNECT_TIMEOUT, VsockStream::connect(addr)).await {
        Ok(Ok(conn)) => conn,
        Ok(Err(error)) => {
            send_exit(1, Some(format!("dial host stdio port: {error}")), false).await;
            return;
        }
        Err(_) => {
            send_exit(1, Some("dial host stdio port: timed out".to_owned()), false).await;
            return;
        }
    };
    spec = spec.piped_stdin();
    if !req.cwd.is_empty() {
        spec = spec.cwd(req.cwd.clone());
    }
    let timeout =
        (req.timeout_seconds > 0).then(|| Duration::from_secs(u64::from(req.timeout_seconds)));
    if let Some(timeout) = timeout {
        spec = spec.timeout(timeout);
    }

    info!(
        %id,
        user = %run_as.name,
        uid = run_as.uid,
        gid = run_as.gid,
        argv = ?req.argv,
        cwd = ?req.cwd,
        "starting exec"
    );

    let mut cmd = match command::spawn_streaming(spec) {
        Ok(cmd) => cmd,
        Err(err) => {
            send_exit(127, Some(err.to_string()), false).await;
            return;
        }
    };

    let (mut conn_reader, mut conn_writer) = tokio::io::split(conn);
    // dropping this also closes the socket if the exec gets cancelled
    let mut child_stdin = cmd.stdin.take().expect("piped_stdin");
    let _stdin_pump = AbortOnDrop(tokio::spawn(async move {
        let _ = tokio::io::copy(&mut conn_reader, &mut child_stdin).await;
    }));

    let (mut events, exit_task) = cmd.into_parts();
    let mut stdio_error = None;
    while let Some(event) = events.recv().await {
        match event.kind {
            OutKind::Stdout => {
                if let Err(error) = conn_writer.write_all(&event.data).await {
                    stdio_error = Some(format!("forward stdout to host: {error}"));
                    // stop these before their pipes fill and block the child
                    events.close();
                    break;
                }
            }
            OutKind::Stderr => {
                let _ = out
                    .send(Message {
                        id: id.clone(),
                        exec_stderr: Some(v1::ExecStderr {
                            data: event.data.into(),
                        }),
                        ..Default::default()
                    })
                    .await;
            }
        }
    }
    // the child might keep reading stdin after it closes stdout
    let exit = match exit_task
        .await
        .unwrap_or_else(|error| Err(anyhow::anyhow!("command supervisor failed: {error}")))
    {
        Ok(exit) => exit,
        Err(err) => {
            send_exit(127, Some(err.to_string()), false).await;
            return;
        }
    };

    // losing stdout fails the exec even if the child exited cleanly
    if let Some(error) = stdio_error {
        send_exit(1, Some(error), false).await;
        return;
    }
    send_exit(exit.exit_code, exit.error, exit.timed_out).await
}

// aborts the task when its exec goes away
struct AbortOnDrop(tokio::task::JoinHandle<()>);

impl Drop for AbortOnDrop {
    fn drop(&mut self) {
        self.0.abort();
    }
}

#[derive(Clone, Debug)]
struct ResolvedUser {
    name: String,
    uid: u32,
    gid: u32,
    home: OsString,
    shell: OsString,
}

impl ResolvedUser {
    fn login_env(&self) -> Vec<(OsString, OsString)> {
        let xdg_cache_home = PathBuf::from(&self.home).join(".cache");
        vec![
            (OsString::from("USER"), OsString::from(&self.name)),
            (OsString::from("LOGNAME"), OsString::from(&self.name)),
            (OsString::from("HOME"), self.home.clone()),
            (
                OsString::from("XDG_CACHE_HOME"),
                xdg_cache_home.into_os_string(),
            ),
            (OsString::from("SHELL"), self.shell.clone()),
        ]
    }

    fn runtime_dir(&self) -> PathBuf {
        PathBuf::from(format!("/run/user/{}", self.uid))
    }
}

fn resolve_user(spec: &str) -> Result<ResolvedUser, String> {
    let spec = spec.trim();
    if spec.is_empty() {
        return resolve_user(DEFAULT_USER);
    }

    let (user_part, group_part) = spec
        .split_once(':')
        .map(|(user, group)| (user, Some(group)))
        .unwrap_or((spec, None));

    let mut user = lookup_user(user_part)?;
    if let Some(group) = group_part.filter(|group| !group.is_empty()) {
        user.gid = lookup_group(group)?;
    }

    if user.uid == 0 || user.gid == 0 {
        return Err(format!("refusing to run exec as privileged user {spec:?}"));
    }

    Ok(user)
}

fn lookup_user(name: &str) -> Result<ResolvedUser, String> {
    match User::from_name(name) {
        Ok(Some(user)) => Ok(ResolvedUser {
            name: name.to_owned(),
            uid: user.uid.as_raw(),
            gid: user.gid.as_raw(),
            home: user.dir.into_os_string(),
            shell: user.shell.into_os_string(),
        }),
        Ok(None) => {
            let uid = name
                .parse::<u32>()
                .map_err(|_| format!("workflow user {name:?} was not found"))?;
            Ok(ResolvedUser {
                name: name.to_owned(),
                uid,
                gid: uid,
                home: OsString::from("/"),
                shell: OsString::from("/bin/sh"),
            })
        }
        Err(error) => Err(format!("lookup workflow user {name:?}: {error}")),
    }
}

fn lookup_group(name: &str) -> Result<u32, String> {
    match Group::from_name(name) {
        Ok(Some(group)) => Ok(group.gid.as_raw()),
        Ok(None) => name
            .parse::<u32>()
            .map_err(|_| format!("workflow group {name:?} was not found")),
        Err(error) => Err(format!("lookup workflow group {name:?}: {error}")),
    }
}

fn parse_env(values: &[String]) -> Vec<(OsString, OsString)> {
    values
        .iter()
        .filter_map(|value| value.split_once('='))
        .map(|(key, value)| (OsString::from(key), OsString::from(value)))
        .collect()
}

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

    #[test]
    fn refuses_root_exec_user() {
        let err = resolve_user("root").unwrap_err();
        assert!(err.contains("refusing to run exec as privileged user"));
    }

    #[test]
    fn refuses_root_exec_group() {
        let err = resolve_user("65534:0").unwrap_err();
        assert!(err.contains("refusing to run exec as privileged user"));
    }
}
