roundabout

use std::{

    collections::{BTreeMap, BTreeSet},

    fs,

    os::unix::fs::MetadataExt,

    path::{Path, PathBuf},

    process::Command,

    time::{SystemTime, UNIX_EPOCH},

};

use serde::{Deserialize, Serialize};

use crate::{

    client,

    config::{Config, Rule},

    paths,

    proto::{self, Event, Reader, Request},

    store, sys,

    units::Bytes,

};

#[derive(Clone, Debug, Serialize, Deserialize)]

pub struct Entry {

    pub root: PathBuf,

    pub rule: String,

    pub last_used: u64,

    pub pinned: bool,

    pub size: Bytes,

    pub sized_at: u64,

}

type Registry = BTreeMap<PathBuf, Entry>;

fn now() -> u64 {

    SystemTime::now()

        .duration_since(UNIX_EPOCH)

        .map_or(0, |d| d.as_secs())

}

fn matches(rule: &Rule, name: &str) -> bool {

    match rule.name.strip_suffix('*') {

        Some(prefix) => name.starts_with(prefix),

        None => name == rule.name,

    }

}

fn discover<'a>(cfg: &'a Config, root: &Path) -> Vec<(PathBuf, &'a Rule)> {

    let Ok(dir) = fs::read_dir(root) else {

        return Vec::new();

    };

    let names: Vec<String> = dir

        .flatten()

        .map(|e| e.file_name().to_string_lossy().into_owned())

        .collect();

    names

        .iter()

        .filter_map(|name| {

            let rule = cfg.artifact_rules().find(|r| {

                matches(r, name) && r.marker.as_ref().is_none_or(|m| root.join(m).exists())

            })?;

            Some((root.join(name), rule))

        })

        .collect()

}

/// symlinks are followed at the top only, so a nix `result` counts its store

/// path. hardlinks count once, the store is full of them

fn size_of(path: &Path) -> Bytes {

    let mut seen = BTreeSet::new();

    let mut total = 0;

    let mut stack = vec![fs::canonicalize(path).unwrap_or_else(|_| path.to_path_buf())];

    while let Some(p) = stack.pop() {

        let Ok(meta) = fs::symlink_metadata(&p) else {

            continue;

        };

        if meta.nlink() > 1 && !seen.insert((meta.dev(), meta.ino())) {

            continue;

        }

        total += meta.blocks() * 512;

        if meta.is_dir()

            && let Ok(dir) = fs::read_dir(&p)

        {

            stack.extend(dir.flatten().map(|e| e.path()));

        }

    }

    Bytes(total)

}

/// only builds bump `last_used`. checking sizes isn't using them

pub fn register(cfg: &Config, root: &Path, used: bool) {

    let found = discover(cfg, root);

    if found.is_empty() {

        return;

    }

    let _ = store::update("artifacts", |reg: &mut Registry| {

        for (path, rule) in found {

            reg.entry(path)

                .and_modify(|e| {

                    if used {

                        e.last_used = now();

                    }

                })

                .or_insert_with(|| Entry {

                    root: root.to_path_buf(),

                    rule: rule.name.clone(),

                    last_used: now(),

                    pinned: false,

                    size: Bytes::ZERO,

                    sized_at: 0,

                });

        }

    });

}

pub fn pin(path: &Path, pinned: bool) -> Result<(), String> {

    let path = fs::canonicalize(path).map_err(|e| format!("{}: {e}", path.display()))?;

    let root = paths::project_root(path.parent().unwrap_or(&path));

    store::update("artifacts", |reg: &mut Registry| {

        reg.entry(path.clone())

            .or_insert_with(|| Entry {

                root: root.clone(),

                rule: "manual".into(),

                last_used: now(),

                pinned,

                size: Bytes::ZERO,

                sized_at: 0,

            })

            .pinned = pinned;

    })

    .map_err(|e| e.to_string())

}

fn active_roots() -> BTreeSet<PathBuf> {

    let Some(mut s) = client::connect() else {

        return BTreeSet::new();

    };

    if proto::send(&mut s, &Request::Status).is_err() {

        return BTreeSet::new();

    }

    match Reader::new(&s).and_then(|mut r| r.recv::<Event>()) {

        Ok(Some(Event::Status(st))) => st

            .leases

            .iter()

            .map(|l| paths::project_root(&l.cwd))

            .collect(),

        _ => BTreeSet::new(),

    }

}

fn refresh(cfg: &Config, root: Option<&Path>, max_age: u64) -> Registry {

    if let Some(root) = root {

        register(cfg, root, false);

    }

    store::update("artifacts", |reg: &mut Registry| {

        reg.retain(|p, _| fs::symlink_metadata(p).is_ok());

        for (path, e) in reg.iter_mut() {

            if now().saturating_sub(e.sized_at) >= max_age {

                e.size = size_of(path);

                e.sized_at = now();

            }

        }

        reg.clone()

    })

    .unwrap_or_default()

}

pub struct Report {

    pub project: Option<(PathBuf, Bytes)>,

    pub global: Bytes,

    pub over_project: bool,

    pub over_global: bool,

    pub entries: Registry,

}

impl Report {

    pub const fn over(&self) -> bool {

        self.over_project || self.over_global

    }

    pub const fn remedy(&self) -> &'static str {

        if self.over_global {

            "roundabout artifacts gc --cap"

        } else {

            "roundabout artifacts gc --project --cap"

        }

    }

    pub fn describe(&self, cfg: &Config) -> String {

        let mut out = Vec::new();

        if let Some((root, size)) = &self.project

            && self.over_project

        {

            out.push(format!(

                "{} holds {size} of build artifacts, cap is {}",

                root.display(),

                cfg.artifacts.project_cap

            ));

        }

        if self.over_global {

            out.push(format!(

                "artifacts total {} across projects, cap is {}",

                self.global, cfg.artifacts.global_cap

            ));

        }

        out.join(". ")

    }

}

const RECHECK_SECS: u64 = 60;

/// the cache can predate someone's `rm`, so being over earns a second look

/// before anyone gets nagged

pub fn check(cfg: &Config, root: Option<&Path>, fresh: bool) -> Report {

    let ttl = if fresh {

        0

    } else {

        cfg.artifacts.size_ttl.0.as_secs()

    };

    let report = measure(cfg, root, ttl);

    if report.over() && ttl > RECHECK_SECS {

        return measure(cfg, root, RECHECK_SECS);

    }

    report

}

fn measure(cfg: &Config, root: Option<&Path>, max_age: u64) -> Report {

    let entries = refresh(cfg, root, max_age);

    let global = entries.values().map(|e| e.size).sum();

    let project = root.map(|r| {

        let size = entries

            .values()

            .filter(|e| e.root == r)

            .map(|e| e.size)

            .sum();

        (r.to_path_buf(), size)

    });

    Report {

        over_project: project

            .as_ref()

            .is_some_and(|(_, s)| *s > cfg.artifacts.project_cap),

        over_global: global > cfg.artifacts.global_cap,

        project,

        global,

        entries,

    }

}

fn remove(cfg: &Config, path: &Path, e: &Entry) -> Result<(), String> {

    let rule = cfg.artifact_rules().find(|r| r.name == e.rule);

    if let Some(clean) = rule.and_then(|r| r.clean.as_ref())

        && let Some((prog, args)) = clean.split_first()

    {

        let ok = Command::new(prog)

            .args(args)

            .current_dir(&e.root)

            .status()

            .is_ok_and(|s| s.success());

        return if ok {

            Ok(())

        } else {

            Err(format!("`{}` failed", clean.join(" ")))

        };

    }

    let meta = fs::symlink_metadata(path).map_err(|err| err.to_string())?;

    if meta.is_dir() {

        fs::remove_dir_all(path)

    } else {

        fs::remove_file(path)

    }

    .map_err(|err| format!("{}: {err}", path.display()))

}

pub struct Gc {

    pub freed: Bytes,

    pub removed: Vec<PathBuf>,

    pub failed: Vec<String>,

}

pub fn gc(cfg: &Config, project: Option<&Path>, to_cap: bool, dry_run: bool) -> Gc {

    let max_age = if to_cap {

        RECHECK_SECS

    } else {

        cfg.artifacts.size_ttl.0.as_secs()

    };

    let entries = refresh(cfg, project, max_age);

    let active = active_roots();

    let idle_after = cfg.artifacts.idle_after.0.as_secs();

    let movable =

        |e: &Entry| !e.pinned && !active.contains(&e.root) && project.is_none_or(|r| e.root == r);

    let mut doomed: Vec<PathBuf> = entries

        .iter()

        .filter(|(_, e)| movable(e) && now().saturating_sub(e.last_used) > idle_after)

        .map(|(p, _)| p.clone())

        .collect();

    if to_cap {

        let mut global: Bytes = entries

            .iter()

            .filter(|(p, _)| !doomed.contains(p))

            .map(|(_, e)| e.size)

            .sum();

        let mut per_root: BTreeMap<&Path, Bytes> = BTreeMap::new();

        for (p, e) in &entries {

            if !doomed.contains(p) {

                *per_root.entry(&e.root).or_default() += e.size;

            }

        }

        let mut lru: Vec<(&PathBuf, &Entry)> = entries

            .iter()

            .filter(|(p, e)| movable(e) && !doomed.contains(p))

            .collect();

        lru.sort_by_key(|(_, e)| e.last_used);

        for (p, e) in lru {

            let root_size = per_root.get(e.root.as_path()).copied().unwrap_or_default();

            if global <= cfg.artifacts.global_cap && root_size <= cfg.artifacts.project_cap {

                continue;

            }

            global = global.saturating_sub(e.size);

            per_root.insert(&e.root, root_size.saturating_sub(e.size));

            doomed.push(p.clone());

        }

    }

    let mut out = Gc {

        freed: Bytes::ZERO,

        removed: Vec::new(),

        failed: Vec::new(),

    };

    for path in doomed {

        let e = &entries[&path];

        if !dry_run && let Err(err) = remove(cfg, &path, e) {

            out.failed.push(err);

            continue;

        }

        out.freed += e.size;

        out.removed.push(path);

    }

    if !dry_run {

        let removed = out.removed.clone();

        let _ = store::update("artifacts", |reg: &mut Registry| {

            for p in &removed {

                reg.remove(p);

            }

        });

        nix_gc(cfg, &mut out);

    }

    out

}

fn nix_gc(cfg: &Config, out: &mut Gc) {

    if !cfg.artifacts.nix_gc {

        return;

    }

    let Some(store) = sys::fs(Path::new("/nix/store")) else {

        return;

    };

    let want = cfg.artifacts.nix_min_free.saturating_sub(store.free);

    if want == Bytes::ZERO {

        return;

    }

    let ok = Command::new("nix")

        .args(["store", "gc", "--max", &want.0.to_string()])

        .status()

        .is_ok_and(|s| s.success());

    if ok {

        let after = sys::fs(Path::new("/nix/store")).map_or(store.free, |s| s.free);

        out.freed += after.saturating_sub(store.free);

    } else {

        out.failed.push("nix store gc failed".into());

    }

}