//! On-disk layout of the image store, the metadata format, and garbage collection.

use std::collections::HashSet;
use std::fs;
use std::io::{self, Write};
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime};

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use super::{ImageError, PreparedImage};
use crate::sandbox::idmap::{IdMap, remove_tree};

/// The default cap on the unpacked size of one image.
pub const DEFAULT_MAX_UNPACKED_BYTES: u64 = 20 * 1024 * 1024 * 1024;

/// Staging files and dirs older than this are assumed to be left behind by a crash and are removed.
const STALE_TEMP_AGE: Duration = Duration::from_secs(60 * 60);
/// Name prefixes of staging entries: `tempfile`'s default prefix, and the unpack staging dirs.
const TEMP_PREFIXES: [&str; 2] = [".tmp", ".unpack-"];

/// A cache of pulled and unpacked images rooted at one directory. See the module docs of
/// [`super`] for the layout.
pub struct ImageStore {
    root: PathBuf,
    max_unpacked_bytes: u64,
    /// The id map that unpacking and removal run under. Single unless set by `with_idmap`.
    pub(super) idmap: IdMap,
    /// `forge-runner` itself, which runs `__sandbox-unpack` and `__sandbox-rm` in subuid mode.
    pub(super) helper_exe: PathBuf,
}

/// What a garbage collection pass removed.
#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct GcStats {
    pub images_removed: u64,
    pub blobs_removed: u64,
    /// Staging leftovers (temp files and `.unpack-*` dirs) from crashed pulls.
    pub temp_removed: u64,
}

/// The per-rootfs metadata stored in `meta/<hex>.json`.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub(super) struct ImageMeta {
    pub digest: String,
    pub env: Vec<String>,
    pub working_dir: Option<String>,
    pub user: Option<String>,
    /// Layer digests in order, as bare hex. Garbage collection keeps these blobs.
    pub layers: Vec<String>,
}

impl ImageStore {
    /// A store rooted at `root`, with the default unpacked size cap.
    pub fn new(root: PathBuf) -> Self {
        Self {
            root,
            max_unpacked_bytes: DEFAULT_MAX_UNPACKED_BYTES,
            idmap: IdMap::single(),
            helper_exe: PathBuf::new(),
        }
    }

    /// Sets the cap on the unpacked size of one image, in bytes.
    pub fn with_max_unpacked_bytes(mut self, bytes: u64) -> Self {
        self.max_unpacked_bytes = bytes;
        self
    }

    /// Unpacks under `map`. A subuid map unpacks through `helper_exe` (`__sandbox-unpack`), and its
    /// rootfs dirs are named after the range, so a changed range unpacks again. A single map keeps
    /// the plain `rootfs/<digest>` layout.
    pub fn with_idmap(mut self, map: IdMap, helper_exe: PathBuf) -> Self {
        self.idmap = map;
        self.helper_exe = helper_exe;
        self
    }

    pub(super) fn max_unpacked_bytes(&self) -> u64 {
        self.max_unpacked_bytes
    }

    pub(super) fn root_dir(&self) -> &Path {
        &self.root
    }

    pub(super) fn blob_path(&self, hex: &str) -> PathBuf {
        self.root.join("blobs/sha256").join(hex)
    }

    pub(super) fn rootfs_path(&self, hex: &str) -> PathBuf {
        self.root
            .join("rootfs")
            .join(rootfs_dir_name(&self.idmap, hex))
    }

    /// Removes a file, symlink or directory tree from the store. Trees go through `remove_tree`,
    /// because subuid-owned files can only be removed by the helper.
    pub(super) fn remove_entry(&self, path: &Path) -> io::Result<()> {
        match fs::symlink_metadata(path) {
            Ok(meta) if meta.is_dir() => remove_tree(&self.idmap, path, &self.helper_exe),
            Ok(_) => fs::remove_file(path),
            Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(()),
            Err(err) => Err(err),
        }
    }

    pub(super) fn meta_path(&self, hex: &str) -> PathBuf {
        self.root.join("meta").join(format!("{hex}.json"))
    }

    pub(super) fn lock_path(&self, hex: &str) -> PathBuf {
        self.root.join("locks").join(format!("{hex}.lock"))
    }

    pub(super) fn stamp_path(&self, hex: &str) -> PathBuf {
        self.root.join("last_used").join(hex)
    }

    fn tag_path(&self, reference: &str) -> PathBuf {
        self.root.join("tags").join(encode_tag(reference))
    }

    /// Returns the image for `hex` when both its rootfs and its metadata are present. The rootfs
    /// is only renamed into place once it is complete, so presence means complete.
    pub(super) fn cached(&self, hex: &str) -> Result<Option<PreparedImage>, ImageError> {
        if !self.rootfs_path(hex).is_dir() || !self.meta_path(hex).is_file() {
            return Ok(None);
        }
        let meta: ImageMeta = serde_json::from_slice(&fs::read(self.meta_path(hex))?)?;
        Ok(Some(self.prepared(hex, meta)))
    }

    pub(super) fn prepared(&self, hex: &str, meta: ImageMeta) -> PreparedImage {
        PreparedImage {
            rootfs: self.rootfs_path(hex),
            digest: meta.digest,
            env: meta.env,
            working_dir: meta.working_dir,
            user: meta.user,
        }
    }

    /// The image that `reference` last resolved to, if it is still cached.
    pub(super) fn cached_for_tag(
        &self,
        reference: &str,
    ) -> Result<Option<PreparedImage>, ImageError> {
        let path = self.tag_path(reference);
        if !path.is_file() {
            return Ok(None);
        }
        let digest = fs::read_to_string(&path)?;
        let hex = super::pull::hex_of(digest.trim())?;
        self.cached(hex)
    }

    pub(super) fn write_tag(&self, reference: &str, digest: &str) -> Result<(), ImageError> {
        write_atomic(&self.tag_path(reference), digest.as_bytes())
    }

    pub(super) fn write_meta(&self, hex: &str, meta: &ImageMeta) -> Result<(), ImageError> {
        write_atomic(&self.meta_path(hex), &serde_json::to_vec(meta)?)
    }

    /// Marks the image as used now, for garbage collection.
    pub(super) fn touch(&self, hex: &str) -> io::Result<()> {
        let path = self.stamp_path(hex);
        if let Some(parent) = path.parent() {
            fs::create_dir_all(parent)?;
        }
        fs::write(path, b"")
    }

    /// Deletes unpacked rootfs directories and metadata that have not been used for `keep_days`,
    /// then deletes layer blobs that no kept image refers to and that are older than `keep_days`.
    /// Staging leftovers older than an hour are removed first. Images are locked while they are
    /// checked, so a concurrent pull of the same image waits.
    pub fn gc(&self, keep_days: u32) -> Result<GcStats, ImageError> {
        let max_age = Duration::from_secs(u64::from(keep_days) * 86_400);
        let now = SystemTime::now();
        let mut stats = GcStats::default();
        let mut live_layers: HashSet<String> = HashSet::new();

        for dir in ["rootfs", "blobs/sha256", "meta", "tags"] {
            for path in entries(&self.root.join(dir))? {
                let is_temp = path
                    .file_name()
                    .and_then(|name| name.to_str())
                    .is_some_and(|name| TEMP_PREFIXES.iter().any(|p| name.starts_with(p)));
                if is_temp && modified(&path).is_some_and(|t| older_than(t, now, STALE_TEMP_AGE)) {
                    self.remove_entry(&path)?;
                    stats.temp_removed += 1;
                }
            }
        }

        for meta_file in entries(&self.root.join("meta"))? {
            let Some(hex) = file_stem(&meta_file) else {
                continue;
            };
            let last_used = modified(&self.stamp_path(&hex)).or_else(|| modified(&meta_file));
            let expired = last_used.is_none_or(|t| older_than(t, now, max_age));
            let _lock = lock_file(&self.lock_path(&hex))?;
            if expired {
                self.remove_entry(&self.rootfs_path(&hex))?;
                remove_file_if_exists(&meta_file)?;
                remove_file_if_exists(&self.stamp_path(&hex))?;
                stats.images_removed += 1;
            } else if let Ok(meta) = serde_json::from_slice::<ImageMeta>(&fs::read(&meta_file)?) {
                live_layers.extend(meta.layers);
            }
        }

        for blob in entries(&self.root.join("blobs/sha256"))? {
            let Some(name) = file_stem(&blob) else {
                continue;
            };
            if live_layers.contains(&name) {
                continue;
            }
            if modified(&blob).is_some_and(|t| older_than(t, now, max_age)) {
                remove_file_if_exists(&blob)?;
                stats.blobs_removed += 1;
            }
        }
        Ok(stats)
    }
}

/// Opens (creating if needed) and exclusively locks a lock file. The lock is released when the
/// returned file is dropped.
pub(super) fn lock_file(path: &Path) -> io::Result<fs::File> {
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)?;
    }
    let file = fs::OpenOptions::new()
        .read(true)
        .write(true)
        .create(true)
        .truncate(false)
        .open(path)?;
    file.lock()?;
    Ok(file)
}

/// Writes `bytes` to `path` through a temp file in the same directory, so readers never see a
/// partial file.
pub(super) fn write_atomic(path: &Path, bytes: &[u8]) -> Result<(), ImageError> {
    let dir = path
        .parent()
        .ok_or_else(|| io::Error::other("path has no parent directory"))?;
    fs::create_dir_all(dir)?;
    let mut tmp = tempfile::NamedTempFile::new_in(dir)?;
    tmp.write_all(bytes)?;
    tmp.as_file().sync_all()?;
    tmp.persist(path).map_err(|err| err.error)?;
    Ok(())
}

/// Turns a reference into one file name. Letters, digits, `.` and `-` are kept, and every other
/// byte becomes `%xx`, so the encoding is injective. Long references are replaced by a hash.
fn encode_tag(reference: &str) -> String {
    let mut out = String::with_capacity(reference.len());
    for byte in reference.bytes() {
        if byte.is_ascii_alphanumeric() || byte == b'.' || byte == b'-' {
            out.push(char::from(byte));
        } else {
            out.push_str(&format!("%{byte:02x}"));
        }
    }
    if out.len() > 200 {
        out = hex::encode(Sha256::digest(reference.as_bytes()));
    }
    out
}

/// The entries of `dir`, or none if it does not exist.
fn entries(dir: &Path) -> io::Result<Vec<PathBuf>> {
    match fs::read_dir(dir) {
        Ok(read) => read.map(|entry| entry.map(|e| e.path())).collect(),
        Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(Vec::new()),
        Err(err) => Err(err),
    }
}

fn file_stem(path: &Path) -> Option<String> {
    path.file_stem()
        .and_then(|name| name.to_str())
        .map(str::to_owned)
}

fn modified(path: &Path) -> Option<SystemTime> {
    fs::metadata(path).and_then(|meta| meta.modified()).ok()
}

fn older_than(time: SystemTime, now: SystemTime, max_age: Duration) -> bool {
    now.duration_since(time).is_ok_and(|age| age > max_age)
}

/// The rootfs dir name for `hex`: `<hex>` in single mode, and `<hex>-s<start>-<n>` for a subuid
/// range, so each range has its own unpacked tree.
fn rootfs_dir_name(map: &IdMap, hex: &str) -> String {
    match map.sub_uid {
        Some((start, n)) if map.is_subuid() => format!("{hex}-s{start}-{n}"),
        _ => hex.to_owned(),
    }
}

fn remove_file_if_exists(path: &Path) -> io::Result<()> {
    match fs::remove_file(path) {
        Err(err) if err.kind() == io::ErrorKind::NotFound => Ok(()),
        other => other,
    }
}

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

    #[test]
    fn rootfs_dir_names_follow_the_mode() {
        let hex = "ab".repeat(32);
        assert_eq!(rootfs_dir_name(&IdMap::single(), &hex), hex);
        let subuid = IdMap {
            uid: 1000,
            gid: 1000,
            sub_uid: Some((524_288, 65_536)),
            sub_gid: Some((524_288, 65_536)),
        };
        assert_eq!(
            rootfs_dir_name(&subuid, &hex),
            format!("{hex}-s524288-65536")
        );
        let other_range = IdMap {
            sub_uid: Some((600_000, 65_536)),
            ..subuid
        };
        assert_ne!(
            rootfs_dir_name(&other_range, &hex),
            rootfs_dir_name(&subuid, &hex)
        );
    }

    #[test]
    fn tag_encoding_is_injective_for_similar_names() {
        assert_ne!(encode_tag("a/b:c"), encode_tag("a_b:c"));
        assert!(!encode_tag("docker.io/library/alpine:latest").contains('/'));
        assert!(encode_tag(&"x".repeat(500)).len() <= 200);
    }

    fn hex_n(n: u64) -> String {
        format!("{n:064x}")
    }

    /// Sets the mtime of a file or directory to `age` in the past.
    fn set_age(path: &Path, age: Duration) {
        let when = SystemTime::now() - age;
        fs::File::open(path)
            .expect("open")
            .set_modified(when)
            .expect("set_modified");
    }

    /// Writes an image that was last used just now, with the given layer blobs.
    fn fresh_image(store: &ImageStore, hex: &str, layers: &[String]) {
        fs::create_dir_all(store.rootfs_path(hex)).expect("rootfs");
        fs::write(store.rootfs_path(hex).join("file"), b"x").expect("file");
        store
            .write_meta(
                hex,
                &ImageMeta {
                    digest: format!("sha256:{hex}"),
                    env: Vec::new(),
                    working_dir: None,
                    user: None,
                    layers: layers.to_vec(),
                },
            )
            .expect("meta");
        store.touch(hex).expect("stamp");
    }

    fn blob(store: &ImageStore, hex: &str) {
        fs::create_dir_all(store.blob_path(hex).parent().expect("parent")).expect("blobs");
        fs::write(store.blob_path(hex), b"layer").expect("blob");
    }

    #[test]
    fn gc_removes_stale_images_and_blobs_and_keeps_fresh_ones() {
        let dir = tempfile::tempdir().expect("tempdir");
        let store = ImageStore::new(dir.path().join("oci"));
        let (stale_img, fresh_img) = (hex_n(1), hex_n(2));
        let (stale_layer, shared, orphan_old, orphan_fresh) =
            (hex_n(3), hex_n(4), hex_n(5), hex_n(6));
        let month = Duration::from_secs(30 * 86_400);

        fresh_image(&store, &stale_img, &[stale_layer.clone(), shared.clone()]);
        set_age(&store.stamp_path(&stale_img), month);
        fresh_image(&store, &fresh_img, std::slice::from_ref(&shared));

        blob(&store, &stale_layer);
        blob(&store, &shared);
        blob(&store, &orphan_old);
        blob(&store, &orphan_fresh);
        set_age(&store.blob_path(&stale_layer), month);
        set_age(&store.blob_path(&shared), month);
        set_age(&store.blob_path(&orphan_old), month);

        let stats = store.gc(7).expect("gc");
        assert_eq!(
            stats,
            GcStats {
                images_removed: 1,
                blobs_removed: 2,
                temp_removed: 0,
            }
        );

        assert!(!store.rootfs_path(&stale_img).exists());
        assert!(!store.meta_path(&stale_img).exists());
        assert!(store.rootfs_path(&fresh_img).is_dir());
        assert!(store.meta_path(&fresh_img).is_file());

        // Stale but still referenced by a fresh image: kept.
        assert!(store.blob_path(&shared).is_file());
        // Stale and only referenced by the removed image: removed.
        assert!(!store.blob_path(&stale_layer).exists());
        // Stale orphan: removed. Fresh orphan: kept.
        assert!(!store.blob_path(&orphan_old).exists());
        assert!(store.blob_path(&orphan_fresh).is_file());
    }

    #[test]
    fn gc_removes_stale_staging_leftovers_only() {
        let dir = tempfile::tempdir().expect("tempdir");
        let store = ImageStore::new(dir.path().join("oci"));
        let root = store.root_dir().to_path_buf();
        let two_hours = Duration::from_secs(2 * 3600);

        let stale_dir = root.join("rootfs/.unpack-stale");
        fs::create_dir_all(stale_dir.join("usr")).expect("stale dir");
        fs::write(stale_dir.join("usr/file"), b"partial").expect("partial");
        let fresh_dir = root.join("rootfs/.unpack-fresh");
        fs::create_dir_all(&fresh_dir).expect("fresh dir");

        let blobs = root.join("blobs/sha256");
        fs::create_dir_all(&blobs).expect("blobs");
        let stale_blob_tmp = blobs.join(".tmpstale");
        fs::write(&stale_blob_tmp, b"partial").expect("tmp");
        let fresh_blob_tmp = blobs.join(".tmpfresh");
        fs::write(&fresh_blob_tmp, b"partial").expect("tmp");

        let meta = root.join("meta");
        fs::create_dir_all(&meta).expect("meta");
        let stale_meta_tmp = meta.join(".tmpstale");
        fs::write(&stale_meta_tmp, b"{}").expect("meta tmp");

        // A non-staging name in the same dir is never touched, however old it is.
        let plain = root.join("rootfs/keep-me");
        fs::create_dir_all(&plain).expect("plain");
        set_age(&plain, two_hours);

        set_age(&stale_dir.join("usr"), two_hours);
        set_age(&stale_dir, two_hours);
        set_age(&stale_blob_tmp, two_hours);
        set_age(&stale_meta_tmp, two_hours);

        let stats = store.gc(7).expect("gc");
        assert_eq!(stats.temp_removed, 3);
        assert_eq!(stats.blobs_removed, 0);
        assert!(!stale_dir.exists());
        assert!(fresh_dir.is_dir());
        assert!(!stale_blob_tmp.exists());
        assert!(fresh_blob_tmp.is_file());
        assert!(!stale_meta_tmp.exists());
        assert!(plain.is_dir());
    }
}