//! `flow`: check pipelines, run them locally, and run jobs for OxygenHub

//! (docs/spec.md, docs/m0.md).

mod cpio;

/// The job-starting code `flow-init` uses in a microVM, shared with the

/// syscall tier.

#[path = "../../flow-init/src/job.rs"]

mod guest;

mod home;

mod microvm;

mod outputs;

mod runner;

mod syscall;

mod wait;

mod workspace;

use anyhow::{bail, Context, Result};

use clap::{Parser, Subcommand};

use flow_core::api::Record;

use flow_core::{Pipeline, Runtime, PIPELINE_PATH};

use home::Home;

use std::collections::BTreeSet;

use std::path::{Path, PathBuf};

use std::time::Duration;

#[derive(Parser)]

#[command(name = "flow", version = env!("OXY_VERSION_LONG"), about = "OxygenFlow: hermetic pipelines for Oxygen repositories")]

struct Cli {

    /// The Flow home (default $FLOW_HOME, else ~/.local/share/flow).

    #[arg(long, global = true)]

    home: Option<PathBuf>,

    #[command(subcommand)]

    cmd: Cmd,

}

#[derive(Subcommand)]

enum RunnerCmd {

    /// Register this runner with a hub using a one-time token from its Admin → Runners page

    Register {

        /// The hub, e.g. https://earth.example:7343

        hub: String,

        /// The registration token (read from stdin when not given)

        #[arg(long)]

        token: Option<String>,

        /// The runner's name (default: the host name)

        #[arg(long)]

        name: Option<String>,

        /// The hub's CA certificate, when self-signed

        #[arg(long)]

        ca: Option<PathBuf>,

        /// A label the runner offers (repeatable)

        #[arg(long)]

        label: Vec<String>,

    },

}

#[derive(Subcommand)]

enum Cmd {

    /// Validate a pipeline and print its jobs and hash

    Check {

        /// The pipeline file (default: .flow/pipeline.kdl)

        file: Option<PathBuf>,

    },

    /// Print a file's BLAKE3, for pinning in flow.toml

    Hash { file: PathBuf },

    /// Rootfs images for sandboxes

    Rootfs {

        #[command(subcommand)]

        cmd: RootfsCmd,

    },

    /// Run the pipeline on this directory's files, locally

    Run {

        /// Only these jobs (and what they depend on)

        jobs: Vec<String>,

        /// The directory to build (default: the current one)

        #[arg(long, default_value = ".")]

        dir: PathBuf,

        /// Minutes before a job is stopped

        #[arg(long, default_value_t = 60)]

        timeout: u64,

        /// A variable for every job, as the hub's settings would give it;

        /// the pipeline's own `env` wins

        #[arg(long = "var", value_name = "NAME=VALUE")]

        vars: Vec<String>,

        /// A secret a job's `secret` block names, as the hub would supply

        /// it; mounted as a file, never in the environment

        #[arg(long = "secret", value_name = "NAME=VALUE")]

        secrets: Vec<String>,

    },

    /// Claim and run jobs for the hub in flow.toml's [runner] (or the

    /// runner.toml `flow runner register` wrote)

    Runner {

        #[command(subcommand)]

        cmd: Option<RunnerCmd>,

    },

    /// Follow the checks of a pushed state on OxygenHub until they finish:

    /// exit 0 when all passed, 1 when one failed, 2 at the time limit.

    /// Uses $OXY_TOKEN and $OXY_CA as oxy does.

    Wait {

        /// The repository's URL on the hub (https://hub/owner/name)

        url: String,

        /// The state hash the push moved the bookmark to (`oxy push --json` prints it)

        state: String,

        /// Minutes to wait at most

        #[arg(long, default_value_t = 60)]

        timeout: u64,

    },

    /// Show which isolation tiers this machine can run jobs in, and why

    /// the others cannot

    Tiers,

    /// The syscall tier's sandbox helper, started by the runner for each job

    #[command(name = "__sandbox", hide = true)]

    Sandbox { spec: PathBuf },

    /// Check that this machine can run the syscall tier

    #[command(name = "__sandbox-probe", hide = true)]

    SandboxProbe,

}

#[derive(Subcommand)]

enum RootfsCmd {

    /// Make an ext4 rootfs image from a root file system tarball and store it by hash

    Import {

        tarball: PathBuf,

        /// Refuse the tarball unless its SHA-256 matches (as its publisher states it)

        #[arg(long)]

        sha256: Option<String>,

        /// The image's own environment, NAME=value (repeatable): what a

        /// Dockerfile's ENV sets and `docker export` loses. Jobs start from

        /// it; their sandbox's and their own `env` override it

        #[arg(long = "env", value_name = "NAME=VALUE")]

        env: Vec<String>,

    },

    /// List the images this home holds

    List,

}

fn main() {

    let cli = Cli::parse();

    if let Err(e) = run(cli) {

        eprintln!("flow: {e:#}");

        std::process::exit(1);

    }

}

fn run(cli: Cli) -> Result<()> {

    match cli.cmd {

        Cmd::Sandbox { spec } => std::process::exit(syscall::sandbox_main(&spec)),

        Cmd::SandboxProbe => std::process::exit(syscall::probe_main()),

        Cmd::Check { file } => check(&file.unwrap_or_else(|| PathBuf::from(PIPELINE_PATH))),

        Cmd::Hash { file } => {

            println!("{}", home::hash_file(&file)?);

            Ok(())

        }

        Cmd::Rootfs { cmd } => {

            let home = Home::locate(cli.home.as_deref())?;

            match cmd {

                RootfsCmd::Import { tarball, sha256, env } => rootfs_import(&home, &tarball, sha256.as_deref(), &env),

                RootfsCmd::List => {

                    for e in std::fs::read_dir(home.root.join("cas"))?.flatten() {

                        let size = e.metadata().map(|m| m.len()).unwrap_or(0);

                        if let Some(h) = e.file_name().to_string_lossy().strip_suffix(".ext4") {

                            println!("cas://blob/{h}  {} MiB", size >> 20);

                        }

                    }

                    Ok(())

                }

            }

        }

        Cmd::Run { jobs, dir, timeout, vars, secrets } => {

            let home = Home::locate(cli.home.as_deref())?;

            let pairs = |what: &str, list: &[String]| {

                list.iter()

                    .map(|e| {

                        e.split_once('=')

                            .map(|(k, v)| (k.to_string(), v.to_string()))

                            .ok_or_else(|| anyhow::anyhow!("--{what} {e:?}: expected NAME=VALUE"))

                    })

                    .collect::<Result<Vec<_>>>()

            };

            let (vars, secrets) = (pairs("var", &vars)?, pairs("secret", &secrets)?);

            run_local(&home, &dir, &jobs, Duration::from_secs(timeout * 60), &vars, &secrets)

        }

        Cmd::Runner { cmd: Some(RunnerCmd::Register { hub, token, name, ca, label }) } => {

            let home = Home::locate(cli.home.as_deref())?;

            if std::env::var_os("OXY_HOME").is_none() {

                std::env::set_var("OXY_HOME", home.root.join("oxy"));

            }

            let name = name.unwrap_or_else(|| {

                std::fs::read_to_string("/etc/hostname").map(|h| h.trim().to_string()).unwrap_or_else(|_| "flow".into())

            });

            let token = match token {

                Some(t) => t,

                None => {

                    eprintln!("Paste the registration token from the hub's Admin → Runners, then press Enter:");

                    let mut line = String::new();

                    std::io::stdin().read_line(&mut line)?;

                    line.trim().to_string()

                }

            };

            let got = runner::register(&home, &hub, &token, &name, ca, label)?;

            println!("registered as {} (runner {}); start it with `flow runner`", got.name, got.id);

            Ok(())

        }

        Cmd::Runner { cmd: None } => {

            let home = Home::locate(cli.home.as_deref())?;

            // The runner's oxy state (identity, user config) stays in its home.

            if std::env::var_os("OXY_HOME").is_none() {

                std::env::set_var("OXY_HOME", home.root.join("oxy"));

            }

            runner::serve(&home)

        }

        Cmd::Tiers => {

            let home = Home::locate(cli.home.as_deref())?;

            match runner::tiers(&home) {

                Ok((have, missing)) => {

                    for t in have {

                        println!("{t}: available");

                    }

                    for why in missing {

                        println!("{}", why.replacen(": ", ": unavailable: ", 1));

                    }

                }

                Err(why) => println!("{why}"),

            }

            Ok(())

        }

        Cmd::Wait { url, state, timeout } => {

            let token = std::env::var("OXY_TOKEN").ok();

            let ca = std::env::var_os("OXY_CA").map(PathBuf::from);

            let code =

                wait::wait(&url, &state, token.as_deref(), ca.as_deref(), Duration::from_secs(timeout * 60), Duration::from_secs(15))?;

            std::process::exit(code);

        }

    }

}

fn load(file: &Path) -> Result<Pipeline> {

    let src = std::fs::read_to_string(file).with_context(|| format!("read {}", file.display()))?;

    Pipeline::parse(&src).map_err(|p| anyhow::anyhow!("{} is not valid:\n{p}", file.display()))

}

fn check(file: &Path) -> Result<()> {

    let p = load(file)?;

    let depths = p.depths();

    for j in p.order() {

        let sb = &p.sandboxes[&j.sandbox];

        let after = if j.depends_on.is_empty() { String::new() } else { format!(" after {}", j.depends_on.join(", ")) };

        println!("{}{}  {} [{}]{after}", "  ".repeat(depths[j.name.as_str()]), j.name, j.exec.join(" "), sb.runtime.as_str());

    }

    println!("pipeline {}", hex::encode(p.hash()));

    Ok(())

}

/// The image environment file flow-init reads (`/etc/flow/env`): one

/// `NAME=value` per line, names checked as a pipeline's are.

fn image_env_file(env: &[String]) -> Result<String> {

    let mut out = String::new();

    for e in env {

        let Some((k, v)) = e.split_once('=') else { bail!("--env {e:?}: expected NAME=value") };

        let ok = !k.is_empty() && !k.starts_with(|c: char| c.is_ascii_digit()) && k.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_');

        if !ok || v.contains('\n') || v.contains('\0') {

            bail!("--env {e:?}: a name is letters, digits and _, and a value is one line");

        }

        out.push_str(&format!("{k}={v}\n"));

    }

    Ok(out)

}

/// The architecture an ELF file was built for, by its header.

fn elf_arch(path: &Path) -> Option<&'static str> {

    let bytes = std::fs::read(path).ok()?;

    if bytes.len() < 20 || &bytes[..4] != b"\x7fELF" {

        return None;

    }

    let little = bytes[5] == 1;

    let machine = if little { u16::from_le_bytes([bytes[18], bytes[19]]) } else { u16::from_be_bytes([bytes[18], bytes[19]]) };

    Some(match machine {

        0x3E => "x86_64",

        0xB7 => "aarch64",

        0xF3 => "riscv64",

        0x03 => "x86",

        0x28 => "arm",

        0x15 => "powerpc64",

        0x16 => "s390x",

        _ => return None,

    })

}

/// The architecture of an unpacked rootfs: its shell's, or the first

/// program's found under bin.

fn tree_arch(tree: &Path) -> Option<&'static str> {

    for rel in ["bin/sh", "bin/busybox", "usr/bin/env", "bin/bash", "usr/bin/sh"] {

        let p = tree.join(rel);

        let p = std::fs::canonicalize(&p).unwrap_or(p);

        if let Some(a) = elf_arch(&p) {

            return Some(a);

        }

    }

    None

}

fn rootfs_import(home: &Home, tarball: &Path, sha256: Option<&str>, env: &[String]) -> Result<()> {

    let env_file = image_env_file(env)?;

    if let Some(want) = sha256 {

        let out = std::process::Command::new("sha256sum").arg(tarball).output().context("run sha256sum")?;

        let got = String::from_utf8_lossy(&out.stdout).split_whitespace().next().unwrap_or("").to_string();

        if got != want.to_lowercase() {

            bail!("{} has SHA-256 {got}, expected {want}", tarball.display());

        }

    }

    let tmp = home.root.join("work").join(format!("rootfs-{}", std::process::id()));

    let _ = std::fs::remove_dir_all(&tmp);

    let tree = tmp.join("tree");

    std::fs::create_dir_all(&tree)?;

    let result = (|| -> Result<String> {

        let status = std::process::Command::new("tar").arg("-xf").arg(tarball).arg("-C").arg(&tree).arg("--no-same-owner").status()?;

        if !status.success() {

            bail!("tar could not unpack {}", tarball.display());

        }

        // The image's architecture must be this machine's: there is no

        // emulation tier (spec §11.2).

        let arch = tree_arch(&tree).unwrap_or("unknown");

        if arch != "unknown" && arch != std::env::consts::ARCH {

            bail!(

                "{} is a {arch} image; this runner is {}, and jobs here run no other architecture",

                tarball.display(),

                std::env::consts::ARCH

            );

        }

        // Mount points flow-init needs (docs/m0.md §2).

        for d in ["src", "tmp", "proc", "dev", "run"] {

            std::fs::create_dir_all(tree.join(d))?;

        }

        if !env_file.is_empty() {

            std::fs::create_dir_all(tree.join("etc/flow"))?;

            std::fs::write(tree.join("etc/flow/env"), &env_file)?;

        }

        let used: u64 = ignore::WalkBuilder::new(&tree)

            .hidden(false)

            .ignore(false)

            .git_ignore(false)

            .build()

            .flatten()

            .filter_map(|e| e.metadata().ok())

            .map(|m| m.len().max(4096))

            .sum();

        let mib = used * 13 / 10 / (1 << 20) + 64;

        let img = tmp.join("rootfs.ext4");

        let status = std::process::Command::new("mkfs.ext4")

            .args(["-q", "-F", "-L", "rootfs", "-E", "root_owner=0:0", "-d"])

            .arg(&tree)

            .arg(&img)

            .arg(format!("{mib}M"))

            .status()

            .context("run mkfs.ext4")?;

        if !status.success() {

            bail!("mkfs.ext4 failed");

        }

        let hash = home::hash_file(&img)?;

        std::fs::rename(&img, home.rootfs(&hash))?;

        let manifest = serde_json::json!({ "arch": if arch == "unknown" { std::env::consts::ARCH } else { arch }, "imported_ms": std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_millis() as u64).unwrap_or(0) });

        std::fs::write(home.image_manifest(&hash), manifest.to_string())?;

        Ok(hash)

    })();

    let _ = std::fs::remove_dir_all(&tmp);

    let hash = result?;

    println!("rootfs \"cas://blob/{hash}\"");

    Ok(())

}

fn run_local(

    home: &Home,

    dir: &Path,

    only: &[String],

    timeout: Duration,

    vars: &[(String, String)],

    secrets: &[(String, String)],

) -> Result<()> {

    let p = load(&dir.join(PIPELINE_PATH))?;

    // The named jobs and everything they depend on.

    let mut wanted: BTreeSet<String> = only.iter().cloned().collect();

    for name in only {

        if !p.jobs.contains_key(name) {

            bail!("no job {name:?} in the pipeline");

        }

    }

    if !only.is_empty() {

        loop {

            let more: Vec<String> = wanted.iter().flat_map(|j| p.jobs[j].depends_on.clone()).filter(|d| !wanted.contains(d)).collect();

            if more.is_empty() {

                break;

            }

            wanted.extend(more);

        }

    }

    let scratch = home.root.join("work").join(format!("local-{}", std::process::id()));

    let _ = std::fs::remove_dir_all(&scratch);

    std::fs::create_dir_all(scratch.join("tree"))?;

    workspace::copy_tree(dir, &scratch.join("tree"))?;

    let mut failed: BTreeSet<String> = BTreeSet::new();

    let mut worst = 0;

    for job in p.order() {

        if !only.is_empty() && !wanted.contains(&job.name) {

            continue;

        }

        if let Some(d) = job.depends_on.iter().find(|d| failed.contains(*d)) {

            println!("── {} skipped: {d} failed", job.name);

            failed.insert(job.name.clone());

            continue;

        }

        let sb = &p.sandboxes[&job.sandbox];

        if sb.runtime == Runtime::Wasm {

            bail!("job {}: runtime {} is not available in this build yet", job.name, sb.runtime.as_str());

        }

        println!("── {} ({})", job.name, job.exec.join(" "));

        let js = scratch.join(&job.name);

        std::fs::create_dir_all(&js)?;

        let job = &job

            .in_sandbox(sb)

            .with_flow_env(&[("FLOW_JOB".into(), job.name.clone()), ("FLOW_RUN".into(), "local".into())])

            .with_flow_env(vars);

        // Identities come from a hub: a local run has none.

        for i in &job.identities {

            println!("── {}: identity {:?} is not available locally; {} will be missing", job.name, i.audience, i.mount_path);

        }

        // Each secret the job declares must be given, as the hub would.

        let inputs = job

            .secrets

            .iter()

            .map(|s| {

                let value = secrets.iter().find(|(k, _)| *k == s.name).map(|(_, v)| v.clone());

                value

                    .map(|value| flow_core::api::SecretInput { name: s.name.clone(), path: s.mount_path.clone(), value })

                    .ok_or_else(|| anyhow::anyhow!("job {}: secret {:?} is not given; pass --secret {}=VALUE", job.name, s.name, s.name))

            })

            .collect::<Result<Vec<_>>>()?;

        // Each job starts from the same files: jobs do not share writes.

        let out = if sb.runtime == Runtime::Syscall {

            workspace::copy_tree(&scratch.join("tree"), &js.join("tree"))?;

            let group = format!("local-{}-{}", std::process::id(), js.file_name().unwrap_or_default().to_string_lossy());