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

/// 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 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 = guest_spec(job);
    // 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 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 })
}

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