//! Tier 2 (spec §4.3, docs/m0.md §2): one Firecracker microVM per job,
//! no network device, the rootfs read-only, the workspace read-write,
//! `flow-init` as PID 1. The host reads the framed console, enforces
//! the time limit, and writes the receipt's environment.

use crate::home::Home;
use anyhow::{bail, Context, Result};
use flow_core::api::{Record, SecretInput};
use flow_core::{Environment, Job, MountSource, Sandbox};
use std::io::{BufRead, BufReader, Read};
use std::path::Path;
use std::sync::mpsc;
use std::time::{Duration, Instant};

pub struct Outcome {
    /// None when the job never reported an exit (killed at the time limit, or the VM failed).
    pub exit_code: Option<i32>,
    pub duration_ms: u64,
    pub environment: Environment,
    pub timed_out: bool,
    /// Stopped because the run was cancelled (oxygenflow#32).
    pub cancelled: bool,
    /// The most memory the sandbox held, in MiB (oxygenflow#41).
    pub peak_memory_mib: Option<u64>,
}

/// Where a job's records go, and whether it is still wanted: asked a few
/// times a second while the job runs, and the microVM is stopped on `true`.
pub trait Sink {
    fn record(&mut self, r: Record);
    fn cancelled(&mut self) -> bool {
        false
    }
}

impl<F: FnMut(Record)> Sink for F {
    fn record(&mut self, r: Record) {
        self(r)
    }
}

/// The receipt environment for a job in this tier.
pub fn environment(home: &Home, sandbox: &Sandbox, job: &Job) -> Result<Environment> {
    let cfg = home.config.microvm.as_ref().context("flow.toml has no [microvm] section")?;
    Ok(Environment {
        runtime: "microvm".into(),
        parts: vec![
            ("firecracker".into(), cfg.firecracker.blake3.clone()),
            ("kernel".into(), cfg.kernel.blake3.clone()),
            ("init".into(), cfg.init.blake3.clone()),
            ("memory_mib".into(), sandbox.memory_mib.to_string()),
            ("cpus".into(), sandbox.cpus.to_string()),
        ]
        .into_iter()
        .chain(job.mounts.iter().map(|m| (format!("mount {}", m.target), format!("{} {} MiB", m.source.as_str(), m.size_mib))))
        .collect(),
        rootfs: sandbox.rootfs.clone(),
        argv: job.exec.clone(),
        env: job.env.iter().map(|(k, v)| (k.clone(), v.clone())).collect(),
    })
}

/// Boot, run, collect. `scratch` is an empty directory this call may fill.
#[allow(clippy::too_many_arguments)]
pub fn run(
    home: &Home,
    sandbox: &Sandbox,
    job: &Job,
    secrets: &[SecretInput],
    workspace: &Path,
    scratch: &Path,
    timeout: Duration,
    sink: &mut dyn Sink,
) -> Result<Outcome> {
    let cfg = home.config.microvm.as_ref().context("flow.toml has no [microvm] section")?;
    let firecracker = home.verified(&home.resolve(&cfg.firecracker.path), &cfg.firecracker.blake3)?;
    let kernel = home.verified(&home.resolve(&cfg.kernel.path), &cfg.kernel.blake3)?;
    let init = home.verified(&home.resolve(&cfg.init.path), &cfg.init.blake3)?;
    let rootfs = home.rootfs(&sandbox.rootfs);
    if !rootfs.is_file() {
        bail!("rootfs {} is not on this runner (import it with `flow rootfs import`)", sandbox.rootfs);
    }
    let rootfs = home.verified(&rootfs, &sandbox.rootfs)?;
    let environment = environment(home, sandbox, job)?;

    let mut spec = guest_spec(job);
    spec["loopback"] = serde_json::Value::Bool(sandbox.network == flow_core::Network::Loopback);
    // Allow-listed egress (oxygenflow#55): a vsock device and the runner's
    // proxy behind it; the guest relays 127.0.0.1:3128 to it.
    let egress = sandbox.network == flow_core::Network::Egress;
    let (egress_tx, egress_rx) = mpsc::channel::<String>();
    let mut _proxy = None;
    if egress {
        if !home.config.runner.as_ref().is_some_and(|r| r.egress) {
            bail!("this runner does not allow network egress (a runner allows it with `egress = true` under [runner] in flow.toml)");
        }
        _proxy = Some(
            crate::egress::start(&scratch.join(format!("v.sock_{}", crate::egress::VSOCK_PORT)), sandbox.egress_allow.clone(), egress_tx)
                .context("start the egress proxy")?,
        );
        spec["egress"] = serde_json::Value::Bool(true);
        add_proxy_env(&mut spec);
        sink.record(Record {
            t: 0,
            stream: "flow".into(),
            text: format!("egress: through this runner, to {} only", sandbox.egress_allow.join(", ")),
        });
    }
    // Secrets ride in the guest's job file, which only flow-init reads,
    // outside the root the job is confined to; it writes each as a file.
    if !secrets.is_empty() {
        spec["secrets"] =
            serde_json::Value::Array(secrets.iter().map(|s| serde_json::json!({ "path": s.path, "value": s.value })).collect());
    }
    // Each `ephemeral://disk` mount is a fresh sparse ext4 disk in the
    // scratch directory, attached after the workspace (vdc, vdd, ...).
    let mut disks = Vec::new();
    for (i, m) in job.mounts.iter().filter(|m| m.source == MountSource::Disk).enumerate() {
        let path = scratch.join(format!("mount-{i}.ext4"));
        // A disk already in place came from the build cache (#28).
        if !path.exists() {
            crate::workspace::empty_ext4(&path, m.size_mib, "scratch")?;
        }
        disks.push(
            serde_json::json!({ "drive_id": format!("mount{i}"), "path_on_host": path, "is_root_device": false, "is_read_only": false }),
        );
    }
    let init_bytes = std::fs::read(&init)?;
    let spec_bytes = serde_json::to_vec(&spec)?;
    let initrd = crate::cpio::newc(&[
        crate::cpio::Entry { name: "init", mode: 0o755, data: &init_bytes },
        crate::cpio::Entry { name: "job.json", mode: 0o644, data: &spec_bytes },
    ]);
    let initrd_path = scratch.join("initrd.cpio");
    std::fs::write(&initrd_path, initrd)?;
    let drives: Vec<serde_json::Value> = [
        serde_json::json!({ "drive_id": "rootfs", "path_on_host": rootfs, "is_root_device": false, "is_read_only": true }),
        serde_json::json!({ "drive_id": "workspace", "path_on_host": workspace, "is_root_device": false, "is_read_only": false }),
    ]
    .into_iter()
    .chain(disks)
    .collect();
    let vm = serde_json::json!({
        "boot-source": {
            "kernel_image_path": kernel,
            "initrd_path": initrd_path,
            "boot_args": "console=ttyS0 reboot=k panic=-1 pci=off quiet loglevel=0 rdinit=/init",
        },
        "drives": drives,
        "machine-config": { "vcpu_count": sandbox.cpus, "mem_size_mib": sandbox.memory_mib },
    });
    let mut vm = vm;
    if egress {
        // Relative to the VMM's working directory (the scratch directory),
        // which keeps the socket path short.
        vm["vsock"] = serde_json::json!({ "guest_cid": 3, "uds_path": "v.sock" });
    }
    let vm_path = scratch.join("vm.json");
    std::fs::write(&vm_path, serde_json::to_vec_pretty(&vm)?)?;

    let started = Instant::now();
    let mut child = std::process::Command::new(&firecracker)
        .arg("--no-api")
        .arg("--config-file")
        .arg(&vm_path)
        .args(["--level", "Error"])
        .current_dir(scratch)
        .stdin(std::process::Stdio::null())
        .stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped())
        .spawn()
        .with_context(|| format!("start {}", firecracker.display()))?;
    let (tx, rx) = mpsc::channel::<(bool, Vec<u8>)>();
    let out = child.stdout.take().unwrap();
    let err = child.stderr.take().unwrap();
    let tx2 = tx.clone();
    let readers = [std::thread::spawn(move || lines(out, true, tx)), std::thread::spawn(move || lines(err, false, tx2))];
    let mut exit_code = None;
    let mut timed_out = false;
    let mut cancelled = false;
    let deadline = started + timeout;
    let t = |s: &Instant| s.elapsed().as_millis() as u64;
    // The VMM's peak resident memory is the guest memory it touched: read
    // while it runs, since it is gone once it exits.
    let mut peak_memory_mib = None;
    loop {
        if let Some(m) = crate::facts::peak_rss_mib(child.id()) {
            peak_memory_mib = Some(m);
        }
        while let Ok(line) = egress_rx.try_recv() {
            sink.record(Record { t: t(&started), stream: "flow".into(), text: line });
        }
        match rx.recv_timeout(Duration::from_millis(200)) {
            Ok((from_console, line)) => {
                let text = String::from_utf8_lossy(&line).trim_end_matches('\r').to_string();
                let rec = |stream: &str, text: &str| Record { t: t(&started), stream: stream.into(), text: text.into() };
                if !from_console {
                    sink.record(rec("console", &format!("firecracker: {text}")));
                } else if let Some(rest) = text.strip_prefix('\u{1e}') {
                    let (kind, body) = rest.split_at(rest.len().min(1));
                    let body = body.strip_prefix(' ').unwrap_or(body);
                    match kind {
                        "O" => sink.record(rec("out", body)),
                        "E" => sink.record(rec("err", body)),
                        "X" => exit_code = body.trim().parse().ok(),
                        "F" => sink.record(rec("flow", &format!("guest init failed: {body}"))),
                        _ => sink.record(rec("console", &text)),
                    }
                } else if !text.is_empty() {
                    sink.record(rec("console", &text));
                }
            }
            Err(mpsc::RecvTimeoutError::Timeout) => {}
            Err(mpsc::RecvTimeoutError::Disconnected) => break,
        }
        if !timed_out && !cancelled && sink.cancelled() {
            cancelled = true;
            let _ = child.kill();
            sink.record(Record { t: t(&started), stream: "flow".into(), text: "stopped: the run was cancelled".into() });
        }
        if Instant::now() >= deadline && !timed_out && !cancelled {
            timed_out = true;
            let _ = child.kill();
            sink.record(Record {
                t: t(&started),
                stream: "flow".into(),
                text: format!("stopped at the time limit of {} s", timeout.as_secs()),
            });
        }
    }
    let status = child.wait()?;
    for r in readers {
        let _ = r.join();
    }
    if exit_code.is_none() && !timed_out && !cancelled {
        sink.record(Record {
            t: t(&started),
            stream: "flow".into(),
            text: format!("the microVM ended without the job's exit code ({status})"),
        });
    }
    Ok(Outcome {
        exit_code: if timed_out || cancelled { None } else { exit_code },
        duration_ms: t(&started),
        environment,
        timed_out,
        cancelled,
        peak_memory_mib,
    })
}

/// The proxy variables a job with egress gets, in both spellings tools
/// read, unless the job sets them itself.
fn add_proxy_env(spec: &mut serde_json::Value) {
    let Some(env) = spec["env"].as_array_mut() else { return };
    let set: Vec<String> = env.iter().filter_map(|kv| kv[0].as_str().map(String::from)).collect();
    for (k, v) in [
        ("HTTPS_PROXY", crate::egress::GUEST_PROXY),
        ("https_proxy", crate::egress::GUEST_PROXY),
        ("HTTP_PROXY", crate::egress::GUEST_PROXY),
        ("http_proxy", crate::egress::GUEST_PROXY),
        ("NO_PROXY", "localhost,127.0.0.1"),
        ("no_proxy", "localhost,127.0.0.1"),
    ] {
        if !set.iter().any(|s| s == k) {
            env.push(serde_json::json!([k, v]));
        }
    }
}

/// What `flow-init` reads from `/job.json`: argv, env, and the mounts
/// with the device each disk mount is attached as (vdc onwards).
pub fn guest_spec(job: &Job) -> serde_json::Value {
    let mut next_disk = b'c';
    let mounts: Vec<serde_json::Value> = job
        .mounts
        .iter()
        .map(|m| match m.source {
            MountSource::Tmpfs => serde_json::json!({ "kind": "tmpfs", "target": m.target, "size_mib": m.size_mib }),
            MountSource::Disk => {
                let dev = format!("/dev/vd{}", next_disk as char);
                next_disk += 1;
                serde_json::json!({ "kind": "disk", "target": m.target, "size_mib": m.size_mib, "device": dev })
            }
        })
        .collect();
    serde_json::json!({ "argv": job.exec, "env": job.env.iter().collect::<Vec<_>>(), "mounts": mounts })
}

fn lines(r: impl Read, console: bool, tx: mpsc::Sender<(bool, Vec<u8>)>) {
    let mut r = BufReader::new(r);
    let mut buf = Vec::new();
    loop {
        buf.clear();
        match r.read_until(b'\n', &mut buf) {
            Ok(0) | Err(_) => break,
            Ok(_) => {
                if buf.last() == Some(&b'\n') {
                    buf.pop();
                }
                if tx.send((console, buf.clone())).is_err() {
                    break;
                }
            }
        }
    }
}

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

    #[test]
    fn a_job_with_egress_gets_the_proxy_variables_unless_it_sets_them() {
        let mut spec = serde_json::json!({ "env": [["HTTPS_PROXY", "http://mine:1"]] });
        super::add_proxy_env(&mut spec);
        let env = spec["env"].as_array().unwrap();
        let get = |k: &str| env.iter().find(|kv| kv[0] == k).map(|kv| kv[1].as_str().unwrap().to_string());
        assert_eq!(get("HTTPS_PROXY").as_deref(), Some("http://mine:1"), "the job's own is kept");
        assert_eq!(get("https_proxy").as_deref(), Some("http://127.0.0.1:3128"));
        assert_eq!(get("NO_PROXY").as_deref(), Some("localhost,127.0.0.1"));
    }

    #[test]
    fn guest_spec_numbers_disks_after_the_workspace() {
        let mut job: Job = serde_json::from_value(serde_json::json!({
            "name": "t", "sandbox": "s", "scope": [], "depends_on": [], "env": {"A": "1"}, "exec": ["true"]
        }))
        .unwrap();
        assert_eq!(guest_spec(&job)["mounts"], serde_json::json!([]), "no mounts");
        job.mounts = vec![
            Mount { source: MountSource::Disk, target: "/src/target".into(), size_mib: 8192, cache: None, key: vec![] },
            Mount { source: MountSource::Tmpfs, target: "/tmp/cache".into(), size_mib: 256, cache: None, key: vec![] },
            Mount { source: MountSource::Disk, target: "/tmp/z".into(), size_mib: 64, cache: None, key: vec![] },
        ];
        let spec = guest_spec(&job);
        assert_eq!(spec["mounts"][0], serde_json::json!({"kind": "disk", "target": "/src/target", "size_mib": 8192, "device": "/dev/vdc"}));
        assert_eq!(spec["mounts"][1], serde_json::json!({"kind": "tmpfs", "target": "/tmp/cache", "size_mib": 256}));
        assert_eq!(spec["mounts"][2]["device"], "/dev/vdd");
        assert_eq!(spec["env"], serde_json::json!([["A", "1"]]));
    }
}