lean_ctx/core/

context_ledger.rs

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use super::context_field::{
    ContextItemId, ContextKind, ContextState, Provenance, ViewCosts, ViewKind,
};

const DEFAULT_CONTEXT_WINDOW: usize = 128_000;

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextLedger {
    pub window_size: usize,
    pub entries: Vec<LedgerEntry>,
    pub total_tokens_sent: usize,
    pub total_tokens_saved: usize,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LedgerEntry {
    pub path: String,
    pub mode: String,
    pub original_tokens: usize,
    pub sent_tokens: usize,
    pub timestamp: i64,
    #[serde(default)]
    pub id: Option<ContextItemId>,
    #[serde(default)]
    pub kind: Option<ContextKind>,
    #[serde(default)]
    pub source_hash: Option<String>,
    #[serde(default)]
    pub state: Option<ContextState>,
    #[serde(default)]
    pub phi: Option<f64>,
    #[serde(default)]
    pub view_costs: Option<ViewCosts>,
    #[serde(default)]
    pub active_view: Option<ViewKind>,
    #[serde(default)]
    pub provenance: Option<Provenance>,
}

#[derive(Debug, Clone)]
pub struct ContextPressure {
    pub utilization: f64,
    pub remaining_tokens: usize,
    pub entries_count: usize,
    pub recommendation: PressureAction,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PressureAction {
    NoAction,
    SuggestCompression,
    ForceCompression,
    EvictLeastRelevant,
}

impl ContextLedger {
    pub fn new() -> Self {
        Self {
            window_size: DEFAULT_CONTEXT_WINDOW,
            entries: Vec::new(),
            total_tokens_sent: 0,
            total_tokens_saved: 0,
        }
    }

    pub fn with_window_size(size: usize) -> Self {
        Self {
            window_size: size,
            entries: Vec::new(),
            total_tokens_sent: 0,
            total_tokens_saved: 0,
        }
    }

    pub fn record(&mut self, path: &str, mode: &str, original_tokens: usize, sent_tokens: usize) {
        let item_id = ContextItemId::from_file(path);
        if let Some(existing) = self.entries.iter_mut().find(|e| e.path == path) {
            self.total_tokens_sent -= existing.sent_tokens;
            self.total_tokens_saved -= existing
                .original_tokens
                .saturating_sub(existing.sent_tokens);
            existing.mode = mode.to_string();
            existing.original_tokens = original_tokens;
            existing.sent_tokens = sent_tokens;
            existing.timestamp = chrono::Utc::now().timestamp();
            existing.active_view = Some(ViewKind::parse(mode));
            if existing.id.is_none() {
                existing.id = Some(item_id);
            }
            if existing.state.is_none() || existing.state == Some(ContextState::Candidate) {
                existing.state = Some(ContextState::Included);
            }
        } else {
            self.entries.push(LedgerEntry {
                path: path.to_string(),
                mode: mode.to_string(),
                original_tokens,
                sent_tokens,
                timestamp: chrono::Utc::now().timestamp(),
                id: Some(item_id),
                kind: Some(ContextKind::File),
                source_hash: None,
                state: Some(ContextState::Included),
                phi: None,
                view_costs: Some(ViewCosts::from_full_tokens(original_tokens)),
                active_view: Some(ViewKind::parse(mode)),
                provenance: None,
            });
        }
        self.total_tokens_sent += sent_tokens;
        self.total_tokens_saved += original_tokens.saturating_sub(sent_tokens);
    }

    /// Record with full CFT metadata including source hash and provenance.
    pub fn upsert(
        &mut self,
        path: &str,
        mode: &str,
        original_tokens: usize,
        sent_tokens: usize,
        source_hash: Option<&str>,
        kind: ContextKind,
        provenance: Option<Provenance>,
    ) {
        self.record(path, mode, original_tokens, sent_tokens);
        if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
            entry.kind = Some(kind);
            if let Some(h) = source_hash {
                if entry.source_hash.as_deref() != Some(h) {
                    if entry.source_hash.is_some() {
                        entry.state = Some(ContextState::Stale);
                    }
                    entry.source_hash = Some(h.to_string());
                }
            }
            if let Some(prov) = provenance {
                entry.provenance = Some(prov);
            }
        }
    }

    /// Update the Phi score for an entry.
    pub fn update_phi(&mut self, path: &str, phi: f64) {
        if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
            entry.phi = Some(phi);
        }
    }

    /// Set the state for an entry.
    pub fn set_state(&mut self, path: &str, state: ContextState) {
        if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
            entry.state = Some(state);
        }
    }

    /// Find an entry by its ContextItemId.
    pub fn find_by_id(&self, id: &ContextItemId) -> Option<&LedgerEntry> {
        self.entries.iter().find(|e| e.id.as_ref() == Some(id))
    }

    /// Get all entries with a specific state.
    pub fn items_by_state(&self, state: ContextState) -> Vec<&LedgerEntry> {
        self.entries
            .iter()
            .filter(|e| e.state == Some(state))
            .collect()
    }

    /// Eviction candidates ordered by Phi (lowest first), falling back to
    /// timestamp for entries without Phi scores.
    pub fn eviction_candidates_by_phi(&self, keep_count: usize) -> Vec<String> {
        if self.entries.len() <= keep_count {
            return Vec::new();
        }
        let mut sorted = self.entries.clone();
        sorted.sort_by(|a, b| {
            let a_phi = a.phi.unwrap_or(0.0);
            let b_phi = b.phi.unwrap_or(0.0);
            a_phi
                .partial_cmp(&b_phi)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then_with(|| a.timestamp.cmp(&b.timestamp))
        });
        sorted
            .iter()
            .filter(|e| e.state != Some(ContextState::Pinned))
            .take(self.entries.len() - keep_count)
            .map(|e| e.path.clone())
            .collect()
    }

    /// Mark entries as stale if their source hash has changed.
    pub fn mark_stale_by_hash(&mut self, path: &str, new_hash: &str) {
        if let Some(entry) = self.entries.iter_mut().find(|e| e.path == path) {
            if let Some(ref old_hash) = entry.source_hash {
                if old_hash != new_hash {
                    entry.state = Some(ContextState::Stale);
                    entry.source_hash = Some(new_hash.to_string());
                }
            }
        }
    }

    pub fn pressure(&self) -> ContextPressure {
        let utilization = self.total_tokens_sent as f64 / self.window_size as f64;

        let pinned_count = self
            .entries
            .iter()
            .filter(|e| e.state == Some(ContextState::Pinned))
            .count();
        let stale_count = self
            .entries
            .iter()
            .filter(|e| e.state == Some(ContextState::Stale))
            .count();
        let pinned_pressure = pinned_count as f64 * 0.02;
        let stale_penalty = stale_count as f64 * 0.01;
        let effective_utilization = (utilization + pinned_pressure + stale_penalty).min(1.0);

        let effective_used = (effective_utilization * self.window_size as f64).round() as usize;
        let remaining = self.window_size.saturating_sub(effective_used);

        let recommendation = if effective_utilization > 0.9 {
            PressureAction::EvictLeastRelevant
        } else if effective_utilization > 0.75 {
            PressureAction::ForceCompression
        } else if effective_utilization > 0.5 {
            PressureAction::SuggestCompression
        } else {
            PressureAction::NoAction
        };

        ContextPressure {
            utilization: effective_utilization,
            remaining_tokens: remaining,
            entries_count: self.entries.len(),
            recommendation,
        }
    }

    pub fn compression_ratio(&self) -> f64 {
        let total_original: usize = self.entries.iter().map(|e| e.original_tokens).sum();
        if total_original == 0 {
            return 1.0;
        }
        self.total_tokens_sent as f64 / total_original as f64
    }

    pub fn files_by_token_cost(&self) -> Vec<(String, usize)> {
        let mut costs: Vec<(String, usize)> = self
            .entries
            .iter()
            .map(|e| (e.path.clone(), e.sent_tokens))
            .collect();
        costs.sort_by_key(|b| std::cmp::Reverse(b.1));
        costs
    }

    pub fn mode_distribution(&self) -> HashMap<String, usize> {
        let mut dist: HashMap<String, usize> = HashMap::new();
        for entry in &self.entries {
            *dist.entry(entry.mode.clone()).or_insert(0) += 1;
        }
        dist
    }

    pub fn eviction_candidates(&self, keep_count: usize) -> Vec<String> {
        if self.entries.len() <= keep_count {
            return Vec::new();
        }
        let mut sorted = self.entries.clone();
        sorted.sort_by_key(|e| e.timestamp);
        sorted
            .iter()
            .take(self.entries.len() - keep_count)
            .map(|e| e.path.clone())
            .collect()
    }

    pub fn remove(&mut self, path: &str) {
        if let Some(idx) = self.entries.iter().position(|e| e.path == path) {
            let entry = &self.entries[idx];
            self.total_tokens_sent -= entry.sent_tokens;
            self.total_tokens_saved -= entry.original_tokens.saturating_sub(entry.sent_tokens);
            self.entries.remove(idx);
        }
    }

    pub fn save(&self) {
        if let Ok(dir) = crate::core::data_dir::lean_ctx_data_dir() {
            let path = dir.join("context_ledger.json");
            if let Ok(json) = serde_json::to_string(self) {
                let _ = std::fs::write(path, json);
            }
        }
    }

    pub fn load() -> Self {
        crate::core::data_dir::lean_ctx_data_dir()
            .ok()
            .map(|d| d.join("context_ledger.json"))
            .and_then(|p| std::fs::read_to_string(p).ok())
            .and_then(|s| serde_json::from_str(&s).ok())
            .unwrap_or_default()
    }

    pub fn format_summary(&self) -> String {
        let pressure = self.pressure();
        format!(
            "CTX: {}/{} tokens ({:.0}%), {} files, ratio {:.2}, action: {:?}",
            self.total_tokens_sent,
            self.window_size,
            pressure.utilization * 100.0,
            self.entries.len(),
            self.compression_ratio(),
            pressure.recommendation,
        )
    }
}

#[derive(Debug, Clone)]
pub struct ReinjectionAction {
    pub path: String,
    pub current_mode: String,
    pub new_mode: String,
    pub tokens_freed: usize,
}

#[derive(Debug, Clone)]
pub struct ReinjectionPlan {
    pub actions: Vec<ReinjectionAction>,
    pub total_tokens_freed: usize,
    pub new_utilization: f64,
}

impl ContextLedger {
    pub fn reinjection_plan(
        &self,
        intent: &super::intent_engine::StructuredIntent,
        target_utilization: f64,
    ) -> ReinjectionPlan {
        let current_util = self.total_tokens_sent as f64 / self.window_size as f64;
        if current_util <= target_utilization {
            return ReinjectionPlan {
                actions: Vec::new(),
                total_tokens_freed: 0,
                new_utilization: current_util,
            };
        }

        let tokens_to_free =
            self.total_tokens_sent - (self.window_size as f64 * target_utilization) as usize;

        let target_set: std::collections::HashSet<&str> = intent
            .targets
            .iter()
            .map(std::string::String::as_str)
            .collect();

        let mut candidates: Vec<(usize, &LedgerEntry)> = self
            .entries
            .iter()
            .enumerate()
            .filter(|(_, e)| !target_set.iter().any(|t| e.path.contains(t)))
            .collect();

        candidates.sort_by(|a, b| {
            let a_phi = a.1.phi.unwrap_or(0.0);
            let b_phi = b.1.phi.unwrap_or(0.0);
            a_phi
                .partial_cmp(&b_phi)
                .unwrap_or_else(|| a.1.timestamp.cmp(&b.1.timestamp))
        });

        let mut actions = Vec::new();
        let mut freed = 0usize;

        for (_, entry) in &candidates {
            if freed >= tokens_to_free {
                break;
            }
            if let Some((new_mode, new_tokens)) = downgrade_mode(&entry.mode, entry.sent_tokens) {
                let saving = entry.sent_tokens.saturating_sub(new_tokens);
                if saving > 0 {
                    actions.push(ReinjectionAction {
                        path: entry.path.clone(),
                        current_mode: entry.mode.clone(),
                        new_mode,
                        tokens_freed: saving,
                    });
                    freed += saving;
                }
            }
        }

        let new_sent = self.total_tokens_sent.saturating_sub(freed);
        let new_utilization = new_sent as f64 / self.window_size as f64;

        ReinjectionPlan {
            actions,
            total_tokens_freed: freed,
            new_utilization,
        }
    }
}

fn downgrade_mode(current_mode: &str, current_tokens: usize) -> Option<(String, usize)> {
    match current_mode {
        "full" => Some(("signatures".to_string(), current_tokens / 5)),
        "aggressive" => Some(("signatures".to_string(), current_tokens / 3)),
        "signatures" => Some(("map".to_string(), current_tokens / 2)),
        "map" => Some(("reference".to_string(), current_tokens / 4)),
        _ => None,
    }
}

impl Default for ContextLedger {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn new_ledger_is_empty() {
        let ledger = ContextLedger::new();
        assert_eq!(ledger.total_tokens_sent, 0);
        assert_eq!(ledger.entries.len(), 0);
        assert_eq!(ledger.pressure().recommendation, PressureAction::NoAction);
    }

    #[test]
    fn record_tracks_tokens() {
        let mut ledger = ContextLedger::with_window_size(10000);
        ledger.record("src/main.rs", "full", 500, 500);
        ledger.record("src/lib.rs", "signatures", 1000, 200);
        assert_eq!(ledger.total_tokens_sent, 700);
        assert_eq!(ledger.total_tokens_saved, 800);
        assert_eq!(ledger.entries.len(), 2);
    }

    #[test]
    fn record_updates_existing_entry() {
        let mut ledger = ContextLedger::with_window_size(10000);
        ledger.record("src/main.rs", "full", 500, 500);
        ledger.record("src/main.rs", "signatures", 500, 100);
        assert_eq!(ledger.entries.len(), 1);
        assert_eq!(ledger.total_tokens_sent, 100);
        assert_eq!(ledger.total_tokens_saved, 400);
    }

    #[test]
    fn pressure_escalates() {
        let mut ledger = ContextLedger::with_window_size(1000);
        ledger.record("a.rs", "full", 600, 600);
        assert_eq!(
            ledger.pressure().recommendation,
            PressureAction::SuggestCompression
        );
        ledger.record("b.rs", "full", 200, 200);
        assert_eq!(
            ledger.pressure().recommendation,
            PressureAction::ForceCompression
        );
        ledger.record("c.rs", "full", 150, 150);
        assert_eq!(
            ledger.pressure().recommendation,
            PressureAction::EvictLeastRelevant
        );
    }

    #[test]
    fn compression_ratio_accurate() {
        let mut ledger = ContextLedger::with_window_size(10000);
        ledger.record("a.rs", "full", 1000, 1000);
        ledger.record("b.rs", "signatures", 1000, 200);
        let ratio = ledger.compression_ratio();
        assert!((ratio - 0.6).abs() < 0.01);
    }

    #[test]
    fn eviction_returns_oldest() {
        let mut ledger = ContextLedger::with_window_size(10000);
        ledger.record("old.rs", "full", 100, 100);
        std::thread::sleep(std::time::Duration::from_millis(10));
        ledger.record("new.rs", "full", 100, 100);
        let candidates = ledger.eviction_candidates(1);
        assert_eq!(candidates, vec!["old.rs"]);
    }

    #[test]
    fn remove_updates_totals() {
        let mut ledger = ContextLedger::with_window_size(10000);
        ledger.record("a.rs", "full", 500, 500);
        ledger.record("b.rs", "full", 300, 300);
        ledger.remove("a.rs");
        assert_eq!(ledger.total_tokens_sent, 300);
        assert_eq!(ledger.entries.len(), 1);
    }

    #[test]
    fn mode_distribution_counts() {
        let mut ledger = ContextLedger::new();
        ledger.record("a.rs", "full", 100, 100);
        ledger.record("b.rs", "signatures", 100, 50);
        ledger.record("c.rs", "full", 100, 100);
        let dist = ledger.mode_distribution();
        assert_eq!(dist.get("full"), Some(&2));
        assert_eq!(dist.get("signatures"), Some(&1));
    }

    #[test]
    fn format_summary_includes_key_info() {
        let mut ledger = ContextLedger::with_window_size(10000);
        ledger.record("a.rs", "full", 500, 500);
        let summary = ledger.format_summary();
        assert!(summary.contains("500/10000"));
        assert!(summary.contains("1 files"));
    }

    #[test]
    fn reinjection_no_action_when_low_pressure() {
        use crate::core::intent_engine::StructuredIntent;

        let mut ledger = ContextLedger::with_window_size(10000);
        ledger.record("a.rs", "full", 100, 100);
        let intent = StructuredIntent::from_query("fix bug in a.rs");
        let plan = ledger.reinjection_plan(&intent, 0.7);
        assert!(plan.actions.is_empty());
        assert_eq!(plan.total_tokens_freed, 0);
    }

    #[test]
    fn reinjection_downgrades_non_target_files() {
        use crate::core::intent_engine::StructuredIntent;

        let mut ledger = ContextLedger::with_window_size(1000);
        ledger.record("src/target.rs", "full", 400, 400);
        std::thread::sleep(std::time::Duration::from_millis(10));
        ledger.record("src/other.rs", "full", 400, 400);
        std::thread::sleep(std::time::Duration::from_millis(10));
        ledger.record("src/utils.rs", "full", 200, 200);

        let intent = StructuredIntent::from_query("fix bug in target.rs");
        let plan = ledger.reinjection_plan(&intent, 0.5);

        assert!(!plan.actions.is_empty());
        assert!(
            plan.actions.iter().all(|a| !a.path.contains("target")),
            "should not downgrade target file"
        );
        assert!(plan.total_tokens_freed > 0);
    }

    #[test]
    fn reinjection_preserves_targets() {
        use crate::core::intent_engine::StructuredIntent;

        let mut ledger = ContextLedger::with_window_size(1000);
        ledger.record("src/auth.rs", "full", 900, 900);
        let intent = StructuredIntent::from_query("fix bug in auth.rs");
        let plan = ledger.reinjection_plan(&intent, 0.5);
        assert!(
            plan.actions.is_empty(),
            "should not downgrade target files even under pressure"
        );
    }

    #[test]
    fn downgrade_mode_chain() {
        assert_eq!(
            downgrade_mode("full", 1000),
            Some(("signatures".to_string(), 200))
        );
        assert_eq!(
            downgrade_mode("signatures", 200),
            Some(("map".to_string(), 100))
        );
        assert_eq!(
            downgrade_mode("map", 100),
            Some(("reference".to_string(), 25))
        );
        assert_eq!(downgrade_mode("reference", 25), None);
    }

    #[test]
    fn record_assigns_item_id() {
        let mut ledger = ContextLedger::new();
        ledger.record("src/main.rs", "full", 500, 500);
        let entry = &ledger.entries[0];
        assert!(entry.id.is_some());
        assert_eq!(entry.id.as_ref().unwrap().as_str(), "file:src/main.rs");
    }

    #[test]
    fn record_sets_state_to_included() {
        let mut ledger = ContextLedger::new();
        ledger.record("src/main.rs", "full", 500, 500);
        assert_eq!(
            ledger.entries[0].state,
            Some(crate::core::context_field::ContextState::Included)
        );
    }

    #[test]
    fn record_generates_view_costs() {
        let mut ledger = ContextLedger::new();
        ledger.record("src/main.rs", "full", 5000, 5000);
        let vc = ledger.entries[0].view_costs.as_ref().unwrap();
        assert_eq!(vc.get(&crate::core::context_field::ViewKind::Full), 5000);
        assert_eq!(
            vc.get(&crate::core::context_field::ViewKind::Signatures),
            1000
        );
    }

    #[test]
    fn update_phi_works() {
        let mut ledger = ContextLedger::new();
        ledger.record("a.rs", "full", 100, 100);
        ledger.update_phi("a.rs", 0.85);
        assert_eq!(ledger.entries[0].phi, Some(0.85));
    }

    #[test]
    fn set_state_works() {
        let mut ledger = ContextLedger::new();
        ledger.record("a.rs", "full", 100, 100);
        ledger.set_state("a.rs", crate::core::context_field::ContextState::Pinned);
        assert_eq!(
            ledger.entries[0].state,
            Some(crate::core::context_field::ContextState::Pinned)
        );
    }

    #[test]
    fn items_by_state_filters() {
        let mut ledger = ContextLedger::new();
        ledger.record("a.rs", "full", 100, 100);
        ledger.record("b.rs", "full", 100, 100);
        ledger.set_state("b.rs", crate::core::context_field::ContextState::Excluded);
        let included = ledger.items_by_state(crate::core::context_field::ContextState::Included);
        assert_eq!(included.len(), 1);
        assert_eq!(included[0].path, "a.rs");
    }

    #[test]
    fn eviction_by_phi_prefers_low_phi() {
        let mut ledger = ContextLedger::with_window_size(10000);
        ledger.record("high.rs", "full", 100, 100);
        ledger.update_phi("high.rs", 0.9);
        ledger.record("low.rs", "full", 100, 100);
        ledger.update_phi("low.rs", 0.1);
        let candidates = ledger.eviction_candidates_by_phi(1);
        assert_eq!(candidates, vec!["low.rs"]);
    }

    #[test]
    fn eviction_by_phi_skips_pinned() {
        let mut ledger = ContextLedger::with_window_size(10000);
        ledger.record("pinned.rs", "full", 100, 100);
        ledger.update_phi("pinned.rs", 0.01);
        ledger.set_state(
            "pinned.rs",
            crate::core::context_field::ContextState::Pinned,
        );
        ledger.record("normal.rs", "full", 100, 100);
        ledger.update_phi("normal.rs", 0.5);
        let candidates = ledger.eviction_candidates_by_phi(1);
        assert_eq!(candidates, vec!["normal.rs"]);
    }

    #[test]
    fn mark_stale_by_hash_detects_change() {
        let mut ledger = ContextLedger::new();
        ledger.record("a.rs", "full", 100, 100);
        ledger.entries[0].source_hash = Some("hash_v1".to_string());
        ledger.mark_stale_by_hash("a.rs", "hash_v2");
        assert_eq!(
            ledger.entries[0].state,
            Some(crate::core::context_field::ContextState::Stale)
        );
    }

    #[test]
    fn find_by_id_works() {
        let mut ledger = ContextLedger::new();
        ledger.record("src/lib.rs", "full", 100, 100);
        let id = crate::core::context_field::ContextItemId::from_file("src/lib.rs");
        assert!(ledger.find_by_id(&id).is_some());
    }

    #[test]
    fn upsert_sets_source_hash_and_kind() {
        let mut ledger = ContextLedger::new();
        ledger.upsert(
            "src/main.rs",
            "full",
            500,
            500,
            Some("sha256_abc"),
            crate::core::context_field::ContextKind::File,
            None,
        );
        let entry = &ledger.entries[0];
        assert_eq!(entry.source_hash.as_deref(), Some("sha256_abc"));
        assert_eq!(
            entry.kind,
            Some(crate::core::context_field::ContextKind::File)
        );
    }
}