//! 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]) {
    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(crate::portable::mode(meta).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> {
    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)?;
                crate::portable::set_mode(&at, 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")));
    }
}