//! 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::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.
pub 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,
}

/// Summaries kept; each is the state's change set, a few hundred kilobytes
/// for a long history.
const SUMMARIES: usize = 2048;

/// 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>,
}

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)) }
    }

    /// 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 mut sorted = heads.to_vec();
        sorted.sort();
        sorted.dedup();
        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)
            }
        });
        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)
    }

    /// What it holds: loaded states, their estimated bytes, and summaries.
    pub fn stats(&self) -> (usize, usize, usize) {
        let (entries, weight) = {
            let guard = self.inner.lock().unwrap();
            (guard.0.len(), guard.0.values().map(|e| e.weight).sum())
        };
        (entries, weight, self.summaries.lock().unwrap().0.len())
    }

    pub fn len(&self) -> usize {
        self.inner.lock().unwrap().0.len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

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);
    }
}