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

}

/// 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()),

        ],

        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 FnMut(Record),

) -> 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 spec = serde_json::json!({ "argv": job.exec, "env": job.env.iter().collect::<Vec<_>>() });

    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 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": [

            { "drive_id": "rootfs", "path_on_host": rootfs, "is_root_device": false, "is_read_only": true },

            { "drive_id": "workspace", "path_on_host": workspace, "is_root_device": false, "is_read_only": false },

        ],

        "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 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(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(rec("out", body)),

                        "E" => sink(rec("err", body)),

                        "X" => exit_code = body.trim().parse().ok(),

                        "F" => sink(rec("flow", &format!("guest init failed: {body}"))),

                        _ => sink(rec("console", &text)),

                    }

                } else if !text.is_empty() {

                    sink(rec("console", &text));

                }

            }

            Err(mpsc::RecvTimeoutError::Timeout) => {}

            Err(mpsc::RecvTimeoutError::Disconnected) => break,

        }

        if Instant::now() >= deadline && !timed_out {

            timed_out = true;

            let _ = child.kill();

            sink(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 {

        sink(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 { None } else { exit_code }, duration_ms: t(&started), environment, timed_out })

}

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;

                }

            }

        }

    }

}