//! Loaded states cached by their heads (hub spec §16.3). A state is a pure

//! function of its heads, so entries never go stale; the cache is only

//! bounded. Shared across threads because `oxy_core::state::State` is

//! `Send + Sync`.

use oxy_core::hash::ChangeHash;

use oxy_core::repo::Repo;

use oxy_core::state::{Snapshot, State};

use std::collections::HashMap;

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

use std::sync::{Arc, Mutex};

/// A state and its materialised snapshot.

pub struct Loaded {

    pub heads: Vec<ChangeHash>,

    pub state: State,

    pub snapshot: Snapshot,

}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]

struct Key {

    repo: [u8; 16],

    heads: Vec<ChangeHash>,

}

struct Entry {

    loaded: Arc<Loaded>,

    last_used: u64,

    /// Its estimated size in bytes (`weight`).

    weight: usize,

}

/// A loaded state's estimated size: its rendered text three times over (the

/// text, and the line graphs it was rendered from), plus a kilobyte per

/// change. Rough, but proportional: a repository of large files weighs

/// more than one of small ones.

fn weight(l: &Loaded) -> usize {

    let text: usize = l

        .snapshot

        .files

        .values()

        .map(|f| match &f.content {

            oxy_core::state::FileContent::Text { bytes, .. } => bytes.len(),

            _ => 64,

        })

        .sum();

    text.saturating_mul(3).saturating_add(l.state.len().saturating_mul(1024))

}

/// What the state cache may hold, in bytes: OXYHUB_STATE_CACHE_MB, else

/// 3 GiB. A count alone let a few large repositories' states grow the hub

/// to 17 GB.

fn budget() -> usize {

    std::env::var("OXYHUB_STATE_CACHE_MB").ok().and_then(|v| v.parse::<usize>().ok()).unwrap_or(3 * 1024).saturating_mul(1 << 20)

}

/// What the bookmark lists and a change's "in bookmarks" need of a state:

/// small, so many are kept, and computed without filling the state cache.

/// A repository with more bookmarks than the state cache holds otherwise

/// cycled through it on every such page, so nothing stayed cached

/// (metalgear's 128 bookmarks: 7.7 s per change page, and every page of

/// every repository slower after it).

pub struct Summary {

    pub changes: std::collections::HashSet<ChangeHash>,

    pub latest_ms: u64,

    pub conflicted: bool,

}

impl Summary {

    pub fn brief(&self) -> Brief {

        Brief { changes: self.changes.len(), latest_ms: self.latest_ms, conflicted: self.conflicted }

    }

}

/// What a bookmark list and the overview show of a state (#145): how many

/// changes, the latest, and whether it is conflicted. A pure function of

/// the heads, so it is kept on disk as well and a restart costs nothing:

/// metalgear's 162 bookmarks took ~31 s to summarise again after one.

#[derive(Clone, Copy, Debug, PartialEq, Eq)]

pub struct Brief {

    pub changes: usize,

    pub latest_ms: u64,

    pub conflicted: bool,

}

/// Summaries kept; each is the state's change set, a few hundred kilobytes

/// for a long history.

const SUMMARIES: usize = 2048;

/// Most summaries computed at once when a page needs many.

const PARALLEL: usize = 4;

/// Summaries by heads, each with when it was last used, and the clock.

type Summaries = (HashMap<Key, (Arc<Summary>, u64)>, u64);

/// A bounded least-recently-used cache of loaded states.

pub struct StateCache {

    capacity: usize,

    /// Bytes of estimated state it may hold (`budget`).

    budget: usize,

    inner: Mutex<(HashMap<Key, Entry>, u64)>,

    summaries: Mutex<Summaries>,

    /// Where briefs are kept on disk, when anywhere.

    briefs: Option<PathBuf>,

}

impl StateCache {

    pub fn new(capacity: usize) -> StateCache {

        StateCache::with_budget(capacity, budget())

    }

    pub fn with_budget(capacity: usize, budget: usize) -> StateCache {

        StateCache {

            capacity: capacity.max(1),

            budget,

            inner: Mutex::new((HashMap::new(), 0)),

            summaries: Mutex::new((HashMap::new(), 0)),

            briefs: None,

        }

    }

    /// Keep briefs on disk under `dir`, one small file per state.

    pub fn with_briefs(mut self, dir: PathBuf) -> StateCache {

        self.briefs = Some(dir);

        self

    }

    fn brief_path(&self, key: &Key) -> Option<PathBuf> {

        let dir = self.briefs.as_ref()?;

        let mut h = blake3::Hasher::new();

        for c in &key.heads {

            h.update(&c.0);

        }

        Some(dir.join(hex::encode(key.repo)).join(h.finalize().to_hex().as_str()))

    }

    fn known_brief(&self, key: &Key) -> Option<Brief> {

        if let Some((s, _)) = self.summaries.lock().unwrap().0.get(key) {

            return Some(s.brief());

        }

        read_brief(&self.brief_path(key)?)

    }

    /// The brief of the state with these heads: from memory, from disk, or

    /// worked out and kept.

    pub fn brief(&self, repo: &Repo, heads: &[ChangeHash]) -> Result<Brief, oxy_core::repo::RepoError> {

        Ok(self.briefs(repo, &[heads.to_vec()])?[0])

    }

    /// The briefs of many states, those not yet known worked out a few at

    /// a time in parallel.

    pub fn briefs(&self, repo: &Repo, all: &[Vec<ChangeHash>]) -> Result<Vec<Brief>, oxy_core::repo::RepoError> {

        let id = repo.id()?;

        let mut out: Vec<Option<Brief>> = all.iter().map(|h| self.known_brief(&Key { repo: id, heads: sorted(h) })).collect();

        let missing: Vec<usize> = (0..all.len()).filter(|&i| out[i].is_none()).collect();

        if missing.is_empty() {

            return Ok(out.into_iter().map(|b| b.expect("known")).collect());

        }

        let next = Mutex::new(missing.iter());

        let found: Mutex<Vec<(usize, Result<Brief, oxy_core::repo::RepoError>)>> = Mutex::new(Vec::new());

        let threads = PARALLEL.min(missing.len()).min(std::thread::available_parallelism().map(|n| n.get()).unwrap_or(1));

        // A `Repo` stays on its thread, so each worker opens its own.

        let root = repo.root().to_path_buf();

        std::thread::scope(|s| {

            for _ in 0..threads {

                s.spawn(|| {

                    let mine = Repo::open(&root);

                    loop {

                        let Some(&i) = next.lock().unwrap().next() else { break };

                        let b = match &mine {

                            Ok(r) => self.summary(r, &all[i]).map(|s| s.brief()),

                            Err(e) => Err(oxy_core::repo::RepoError::Msg(format!("open {}: {e}", root.display()))),

                        };

                        found.lock().unwrap().push((i, b));

                    }

                });

            }

        });

        for (i, b) in found.into_inner().unwrap() {

            out[i] = Some(b?);

        }

        Ok(out.into_iter().map(|b| b.expect("worked out")).collect())

    }

    /// The summary of the state with these heads: from the state cache when

    /// it holds the state, else loaded and dropped, the state cache left as

    /// it was.

    pub fn summary(&self, repo: &Repo, heads: &[ChangeHash]) -> Result<Arc<Summary>, oxy_core::repo::RepoError> {

        let sorted = sorted(heads);

        let key = Key { repo: repo.id()?, heads: sorted.clone() };

        {

            let mut guard = self.summaries.lock().unwrap();

            let (map, tick) = &mut *guard;

            *tick += 1;

            if let Some(e) = map.get_mut(&key) {

                e.1 = *tick;

                return Ok(e.0.clone());

            }

        }

        let held = self.inner.lock().unwrap().0.get(&key).map(|e| e.loaded.clone());

        let summary = Arc::new(match held {

            Some(l) => summarise(&l.state, &l.snapshot),

            None => {

                let state = repo.load_state(&sorted)?;

                let snapshot = state.materialize();

                summarise(&state, &snapshot)

            }

        });

        if let Some(path) = self.brief_path(&key) {

            write_brief(&path, summary.brief());

        }

        let mut guard = self.summaries.lock().unwrap();

        let (map, tick) = &mut *guard;

        if map.len() >= SUMMARIES {

            if let Some(oldest) = map.iter().min_by_key(|(_, e)| e.1).map(|(k, _)| k.clone()) {

                map.remove(&oldest);

            }

        }

        map.insert(key, (summary.clone(), *tick));

        Ok(summary)

    }

    /// The state with these heads, loaded once per hub process.

    pub fn load(&self, repo: &Repo, heads: &[ChangeHash]) -> Result<Arc<Loaded>, oxy_core::repo::RepoError> {

        let mut sorted = heads.to_vec();

        sorted.sort();

        sorted.dedup();

        let key = Key { repo: repo.id()?, heads: sorted.clone() };

        {

            let mut guard = self.inner.lock().unwrap();

            let (map, tick) = &mut *guard;

            *tick += 1;

            if let Some(e) = map.get_mut(&key) {

                e.last_used = *tick;

                return Ok(e.loaded.clone());

            }

        }

        let state = repo.load_state(&sorted)?;

        let snapshot = state.materialize();

        let loaded = Arc::new(Loaded { heads: sorted, state, snapshot });

        let w = weight(&loaded);

        let mut guard = self.inner.lock().unwrap();

        let (map, tick) = &mut *guard;

        // The oldest go until it fits by count and by size; the new one

        // stays even when it alone is over the budget.

        let budget = self.budget;

        let mut total: usize = map.values().map(|e| e.weight).sum::<usize>() + w;

        while !map.is_empty() && (map.len() >= self.capacity || total > budget) {

            let Some(oldest) = map.iter().min_by_key(|(_, e)| e.last_used).map(|(k, _)| k.clone()) else {

                break;

            };

            if let Some(e) = map.remove(&oldest) {

                total -= e.weight;

            }

        }

        map.insert(key, Entry { loaded: loaded.clone(), last_used: *tick, weight: w });

        Ok(loaded)

    }

    pub fn len(&self) -> usize {

        self.inner.lock().unwrap().0.len()

    }

    pub fn is_empty(&self) -> bool {

        self.len() == 0

    }

}

fn sorted(heads: &[ChangeHash]) -> Vec<ChangeHash> {

    let mut v = heads.to_vec();

    v.sort();

    v.dedup();

    v

}

/// A brief on disk: `changes latest_ms conflicted`.

fn read_brief(path: &Path) -> Option<Brief> {

    let text = std::fs::read_to_string(path).ok()?;

    let mut f = text.split_whitespace();

    Some(Brief { changes: f.next()?.parse().ok()?, latest_ms: f.next()?.parse().ok()?, conflicted: f.next()? == "1" })

}

/// Kept when it can be; a brief that is not kept is only worked out again.

fn write_brief(path: &Path, b: Brief) {

    let Some(dir) = path.parent() else { return };

    if std::fs::create_dir_all(dir).is_err() {

        return;

    }

    let tmp = path.with_extension(format!("tmp{}", std::process::id()));

    if std::fs::write(&tmp, format!("{} {} {}", b.changes, b.latest_ms, u8::from(b.conflicted))).is_ok() {

        std::fs::rename(&tmp, path).ok();

    }

}

fn summarise(state: &State, snapshot: &Snapshot) -> Summary {

    Summary {

        changes: state.changes().map(|(h, _)| *h).collect(),

        latest_ms: state.changes().map(|(_, m)| m.hlc.wall_ms).max().unwrap_or(0),

        conflicted: snapshot.is_conflicted(),

    }

}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn caches_by_heads_and_evicts_the_oldest() {

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

        let repo = Repo::init_bare(dir.path()).unwrap();

        let cache = StateCache::new(2);

        let a = cache.load(&repo, &[]).unwrap();

        let b = cache.load(&repo, &[]).unwrap();

        assert!(Arc::ptr_eq(&a, &b), "same heads, same entry");

        assert_eq!(cache.len(), 1);

        let h1 = oxy_core::hash::ObjectRef([1; 32]);

        let h2 = oxy_core::hash::ObjectRef([2; 32]);

        // Unknown heads fail to load and are not cached.

        assert!(cache.load(&repo, &[h1]).is_err());

        assert_eq!(cache.len(), 1);

        let _ = (h2, dir);

    }

    #[test]

    fn summaries_leave_the_state_cache_alone() {

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

        let repo = Repo::init_bare(dir.path()).unwrap();

        let cache = StateCache::new(2);

        let a = cache.summary(&repo, &[]).unwrap();

        assert_eq!(cache.len(), 0, "a summary does not fill the state cache");

        assert!(Arc::ptr_eq(&a, &cache.summary(&repo, &[]).unwrap()), "kept");

        assert!(a.changes.is_empty() && !a.conflicted);

        assert!(cache.summary(&repo, &[oxy_core::hash::ObjectRef([3; 32])]).is_err());

    }

    #[test]

    fn a_size_budget_evicts_the_oldest_and_keeps_the_newest() {

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

        let repo = Repo::init(dir.path()).unwrap();

        std::fs::write(dir.path().join("a.txt"), "a\n".repeat(1000)).unwrap();

        let one = repo.commit("one", &[]).unwrap();

        std::fs::write(dir.path().join("b.txt"), "b\n".repeat(1000)).unwrap();

        let two = repo.commit("two", &[]).unwrap();

        // Room for nothing: each load stays, the one before goes.

        let cache = StateCache::with_budget(64, 0);

        cache.load(&repo, &[one]).unwrap();

        assert_eq!(cache.len(), 1, "kept though over the budget");

        cache.load(&repo, &[two]).unwrap();

        assert_eq!(cache.len(), 1, "the older one went");

        // Room for plenty: both stay.

        let roomy = StateCache::with_budget(64, 1 << 30);

        roomy.load(&repo, &[one]).unwrap();

        roomy.load(&repo, &[two]).unwrap();

        assert_eq!(roomy.len(), 2);

    }

    #[test]

    fn briefs_are_worked_out_in_parallel_and_survive_a_restart() {

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

        let repo = Repo::init(dir.path()).unwrap();

        std::fs::write(dir.path().join("a.txt"), "a\n").unwrap();

        let one = repo.commit("one", &[]).unwrap();

        std::fs::write(dir.path().join("b.txt"), "b\n").unwrap();

        let two = repo.commit("two", &[]).unwrap();

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

        let cache = StateCache::new(8).with_briefs(kept.path().to_path_buf());

        let got = cache.briefs(&repo, &[vec![one], vec![one, two], vec![]]).unwrap();

        assert_eq!(got.iter().map(|b| b.changes).collect::<Vec<_>>(), vec![1, 2, 0]);

        assert!(got[1].latest_ms >= got[0].latest_ms && !got[1].conflicted);

        // A new cache (a restarted hub) reads them back without loading a

        // state: a repository that no longer has the changes still answers.

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

        let empty = Repo::init(gone.path()).unwrap();

        let empty_id = empty.id().unwrap();

        for f in std::fs::read_dir(kept.path()).unwrap() {

            let f = f.unwrap().path();

            std::fs::rename(&f, kept.path().join(hex::encode(empty_id))).unwrap();

        }

        let restarted = StateCache::new(8).with_briefs(kept.path().to_path_buf());

        assert_eq!(restarted.brief(&empty, &[two, one]).unwrap(), got[1], "from disk, whatever the heads' order");

        assert!(StateCache::new(8).brief(&empty, &[one, two]).is_err(), "without the disk it would have to load");

        // Unknown heads are an error, not a guess.

        assert!(cache.briefs(&repo, &[vec![one], vec![oxy_core::hash::ObjectRef([9; 32])]]).is_err());

        // A bare repository, as the hub keeps them, works across threads too.

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

        let bare = Repo::init_bare(bare_dir.path()).unwrap();

        assert_eq!(cache.briefs(&bare, &[vec![], vec![]]).unwrap().len(), 2);

    }

}