//! The build cache (oxygenflow#28, spec §5.4): a disk mount that names a
//! `cache` is started from a disk kept by an earlier passing job, rather
//! than empty, so a build carries its compiled dependencies forward.
//!
//! Each runner keeps a pool under `<home>/cache`, by repository, job,
//! cache name and bookmark, each disk named by its key: the BLAKE3 of the
//! image, the job, the cache name and the content of its key files. A job
//! takes, in order:
//!
//! 1. its bookmark's disk with the same key (moved: no copying);
//! 2. its bookmark's newest disk (moved; cargo rebuilds what changed);
//! 3. another bookmark's disk with the same key, then its newest (copied,
//!    so a proposal never takes main's cache away);
//! 4. else an empty disk, as without a cache.
//!
//! A passing job's disk goes back into its bookmark's place; a failing
//! job's is thrown away, so nothing it wrote reaches a later job. The
//! receipt names what the job started from. The pool is kept under its
//! size limit by dropping the disks used longest ago.
//!
//! Beside each disk, `<key>.sources` lists the sources it was built from
//! (oxygenflow#52): each file's content hash and the mtime it had. A job
//! that takes the disk gives every file whose content is unchanged that
//! mtime back, so a build tool that judges freshness by time (cargo)
//! rebuilds only what changed; a changed file keeps the time of its
//! checkout, newer than anything on the disk, so it is always rebuilt.

use anyhow::{Context, Result};
use flow_core::pipeline::Mount;
use std::os::unix::fs::MetadataExt;
use std::path::{Path, PathBuf};
use std::time::SystemTime;

/// The pool's default size limit.
pub const DEFAULT_GB: u64 = 64;

/// The sources a cached disk was built from: path to content hash and
/// mtime, in whole seconds (what the workspace image keeps).
pub type Sources = std::collections::BTreeMap<String, (String, i64)>;

/// What a job started a cached disk from.
pub struct Taken {
    /// For the log and the receipt.
    pub how: String,
    /// The sources it was built from, when they were kept with it.
    pub sources: Option<Sources>,
}

fn sources_of(disk: &Path) -> PathBuf {
    disk.with_extension("sources")
}

fn read_sources(disk: &Path) -> Option<Sources> {
    serde_json::from_slice(&std::fs::read(sources_of(disk)).ok()?).ok()
}

/// Where a job's cached disks come from and go back to.
pub struct Pool {
    root: PathBuf,
    limit: u64,
}

/// Which job: the pool's coordinates other than the key.
pub struct Ctx<'a> {
    pub repo: &'a str,
    pub bookmark: &'a str,
    pub job: &'a str,
    pub image: &'a str,
}

use flow_core::api::segment;

impl Pool {
    pub fn new(home_root: &Path, gb: Option<u64>) -> Pool {
        Pool { root: home_root.join("cache"), limit: gb.unwrap_or(DEFAULT_GB).saturating_mul(1 << 30) }
    }

    fn cache_dir(&self, ctx: &Ctx, name: &str) -> PathBuf {
        self.root.join(segment(ctx.repo)).join(segment(ctx.job)).join(segment(name))
    }

    /// A mount's key: the image, job, cache name and its key files'
    /// content hashes (`hash_of` gives a file's, `None` when absent).
    pub fn key(ctx: &Ctx, m: &Mount, hash_of: impl Fn(&str) -> Option<String>) -> String {
        let mut h = blake3::Hasher::new();
        for part in [ctx.image, ctx.job, m.cache.as_deref().unwrap_or("")] {
            h.update(part.as_bytes());
            h.update(b"\n");
        }
        for k in &m.key {
            h.update(k.as_bytes());
            h.update(b"=");
            h.update(hash_of(k).unwrap_or_else(|| "absent".into()).as_bytes());
            h.update(b"\n");
        }
        h.finalize().to_hex()[..32].to_string()
    }

    /// Start `dest` from the pool for mount `m`, a disk of `m.size_mib`.
    /// What it started from, or `None` (the caller makes an empty disk).
    pub fn take(&self, ctx: &Ctx, m: &Mount, key: &str, dest: &Path) -> Option<Taken> {
        let name = m.cache.as_deref()?;
        let dir = self.cache_dir(ctx, name);
        let size = m.size_mib as u64 * (1 << 20);
        let fits = |p: &Path| std::fs::metadata(p).is_ok_and(|md| md.len() == size);
        let mine = dir.join(segment(ctx.bookmark));
        let exact = mine.join(format!("{key}.ext4"));
        // A moved disk takes its sources with it.
        if fits(&exact) {
            let sources = read_sources(&exact);
            if std::fs::rename(&exact, dest).is_ok() {
                let _ = std::fs::remove_file(sources_of(&exact));
                return Some(Taken { how: format!("cache {name}: this bookmark's, same key"), sources });
            }
        }
        if let Some(p) = newest(&mine, |p| fits(p)) {
            let sources = read_sources(&p);
            if std::fs::rename(&p, dest).is_ok() {
                let _ = std::fs::remove_file(sources_of(&p));
                return Some(Taken { how: format!("cache {name}: this bookmark's, an older key"), sources });
            }
        }
        let others: Vec<PathBuf> =
            std::fs::read_dir(&dir).into_iter().flatten().flatten().map(|e| e.path()).filter(|p| p.is_dir() && *p != mine).collect();
        let pick = others
            .iter()
            .map(|d| d.join(format!("{key}.ext4")))
            .filter(|p| fits(p))
            .max_by_key(|p| modified(p))
            .map(|p| (p, "same key"))
            .or_else(|| others.iter().filter_map(|d| newest(d, |p| fits(p))).max_by_key(|p| modified(p)).map(|p| (p, "an older key")));
        if let Some((p, which)) = pick {
            let from = p.parent().and_then(|d| d.file_name()).map(|n| n.to_string_lossy().into_owned()).unwrap_or_default();
            if sparse_copy(&p, dest).is_ok() {
                return Some(Taken {
                    how: format!("cache {name}: copied from another bookmark's ({from}), {which}"),
                    sources: read_sources(&p),
                });
            }
        }
        None
    }

    /// Keep a passing job's disk for mount `m` under its key, then keep
    /// the pool under its limit.
    pub fn keep(&self, ctx: &Ctx, m: &Mount, key: &str, disk: &Path, sources: &Sources) -> Result<()> {
        let Some(name) = m.cache.as_deref() else { return Ok(()) };
        let mine = self.cache_dir(ctx, name).join(segment(ctx.bookmark));
        std::fs::create_dir_all(&mine)?;
        let dest = mine.join(format!("{key}.ext4"));
        if std::fs::rename(disk, &dest).is_err() {
            sparse_copy(disk, &dest).with_context(|| format!("keep {}", disk.display()))?;
        }
        // Moved disks keep their old time: mark it used now.
        std::fs::File::options().write(true).open(&dest)?.set_modified(SystemTime::now())?;
        let tmp = sources_of(&dest).with_extension("sources-part");
        std::fs::write(&tmp, serde_json::to_vec(sources)?)?;
        std::fs::rename(&tmp, sources_of(&dest))?;
        // One disk per key is enough; older keys of this bookmark go once
        // the pool is over its limit.
        self.evict();
        Ok(())
    }

    /// The repository and job slots holding at least one disk, as
    /// [`flow_core::api::cache_slot`] names them, for the claim
    /// (oxygenflow#51), each with when its newest disk was kept (Unix ms):
    /// a disk's time is set when it is kept, so the newest is the freshest
    /// build (oxygenflow#54).
    pub fn slot_times(&self) -> std::collections::BTreeMap<String, u64> {
        let mut out = std::collections::BTreeMap::new();
        for repo in std::fs::read_dir(&self.root).into_iter().flatten().flatten() {
            for job in std::fs::read_dir(repo.path()).into_iter().flatten().flatten() {
                let mut newest: Option<SystemTime> = None;
                walk(&job.path(), &mut |p| {
                    if p.extension().is_some_and(|e| e == "ext4") {
                        let t = modified(p);
                        newest = Some(newest.map_or(t, |n| n.max(t)));
                    }
                });
                if let Some(t) = newest {
                    let ms = t.duration_since(SystemTime::UNIX_EPOCH).map(|d| d.as_millis() as u64).unwrap_or(0);
                    out.insert(format!("{}/{}", repo.file_name().to_string_lossy(), job.file_name().to_string_lossy()), ms);
                }
            }
        }
        out
    }

    /// Drop the disks used longest ago until the pool's blocks fit.
    pub fn evict(&self) -> u64 {
        let mut disks: Vec<(SystemTime, u64, PathBuf)> = Vec::new();
        walk(&self.root, &mut |p| {
            if p.extension().is_some_and(|e| e == "ext4") {
                if let Ok(md) = std::fs::metadata(p) {
                    disks.push((md.modified().unwrap_or(SystemTime::UNIX_EPOCH), md.blocks() * 512, p.to_path_buf()));
                }
            }
        });
        let mut total: u64 = disks.iter().map(|d| d.1).sum();
        disks.sort();
        let mut freed = 0;
        for (_, bytes, p) in disks {
            if total <= self.limit {
                break;
            }
            if std::fs::remove_file(&p).is_ok() {
                let _ = std::fs::remove_file(sources_of(&p));
                total -= bytes;
                freed += bytes;
            }
        }
        freed
    }
}

fn modified(p: &Path) -> SystemTime {
    std::fs::metadata(p).and_then(|m| m.modified()).unwrap_or(SystemTime::UNIX_EPOCH)
}

fn newest(dir: &Path, ok: impl Fn(&Path) -> bool) -> Option<PathBuf> {
    std::fs::read_dir(dir)
        .ok()?
        .flatten()
        .map(|e| e.path())
        .filter(|p| p.extension().is_some_and(|e| e == "ext4") && ok(p))
        .max_by_key(|p| modified(p))
}

fn walk(dir: &Path, f: &mut dyn FnMut(&Path)) {
    for e in std::fs::read_dir(dir).into_iter().flatten().flatten() {
        let p = e.path();
        if p.is_dir() {
            walk(&p, f);
        } else {
            f(&p);
        }
    }
}

/// Copy a sparse disk without writing its holes.
fn sparse_copy(from: &Path, to: &Path) -> Result<()> {
    let tmp = to.with_extension("part");
    let st = std::process::Command::new("cp").arg("--sparse=always").arg(from).arg(&tmp).status().context("run cp")?;
    if !st.success() {
        let _ = std::fs::remove_file(&tmp);
        anyhow::bail!("cp {} failed", from.display());
    }
    std::fs::rename(&tmp, to)?;
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use flow_core::pipeline::MountSource;

    fn mount(key: &[&str]) -> Mount {
        Mount {
            source: MountSource::Disk,
            target: "/src/target".into(),
            size_mib: 1,
            cache: Some("target".into()),
            key: key.iter().map(|s| s.to_string()).collect(),
        }
    }

    fn disk(path: &Path, fill: &[u8]) {
        let f = std::fs::File::create(path).unwrap();
        f.set_len(1 << 20).unwrap();
        std::os::unix::fs::FileExt::write_at(&f, fill, 0).unwrap();
    }

    fn first(path: &Path) -> Vec<u8> {
        std::fs::read(path).unwrap()[..4].to_vec()
    }

    #[test]
    fn a_kept_disk_starts_the_next_job_of_its_bookmark_then_others_copy_it() {
        let home = tempfile::tempdir().unwrap();
        let work = tempfile::tempdir().unwrap();
        let pool = Pool::new(home.path(), None);
        let m = mount(&["Cargo.lock"]);
        let main = Ctx { repo: "andy/hub", bookmark: "main", job: "test", image: "img" };
        let key = Pool::key(&main, &m, |_| Some("lock-1".into()));
        // Nothing yet: the caller makes an empty disk.
        assert!(pool.take(&main, &m, &key, &work.path().join("a.ext4")).is_none());
        // A passing job's disk is kept and moved into the next job's place.
        disk(&work.path().join("a.ext4"), b"MAIN");
        let built: Sources = [("src/lib.rs".to_string(), ("h1".to_string(), 1_000))].into();
        pool.keep(&main, &m, &key, &work.path().join("a.ext4"), &built).unwrap();
        assert!(!work.path().join("a.ext4").exists(), "moved, not copied");
        let taken = pool.take(&main, &m, &key, &work.path().join("b.ext4")).unwrap();
        assert!(taken.how.contains("same key"), "{}", taken.how);
        assert_eq!(taken.sources.as_ref(), Some(&built), "its sources come with it (#52)");
        assert_eq!(first(&work.path().join("b.ext4")), b"MAIN");
        pool.keep(&main, &m, &key, &work.path().join("b.ext4"), &built).unwrap();
        // A changed key takes the bookmark's newest disk all the same.
        let key2 = Pool::key(&main, &m, |_| Some("lock-2".into()));
        assert_ne!(key, key2);
        let taken = pool.take(&main, &m, &key2, &work.path().join("c.ext4")).unwrap();
        assert!(taken.how.contains("older key"), "{}", taken.how);
        assert_eq!(taken.sources.as_ref(), Some(&built));
        pool.keep(&main, &m, &key2, &work.path().join("c.ext4"), &built).unwrap();
        assert_eq!(
            pool.slot_times().into_keys().collect::<Vec<_>>(),
            vec![flow_core::api::cache_slot("andy/hub", "test")],
            "what the claim reports (#51)"
        );
        // With when it was last kept (#54): just now.
        let now = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_millis() as u64;
        let kept = pool.slot_times()[&flow_core::api::cache_slot("andy/hub", "test")];
        assert!(now.abs_diff(kept) < 60_000, "kept {kept}, now {now}");
        // Another bookmark copies it, leaving main's in place.
        let feature = Ctx { bookmark: "feature", ..main };
        let taken = pool.take(&feature, &m, &key2, &work.path().join("d.ext4")).unwrap();
        assert!(taken.how.contains("copied from another bookmark"), "{}", taken.how);
        assert_eq!(taken.sources.as_ref(), Some(&built), "a copy reads them and leaves them");
        assert_eq!(first(&work.path().join("d.ext4")), b"MAIN");
        let taken = pool.take(&main, &m, &key2, &work.path().join("e.ext4")).unwrap();
        assert!(taken.how.contains("this bookmark's") && taken.sources.is_some(), "main's is still there, with its sources");
        // Another job or a different size never shares it.
        let other_job = Ctx { job: "build", ..main };
        assert!(pool.take(&other_job, &m, &key2, &work.path().join("f.ext4")).is_none());
        let mut bigger = mount(&["Cargo.lock"]);
        bigger.size_mib = 2;
        assert!(pool.take(&feature, &bigger, &key2, &work.path().join("g.ext4")).is_none());
    }

    #[test]
    fn the_key_follows_the_image_the_job_and_the_key_files() {
        let ctx = Ctx { repo: "r", bookmark: "main", job: "test", image: "img-1" };
        let m = mount(&["Cargo.lock", "rust-toolchain.toml"]);
        let k = Pool::key(&ctx, &m, |p| Some(format!("{p}-v1")));
        assert_eq!(k, Pool::key(&ctx, &m, |p| Some(format!("{p}-v1"))));
        assert_ne!(k, Pool::key(&ctx, &m, |p| Some(format!("{p}-v2"))), "a key file changed");
        assert_ne!(k, Pool::key(&Ctx { image: "img-2", ..ctx }, &m, |p| Some(format!("{p}-v1"))), "another image");
        assert_ne!(k, Pool::key(&ctx, &m, |_| None), "a key file missing");
    }

    #[test]
    fn the_pool_drops_the_disks_used_longest_ago() {
        let home = tempfile::tempdir().unwrap();
        let work = tempfile::tempdir().unwrap();
        let pool = Pool { root: home.path().join("cache"), limit: 0 };
        let m = mount(&[]);
        let ctx = Ctx { repo: "r", bookmark: "main", job: "j", image: "i" };
        let full = |p: &Path| {
            let f = std::fs::File::create(p).unwrap();
            std::os::unix::fs::FileExt::write_at(&f, &vec![7u8; 1 << 20], 0).unwrap();
        };
        full(&work.path().join("a.ext4"));
        // Over a limit of nothing: kept, then dropped at once.
        pool.keep(&ctx, &m, "k1", &work.path().join("a.ext4"), &Sources::new()).unwrap();
        assert!(pool.take(&ctx, &m, "k1", &work.path().join("b.ext4")).is_none());
        let roomy = Pool { root: home.path().join("cache"), limit: 1 << 30 };
        full(&work.path().join("c.ext4"));
        roomy.keep(&ctx, &m, "k2", &work.path().join("c.ext4"), &Sources::new()).unwrap();
        assert!(roomy.take(&ctx, &m, "k2", &work.path().join("d.ext4")).is_some(), "under the limit it stays");
    }
}