//! Checkpoints (oxygenflow#77, docs/handover.md §7): the directories a job

//! declares with `checkpoint "/src/…"`, saved when its runner hands it back

//! and restored before the next attempt runs the same exec, so a tool that

//! skips finished work (cargo, make) carries on where it stopped.

//!

//! The format is the runner's own, so CI builds need no new crate: a magic

//! line, then entries (a kind byte, the path under /src, mode, modification

//! time and, for a file, its bytes). Files keep their times, since cargo's

//! fingerprints read them. Links are left out. Unpacking accepts only paths

//! under the job's own checkpoints, without `.` or `..`.

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

use flow_core::pipeline::{Job, MountSource};

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

use std::time::{Duration, SystemTime, UNIX_EPOCH};

/// At most this much is saved: the hub keeps it as it keeps an artifact.

pub const MAX_BYTES: usize = 512 << 20;

const MAGIC: &[u8] = b"OXYFLOW-CHECKPOINT 1\n";

const DIR: u8 = 0;

const FILE: u8 = 1;

/// What was saved or restored.

#[derive(Debug, Default, PartialEq, Eq)]

pub struct Counts {

    pub files: usize,

    pub dirs: usize,

    pub bytes: u64,

}

/// Where checkpoint `path` (under /src) is on this host, for a job of a

/// directory tier (WASM, syscall): in a disk mount's directory when a mount

/// holds it, else in the job's tree. `None` for one under a tmpfs, which

/// the syscall tier keeps in memory.

pub fn host_dir(path: &str, job: &Job, tree: &Path, scratch: &Path, wasm: bool) -> Option<PathBuf> {

    let under = |t: &str| path == t || path.starts_with(&format!("{}/", t.trim_end_matches('/')));

    let mut disks = 0;

    let mut best: Option<(usize, PathBuf)> = None;

    for (i, m) in job.mounts.iter().enumerate() {

        let dir = match m.source {

            MountSource::Disk => {

                disks += 1;

                Some(crate::syscall::disk_dir(scratch, disks - 1))

            }

            MountSource::Tmpfs if wasm => Some(scratch.join(format!("tmpfs-{i}"))),

            MountSource::Tmpfs => None,

        };

        if under(&m.target) && best.as_ref().is_none_or(|(len, _)| m.target.len() > *len) {

            let rest = path[m.target.trim_end_matches('/').len()..].trim_start_matches('/');

            best = Some((m.target.len(), dir.map(|d| d.join(rest)).unwrap_or_default()));

        }

    }

    match best {

        Some((_, d)) if d.as_os_str().is_empty() => None,

        Some((_, d)) => Some(d),

        None => Some(tree.join(path.trim_start_matches("/src").trim_start_matches('/'))),

    }

}

/// Pack `dirs` (each a checkpoint path and where it is on this host).

pub fn pack(dirs: &[(String, PathBuf)]) -> Result<(Vec<u8>, Counts)> {

    let mut out = MAGIC.to_vec();

    let mut n = Counts::default();

    for (path, host) in dirs {

        if !host.is_dir() {

            continue;

        }

        let mut stack = vec![(host.clone(), path.trim_end_matches('/').to_string())];

        while let Some((at, name)) = stack.pop() {

            let meta = std::fs::symlink_metadata(&at)?;

            if meta.is_dir() {

                entry(&mut out, DIR, &name, &meta, &[]);

                n.dirs += 1;

                let mut children: Vec<_> = std::fs::read_dir(&at)?.collect::<std::io::Result<_>>()?;

                children.sort_by_key(|e| e.file_name());

                for c in children.into_iter().rev() {

                    let Some(child) = c.file_name().to_str().map(String::from) else { continue };

                    stack.push((c.path(), format!("{name}/{child}")));

                }

            } else if meta.is_file() {

                let bytes = std::fs::read(&at).with_context(|| format!("read {}", at.display()))?;

                n.files += 1;

                n.bytes += bytes.len() as u64;

                entry(&mut out, FILE, &name, &meta, &bytes);

            }

            if out.len() > MAX_BYTES {

                bail!("the checkpoint is over {} MiB", MAX_BYTES >> 20);

            }

        }

    }

    Ok((out, n))

}

fn entry(out: &mut Vec<u8>, kind: u8, path: &str, meta: &std::fs::Metadata, bytes: &[u8]) {

    use std::os::unix::fs::PermissionsExt;

    let since = meta.modified().ok().and_then(|t| t.duration_since(UNIX_EPOCH).ok()).unwrap_or_default();

    out.push(kind);

    out.extend((path.len() as u32).to_le_bytes());

    out.extend(path.as_bytes());

    out.extend(meta.permissions().mode().to_le_bytes());

    out.extend(since.as_secs().to_le_bytes());

    out.extend(since.subsec_nanos().to_le_bytes());

    out.extend((bytes.len() as u64).to_le_bytes());

    out.extend(bytes);

}

/// Restore a checkpoint: each entry at `place(checkpoint)` joined with the

/// rest of its path, for the job's own `checkpoints` only.

pub fn unpack(bytes: &[u8], checkpoints: &[String], place: impl Fn(&str) -> Option<PathBuf>) -> Result<Counts> {

    use std::os::unix::fs::PermissionsExt;

    let Some(mut rest) = bytes.strip_prefix(MAGIC) else { bail!("not a checkpoint") };

    let mut n = Counts::default();

    let take = |rest: &mut &[u8], k: usize| -> Result<Vec<u8>> {

        if rest.len() < k {

            bail!("the checkpoint is cut short");

        }

        let (a, b) = rest.split_at(k);

        *rest = b;

        Ok(a.to_vec())

    };

    while !rest.is_empty() {

        let kind = take(&mut rest, 1)?[0];

        let len = u32::from_le_bytes(take(&mut rest, 4)?.try_into().unwrap()) as usize;

        let path = String::from_utf8(take(&mut rest, len)?).context("a path that is not UTF-8")?;

        let mode = u32::from_le_bytes(take(&mut rest, 4)?.try_into().unwrap());

        let secs = u64::from_le_bytes(take(&mut rest, 8)?.try_into().unwrap());

        let nanos = u32::from_le_bytes(take(&mut rest, 4)?.try_into().unwrap());

        let size = u64::from_le_bytes(take(&mut rest, 8)?.try_into().unwrap()) as usize;

        let data = take(&mut rest, size)?;

        // Under one of the job's checkpoints, without . or .. on the way.

        let Some(root) = checkpoints.iter().map(|c| c.trim_end_matches('/')).find(|c| path == *c || path.starts_with(&format!("{c}/")))

        else {

            bail!("{path} is not under the job's checkpoints");

        };

        let below = path[root.len()..].trim_start_matches('/');

        if below.split('/').any(|p| p == "." || p == ".." || (p.is_empty() && !below.is_empty())) {

            bail!("{path}: . or .. in a checkpoint path");

        }

        let Some(base) = place(root) else { continue };

        let at = if below.is_empty() { base } else { base.join(below) };

        match kind {

            DIR => {

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

                n.dirs += 1;

            }

            FILE => {

                if let Some(parent) = at.parent() {

                    std::fs::create_dir_all(parent)?;

                }

                // Never through a link already there.

                if std::fs::symlink_metadata(&at).is_ok_and(|m| !m.is_file()) {

                    bail!("{path}: something other than a file is in the way");

                }

                std::fs::write(&at, &data)?;

                let f = std::fs::File::options().write(true).open(&at)?;

                f.set_permissions(std::fs::Permissions::from_mode(mode & 0o777))?;

                f.set_modified(SystemTime::UNIX_EPOCH + Duration::new(secs, nanos))?;

                n.files += 1;

                n.bytes += data.len() as u64;

            }

            other => bail!("{path}: unknown entry kind {other}"),

        }

    }

    Ok(n)

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn a_checkpoint_comes_back_with_its_times_and_only_where_it_belongs() {

        let from = tempfile::tempdir().unwrap();

        let target = from.path().join("target");

        std::fs::create_dir_all(target.join("debug/deps")).unwrap();

        std::fs::write(target.join("debug/deps/libx.rlib"), b"built").unwrap();

        std::fs::write(target.join("debug/.fingerprint"), b"fp").unwrap();

        let old = SystemTime::UNIX_EPOCH + Duration::new(1_700_000_000, 123_000_000);

        std::fs::File::options().write(true).open(target.join("debug/deps/libx.rlib")).unwrap().set_modified(old).unwrap();

        let (bytes, n) = pack(&[("/src/target".into(), target.clone())]).unwrap();

        assert_eq!((n.files, n.dirs, n.bytes), (2, 3, 7));

        let to = tempfile::tempdir().unwrap();

        let place = |root: &str| (root == "/src/target").then(|| to.path().join("t"));

        let back = unpack(&bytes, &["/src/target".into()], place).unwrap();

        assert_eq!(back, n);

        let restored = to.path().join("t/debug/deps/libx.rlib");

        assert_eq!(std::fs::read(&restored).unwrap(), b"built");

        assert_eq!(std::fs::metadata(&restored).unwrap().modified().unwrap(), old, "cargo reads the times");

        // Not the job's checkpoint, or a path that climbs out: refused.

        assert!(unpack(&bytes, &["/src/other".into()], place).is_err());

        let mut evil = MAGIC.to_vec();

        let meta = std::fs::metadata(target.join("debug/.fingerprint")).unwrap();

        entry(&mut evil, FILE, "/src/target/../escape", &meta, b"x");

        assert!(unpack(&evil, &["/src/target".into()], place).is_err());

        assert!(!to.path().join("escape").exists());

        assert!(unpack(b"junk", &["/src/target".into()], place).is_err());

    }

    #[test]

    fn a_checkpoint_under_a_disk_mount_is_in_the_mount_s_directory() {

        let p = flow_core::Pipeline::parse(

            "schema-version 1\nsandbox \"w\" {\n    runtime \"wasm\"\n}\njob \"j\" {\n    sandbox \"w\"\n    mount {\n        source \"ephemeral://disk\"\n        target \"/src/target\"\n        size \"1GiB\"\n    }\n    checkpoint \"/src/target\" \"/src/state\"\n    exec \"x.wasm\"\n}\n",

        )

        .unwrap();

        let job = &p.jobs["j"];

        assert_eq!(job.checkpoints, ["/src/state", "/src/target"]);

        let (tree, scratch) = (Path::new("/w/tree"), Path::new("/w"));

        assert_eq!(host_dir("/src/target", job, tree, scratch, true), Some(crate::syscall::disk_dir(scratch, 0)));

        assert_eq!(host_dir("/src/state", job, tree, scratch, true), Some(PathBuf::from("/w/tree/state")));

    }

}