//! Native Cargo execution, isolated workspace copies, streamed output, and bounded disk use.
mod files;
#[cfg(test)]
mod tests;
use anyhow::{Context, Result, ensure};
use sausage_vfs::Layer;
use serde::{Deserialize, Serialize};
use std::{
    path::{Path, PathBuf},
    process::Stdio,
    sync::Arc,
    time::Duration,
};
use tokio::{
    io::AsyncReadExt,
    process::Command,
    sync::{Mutex, mpsc},
    time::Instant,
};

#[derive(Deserialize)]
pub struct Request {
    pub argv: Vec<String>,
    pub cwd: String,
    pub timeout_secs: Option<u64>,
}
#[derive(Serialize, PartialEq, Eq, PartialOrd, Ord)]
pub struct Executable {
    pub path: String,
    pub cwd: String,
}
#[derive(Serialize)]
#[serde(tag = "event", rename_all = "snake_case")]
pub enum Event {
    Stdout {
        bytes: Vec<u8>,
    },
    Stderr {
        bytes: Vec<u8>,
    },
    Exit {
        code: i32,
        executables: Vec<Executable>,
    },
}
pub struct Builder {
    root: Arc<dyn Layer>,
    directory: PathBuf,
    lock: Mutex<()>,
}
impl Builder {
    pub fn new(root: Arc<dyn Layer>, host: &Path) -> Result<Self> {
        Ok(Self {
            root,
            directory: std::fs::canonicalize(host)?.join("builds"),
            lock: Mutex::new(()),
        })
    }
    pub async fn execute(&self, request: Request, tx: &mpsc::Sender<Event>) -> Result<()> {
        let _guard = self.lock.lock().await;
        let action = request
            .argv
            .first()
            .context("missing Cargo command")?
            .as_str();
        ensure!(
            [
                "build", "check", "clippy", "test", "run", "fmt", "fetch", "add", "remove",
                "update", "tree", "metadata", "clean"
            ]
            .contains(&action),
            "unsupported Cargo command"
        );
        let compiles = ["build", "check", "clippy", "test", "run"].contains(&action);
        let mut args = request.argv.iter().take_while(|arg| arg.as_str() != "--");
        while let Some(arg) = args.next() {
            ensure!(
                arg != "--target-dir" && !arg.starts_with("--target-dir="),
                "the builder controls the target directory"
            );
            let target = if arg == "--target" {
                args.next().map(String::as_str)
            } else {
                arg.strip_prefix("--target=")
            };
            if compiles && let Some(target) = target {
                ensure!(
                    target == "wasm32-wasip2",
                    "the build target is always wasm32-wasip2"
                );
            }
        }
        let seconds = request.timeout_secs.unwrap_or(1800);
        ensure!(seconds > 0, "timeout must be positive");
        let deadline = Instant::now()
            .checked_add(Duration::from_secs(seconds))
            .context("timeout too large")?;
        let world_cwd = sausage_vfs::resolve("/", &request.cwd);
        let workspace = files::workspace(self.root.as_ref(), &world_cwd).await?;
        let world_target = sausage_vfs::resolve(&workspace, "target");
        use std::hash::{Hash, Hasher};
        let mut hash = std::collections::hash_map::DefaultHasher::new();
        workspace.hash(&mut hash);
        let id = format!("{:016x}", hash.finish());
        let base = self.directory.join(id);
        let source = base.join("source");
        // Cargo clean must remove its directory, not a namespace mount point.
        let target_mount = base.join("target");
        let target = target_mount.join("cargo");
        // Keep /tmp's namespace mountpoints outside the directory Cargo clean removes.
        let temporary = target_mount.join("tmp");
        let registry = base.join("cargo");
        let cache = base.join("kache");
        if source.exists() {
            std::fs::remove_dir_all(&source)?;
        }
        let before = files::copy_out(self.root.as_ref(), &workspace, &source).await?;
        for path in [&target, &registry, &cache, &temporary] {
            std::fs::create_dir_all(path)?;
        }
        std::fs::write(
            target.join("CACHEDIR.TAG"),
            b"Signature: 8a477f597d28d172789f06886806bc55\n",
        )?;
        std::fs::write(base.join("last-used"), [])?;
        let relative = world_cwd
            .strip_prefix(&workspace)
            .expect("workspace is an ancestor of cwd")
            .trim_start_matches('/');
        let cwd = if relative.is_empty() {
            source.clone()
        } else {
            source.join(relative)
        };
        let mut fetch = Command::new("cargo");
        if action == "add" {
            fetch.args(&request.argv).arg("--dry-run");
        } else {
            fetch.arg("fetch");
        }
        fetch.current_dir(&cwd).env("CARGO_HOME", &registry);
        let mut code = if matches!(action, "clean" | "fmt" | "remove") {
            0
        } else {
            run(fetch, tx, deadline).await?.0
        };
        let mut executables = Vec::new();
        if code == 0 {
            let tools = base.join("tools");
            std::fs::create_dir_all(&tools)?;
            for (name, bytes) in [
                (
                    "init.rs",
                    include_bytes!("../../../userland/cwd-init/init.rs").as_slice(),
                ),
                (
                    "exit.wit",
                    include_bytes!("../../../userland/cwd-init/exit.wit").as_slice(),
                ),
                (
                    "build.nu",
                    include_bytes!("../../../userland/cwd-init/build.nu").as_slice(),
                ),
                (
                    "link.nu",
                    include_bytes!("../../../userland/cwd-init/link.nu").as_slice(),
                ),
                ("run.nu", include_bytes!("run.nu").as_slice()),
            ] {
                std::fs::write(tools.join(name), bytes)?;
            }
            let linker = std::fs::read_to_string(tools.join("link.nu"))?.replacen(
                "#!/usr/bin/env nu",
                &format!("#!{}", executable("nu")?.display()),
                1,
            );
            std::fs::write(tools.join("link.nu"), linker)?;
            use std::os::unix::fs::PermissionsExt;
            std::fs::set_permissions(
                tools.join("link.nu"),
                std::fs::Permissions::from_mode(0o755),
            )?;
            let mut startup = Command::new("nu");
            startup.arg(tools.join("build.nu"));
            let (startup_code, _) = run(startup, tx, deadline).await?;
            if startup_code != 0 {
                files::copy_back(self.root.as_ref(), &source, &workspace, Some(&before)).await?;
                self.evict(20 * 1024 * 1024 * 1024)?;
                tx.send(Event::Exit {
                    code: startup_code,
                    executables,
                })
                .await?;
                return Ok(());
            }
            let kache = executable("kache")?;
            let run_path = target_mount.join("run-path");
            if action == "run" && run_path.exists() {
                std::fs::remove_file(&run_path)?;
            }
            let mut cmd = Command::new("bwrap");
            cmd.args([
                "--unshare-user",
                "--unshare-pid",
                "--unshare-net",
                "--die-with-parent",
                "--ro-bind",
                "/nix/store",
                "/nix/store",
            ])
            .arg("--bind")
            .arg(&temporary)
            .arg("/tmp")
            .arg("--ro-bind")
            .arg(executable("sh")?)
            .arg("/bin/sh");
            for path in [&source, &target_mount, &cache] {
                cmd.arg("--bind").arg(path).arg(path);
            }
            for path in [&registry, &tools] {
                cmd.arg("--ro-bind").arg(path).arg(path);
            }
            cmd.args([
                "--proc",
                "/proc",
                "--dev",
                "/dev",
                "--remount-ro",
                "/",
                "--chdir",
            ])
            .arg(&cwd)
            .args(["--clearenv", "--setenv", "PATH"])
            .arg(std::env::var("PATH")?)
            .args(["--setenv", "HOME", "/nonexistent", "--setenv", "CARGO_HOME"])
            .arg(&registry)
            .args(["--setenv", "CARGO_TARGET_DIR"])
            .arg(&target)
            .args(["--setenv", "KACHE_CACHE_DIR"])
            .arg(&cache)
            .args(["--setenv", "SAUSAGE_CWD_INIT"])
            .arg(tools.join("init.o"))
            .args(["--setenv", "SAUSAGE_RUN_PATH"])
            .arg(&run_path)
            .args(["--setenv", "CARGO_BUILD_JOBS", "2"]);
            for name in [
                "CC_wasm32_wasip2",
                "AR_wasm32_wasip2",
                "CFLAGS_wasm32_wasip2",
            ] {
                if let Ok(value) = std::env::var(name) {
                    cmd.args(["--setenv", name]).arg(value);
                }
            }
            cmd.args(["cargo", "--offline", "--config"])
                .arg(format!("build.rustc-wrapper={:?}", kache))
                .arg("--config")
                .arg(format!(
                    "target.wasm32-wasip2.linker={:?}",
                    tools.join("link.nu")
                ))
                .args([
                    "--config",
                    "target.wasm32-wasip2.rustflags=['-C','link-arg=sausage-cwd-init']",
                ]);
            let mut args = request.argv.clone();
            if matches!(action, "run" | "test")
                && let Some(index) = args.iter().position(|a| a == "--")
            {
                args.truncate(index);
            }
            if action == "test" {
                // Cargo's positional test filter belongs to the world runner, not the compiler.
                args = test_build_args(&args);
                if !args.iter().any(|a| a == "--no-run") {
                    args.push("--no-run".into());
                }
            }
            let split = args.iter().position(|a| a == "--").unwrap_or(args.len());
            if action == "run" {
                cmd.arg("--config").arg(format!(
                    "target.wasm32-wasip2.runner=[{:?},{:?}]",
                    executable("nu")?,
                    tools.join("run.nu")
                ));
            }
            // Keep all global config flags before the subcommand so Cargo retains them.
            cmd.args(&args[..split]);
            if compiles {
                cmd.args([
                    "--target",
                    "wasm32-wasip2",
                    "--message-format=json-render-diagnostics",
                ]);
            }
            cmd.args(&args[split..]);
            let (status, stdout) = run(cmd, tx, deadline).await?;
            code = status;
            if action == "clean" && code == 0 {
                files::clean_artifacts(self.root.as_ref(), &target, &world_target).await?;
            }
            let artifacts = target.join("wasm32-wasip2");
            if artifacts.exists() {
                files::copy_back(
                    self.root.as_ref(),
                    &artifacts,
                    &format!("{world_target}/wasm32-wasip2"),
                    None,
                )
                .await?;
            }
            if code == 0 && action == "run" && run_path.exists() {
                let path = std::fs::read_to_string(&run_path)?;
                let relative = Path::new(&path).strip_prefix(&target)?;
                executables.push(Executable {
                    path: format!("{world_target}/{}", relative.display()),
                    cwd: world_cwd.clone(),
                });
            }
            if code == 0
                && action == "test"
                && !request
                    .argv
                    .iter()
                    .take_while(|a| a.as_str() != "--")
                    .any(|a| a == "--no-run")
            {
                for line in stdout.split(|b| *b == b'\n') {
                    if let Ok(value) = serde_json::from_slice::<serde_json::Value>(line)
                        && value["reason"] == "compiler-artifact"
                        && value["profile"]["test"] == true
                        && let Some(path) = value["executable"].as_str()
                        && let Ok(relative) = Path::new(path).strip_prefix(&target)
                    {
                        let manifest = Path::new(
                            value["manifest_path"]
                                .as_str()
                                .context("missing test manifest")?,
                        );
                        let package = manifest.parent().unwrap().strip_prefix(&source)?;
                        let cwd = sausage_vfs::resolve(&workspace, &package.to_string_lossy());
                        executables.push(Executable {
                            path: format!("{world_target}/{}", relative.display()),
                            cwd,
                        });
                    }
                }
                executables.sort();
                executables.dedup();
            }
        }
        files::copy_back(self.root.as_ref(), &source, &workspace, Some(&before)).await?;
        self.evict(20 * 1024 * 1024 * 1024)?;
        tx.send(Event::Exit { code, executables }).await?;
        Ok(())
    }
    fn evict(&self, limit: u64) -> Result<()> {
        let mut entries = Vec::new();
        let mut total = 0;
        for entry in std::fs::read_dir(&self.directory)? {
            let path = entry?.path();
            let size = files::size(&path)?;
            total += size;
            entries.push((
                std::fs::metadata(path.join("last-used"))?.modified()?,
                path,
                size,
            ));
        }
        entries.sort_by_key(|e| e.0);
        for (_, path, size) in entries {
            if total <= limit {
                break;
            }
            std::fs::remove_dir_all(path)?;
            total -= size;
        }
        Ok(())
    }
}
fn executable(name: &str) -> Result<PathBuf> {
    let path = std::env::split_paths(&std::env::var_os("PATH").context("PATH missing")?)
        .map(|p| p.join(name))
        .find(|p| p.is_file())
        .context(format!("missing {name}"))?;
    Ok(std::fs::canonicalize(path)?)
}
fn test_build_args(args: &[String]) -> Vec<String> {
    let mut out = vec!["test".into()];
    let mut value = false;
    for arg in &args[1..] {
        if value || arg.starts_with('-') {
            out.push(arg.clone());
        }
        value = [
            "-p",
            "--package",
            "--test",
            "--bin",
            "--example",
            "--features",
            "--profile",
            "--manifest-path",
            "--exclude",
            "-j",
            "--jobs",
            "-F",
            "--config",
            "--target",
            "--target-dir",
            "--message-format",
            "--color",
            "--lockfile-path",
            "-Z",
        ]
        .contains(&arg.as_str());
    }
    out
}
// The guard kills native descendants when Cargo exits or the caller cancels the command.
struct ProcessGroup(u32);
impl Drop for ProcessGroup {
    fn drop(&mut self) {
        unsafe {
            libc::kill(-(self.0 as i32), libc::SIGKILL);
        }
    }
}
async fn pump(
    mut reader: impl tokio::io::AsyncRead + Unpin,
    tx: &mpsc::Sender<Event>,
    stderr: bool,
) -> Result<Vec<u8>> {
    let mut all = Vec::new();
    let mut bytes = vec![0; 8192];
    loop {
        let n = reader.read(&mut bytes).await?;
        if n == 0 {
            return Ok(all);
        }
        if !stderr {
            all.extend_from_slice(&bytes[..n]);
        }
        tx.send(if stderr {
            Event::Stderr {
                bytes: bytes[..n].to_vec(),
            }
        } else {
            Event::Stdout {
                bytes: bytes[..n].to_vec(),
            }
        })
        .await?;
    }
}
async fn run(
    mut cmd: Command,
    tx: &mpsc::Sender<Event>,
    deadline: Instant,
) -> Result<(i32, Vec<u8>)> {
    cmd.stdin(Stdio::null())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .process_group(0)
        .kill_on_drop(true);
    let mut child = cmd.spawn()?;
    let group = ProcessGroup(child.id().unwrap());
    let out = child.stdout.take().unwrap();
    let err = child.stderr.take().unwrap();
    let collect = async {
        let (status, out, _) = tokio::try_join!(
            async { Ok::<_, anyhow::Error>(child.wait().await?) },
            pump(out, tx, false),
            pump(err, tx, true)
        )?;
        Ok::<_, anyhow::Error>((status.code().unwrap_or(1), out))
    };
    let (result, timed_out) = tokio::select! {
        result = collect => (result, false),
        _ = tx.closed() => anyhow::bail!("caller exited"),
        _ = tokio::time::sleep_until(deadline) => (Ok((124, Vec::new())), true),
    };
    // Kill the command before a full output channel can delay its timeout message.
    drop(group);
    if timed_out {
        tx.send(Event::Stderr {
            bytes: b"cargo: command timed out\n".to_vec(),
        })
        .await?;
    }
    result
}