//! 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>,

    /// Each cache disk the guest pruned before it stopped, by its target

    /// (oxygenflow#72); the syscall tier's are pruned by the runner.

    pub pruned: Vec<(String, flow_init::prune::Pruned)>,

}

/// 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 percent = home.config.runner.as_ref().and_then(|r| r.cache_prune_percent).unwrap_or(flow_init::prune::DEFAULT_PERCENT);

    let mut spec = guest_spec(job, percent);

    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;

    let mut pruned = Vec::new();

    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}"))),

                        "C" => match pruned_frame(body) {

                            Some(p) => pruned.push(p),

                            None => sink.record(rec("console", &text)),

                        },

                        _ => 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,

        pruned,

    })

}

/// A `\x1eC` frame: a cache disk the guest pruned (oxygenflow#72).

fn pruned_frame(body: &str) -> Option<(String, flow_init::prune::Pruned)> {

    #[derive(serde::Deserialize)]

    struct Frame {

        target: String,

        pruned: flow_init::prune::Pruned,

    }

    let f: Frame = serde_json::from_str(body).ok()?;

    Some((f.target, f.pruned))

}

/// 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). A cache

/// disk is pruned to `percent` of its size after a passing job

/// (oxygenflow#72), leaving the job's outputs on it.

pub fn guest_spec(job: &Job, percent: u8) -> 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;

                let mut disk = serde_json::json!({ "kind": "disk", "target": m.target, "size_mib": m.size_mib, "device": dev });

                if m.cache.is_some() {

                    let prefix = format!("{}/", m.target);

                    let keep: Vec<&str> = job.outputs.iter().filter_map(|o| o.path.strip_prefix(&prefix)).collect();

                    disk["prune"] = percent.into();

                    disk["keep"] = keep.into();

                }

                disk

            }

        })

        .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, 60)["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, 60);

        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"]]));

    }

    /// oxygenflow#72: a cache disk is pruned in the guest, leaving the

    /// job's outputs on it; a plain disk is not, and the guest's report

    /// comes back as a frame.

    #[test]

    fn a_cache_disk_is_pruned_in_the_guest_but_its_outputs_stay() {

        let mut job: Job = serde_json::from_value(serde_json::json!({

            "name": "t", "sandbox": "s", "scope": [], "depends_on": [], "env": {}, "exec": ["true"],

            "outputs": [{"name": "bin", "path": "/src/target/release/tool"}, {"name": "w", "path": "/src/w"}]

        }))

        .unwrap();

        job.mounts = vec![

            Mount { source: MountSource::Disk, target: "/src/target".into(), size_mib: 8192, cache: Some("t".into()), key: vec![] },

            Mount { source: MountSource::Disk, target: "/tmp/z".into(), size_mib: 64, cache: None, key: vec![] },

        ];

        let spec = guest_spec(&job, 45);

        assert_eq!((spec["mounts"][0]["prune"].as_u64(), &spec["mounts"][0]["keep"]), (Some(45), &serde_json::json!(["release/tool"])));

        assert!(spec["mounts"][1].get("prune").is_none() && spec["mounts"][1].get("keep").is_none());

        let (target, p) =

            pruned_frame(r#"{"target":"/src/target","pruned":{"size":100,"percent":60,"before":90,"after":50,"freed":40,"files":3}}"#)

                .unwrap();

        assert_eq!((target.as_str(), p.after, p.cap(), p.fits(), p.unpruned), ("/src/target", 50, 60, true, None));

        assert!(pruned_frame("not json").is_none());

    }

}