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

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,

}

/// 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.

pub fn run(

    home: &Home,

    sandbox: &Sandbox,

    job: &Job,

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

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

        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 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;

    loop {

        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,

    })

}

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

            Mount { source: MountSource::Tmpfs, target: "/tmp/cache".into(), size_mib: 256 },

            Mount { source: MountSource::Disk, target: "/tmp/z".into(), size_mib: 64 },

        ];

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

    }

}