sqz_engine/

engine.rs

use std::path::Path;
use std::sync::{Arc, Mutex};
use std::time::Duration;

use crate::ast_parser::AstParser;
use crate::budget_tracker::{BudgetTracker, UsageReport};
use crate::cache_manager::CacheManager;
use crate::confidence_router::ConfidenceRouter;
use crate::cost_calculator::{CostCalculator, SessionCostSummary};
use crate::ctx_format::CtxFormat;
use crate::error::{Result, SqzError};
use crate::model_router::ModelRouter;
use crate::pin_manager::PinManager;
use crate::pipeline::CompressionPipeline;
use crate::plugin_api::PluginLoader;
use crate::preset::{Preset, PresetParser};
use crate::session_store::{SessionStore, SessionSummary};
use crate::terse_mode::TerseMode;
use crate::types::{CompressedContent, PinEntry, Provenance, SessionId};
use crate::verifier::Verifier;

/// Top-level facade that wires all sqz_engine modules together.
///
/// # Concurrency design
///
/// `SqzEngine` is designed for single-threaded use on the main thread.
/// The only cross-thread sharing happens during preset hot-reload: the
/// file-watcher callback runs on a background thread and needs to update
/// the preset, pipeline, and model router. These three fields are wrapped
/// in `Arc<Mutex<>>` specifically for that purpose. All other fields are
/// owned directly — no unnecessary synchronization.
pub struct SqzEngine {
    // --- Hot-reloadable state (shared with file-watcher thread) ---
    preset: Arc<Mutex<Preset>>,
    pipeline: Arc<Mutex<CompressionPipeline>>,
    model_router: Arc<Mutex<ModelRouter>>,

    // --- Single-owner state (no cross-thread sharing needed) ---
    session_store: SessionStore,
    cache_manager: CacheManager,
    budget_tracker: BudgetTracker,
    cost_calculator: CostCalculator,
    ast_parser: AstParser,
    terse_mode: TerseMode,
    pin_manager: PinManager,
    confidence_router: ConfidenceRouter,
    _plugin_loader: PluginLoader,
}

impl SqzEngine {
    /// Create a new engine with the default preset and a persistent session store.
    ///
    /// Sessions are stored in `~/.sqz/sessions.db` for cross-session continuity.
    /// Falls back to a temp-file store if the home directory is unavailable.
    pub fn new() -> Result<Self> {
        let preset = Preset::default();
        let store_path = Self::default_store_path();
        Self::with_preset_and_store(preset, &store_path)
    }

    /// Resolve the default session store path.
    ///
    /// `SQZ_DB_PATH` overrides everything — set it to keep the database
    /// (stats, dedup cache, sessions) per-project instead of global.
    /// Relative values resolve against the process working directory, so
    /// prefer absolute paths when hooks may run from subdirectories.
    /// Otherwise `~/.sqz/sessions.db`, falling back to a temp-file path
    /// if the home dir is unavailable.
    pub fn default_store_path() -> std::path::PathBuf {
        if let Some(custom) = Self::store_path_from_env(std::env::var("SQZ_DB_PATH").ok()) {
            return custom;
        }
        if let Some(home) = dirs_next::home_dir() {
            let sqz_dir = home.join(".sqz");
            if std::fs::create_dir_all(&sqz_dir).is_ok() {
                // Harden permissions on Unix: ~/.sqz/ contains session data
                // and cached content that may include sensitive output.
                #[cfg(unix)]
                {
                    use std::os::unix::fs::PermissionsExt;
                    let _ = std::fs::set_permissions(
                        &sqz_dir,
                        std::fs::Permissions::from_mode(0o700),
                    );
                }
                return sqz_dir.join("sessions.db");
            }
        }
        // Fallback: temp dir with unique name
        let dir = std::env::temp_dir();
        dir.join(format!(
            "sqz_session_{}_{}.db",
            std::process::id(),
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map(|d| d.as_nanos())
                .unwrap_or(0)
        ))
    }

    /// Turn an `SQZ_DB_PATH` value into a usable store path, creating the
    /// parent directory (0700 on Unix, matching `~/.sqz`). Empty or
    /// whitespace-only values are treated as unset.
    fn store_path_from_env(value: Option<String>) -> Option<std::path::PathBuf> {
        let raw = value?;
        let trimmed = raw.trim();
        if trimmed.is_empty() {
            return None;
        }
        let path = std::path::PathBuf::from(trimmed);
        if let Some(parent) = path.parent() {
            if !parent.as_os_str().is_empty() && !parent.exists() {
                // Harden only directories we create ourselves — never
                // chmod a pre-existing directory the user pointed into.
                if std::fs::create_dir_all(parent).is_ok() {
                    #[cfg(unix)]
                    {
                        use std::os::unix::fs::PermissionsExt;
                        let _ = std::fs::set_permissions(
                            parent,
                            std::fs::Permissions::from_mode(0o700),
                        );
                    }
                }
            }
        }
        Some(path)
    }

    /// Whether informational stderr lines should be suppressed.
    ///
    /// `SQZ_QUIET=1` silences the per-compression progress lines ("N/M
    /// tokens", "dedup hit", "fallback: verifier confidence …"). Errors
    /// still print. Used by `sqz discover --replay`, which compresses
    /// thousands of outputs in one run, and available to users whose
    /// hook setup surfaces stderr where it isn't wanted.
    pub fn quiet_stderr() -> bool {
        match std::env::var("SQZ_QUIET") {
            Ok(v) => !v.is_empty() && v != "0",
            Err(_) => false,
        }
    }

    /// How long a dedup ref stays valid after the original was served.
    ///
    /// `SQZ_REF_TTL_SECS` overrides the 30-minute default. Clients that
    /// compact context aggressively (or have no PreCompact hook) can set
    /// this lower so refs stop pointing at content the model has already
    /// dropped. Unparseable or empty values fall back to the default.
    fn ref_ttl_from_env(value: Option<String>) -> Duration {
        const DEFAULT: Duration = Duration::from_secs(30 * 60);
        match value.as_deref().map(str::trim) {
            Some(raw) if !raw.is_empty() => raw
                .parse::<u64>()
                .map(Duration::from_secs)
                .unwrap_or(DEFAULT),
            _ => DEFAULT,
        }
    }

    /// Create with a custom preset and a file-backed session store.
    ///
    /// Opens a single SQLite connection for the session store. The cache
    /// manager and pin manager share the same store via separate connections
    /// (SQLite WAL mode supports concurrent readers).
    pub fn with_preset_and_store(preset: Preset, store_path: &Path) -> Result<Self> {
        let pipeline = CompressionPipeline::new(&preset);
        let window_size = preset.budget.default_window_size;

        // One connection per consumer. SQLite WAL mode handles concurrency.
        let session_store = SessionStore::open_or_create(store_path)?;
        let cache_store = SessionStore::open_or_create(store_path)?;
        let pin_store = SessionStore::open_or_create(store_path)?;

        Ok(SqzEngine {
            preset: Arc::new(Mutex::new(preset.clone())),
            pipeline: Arc::new(Mutex::new(pipeline)),
            model_router: Arc::new(Mutex::new(ModelRouter::new(&preset))),
            session_store,
            cache_manager: CacheManager::with_ref_age_duration(
                cache_store,
                512 * 1024 * 1024,
                Self::ref_ttl_from_env(std::env::var("SQZ_REF_TTL_SECS").ok()),
            ),
            budget_tracker: BudgetTracker::new(window_size, &preset),
            cost_calculator: CostCalculator::with_defaults(),
            ast_parser: AstParser::new(),
            terse_mode: TerseMode,
            pin_manager: PinManager::new(pin_store),
            confidence_router: ConfidenceRouter::new(),
            _plugin_loader: PluginLoader::new(Path::new("plugins")),
        })
    }

    /// Compress input text using the current preset.
    ///
    /// Two-pass pipeline:
    /// 1. Route to compression mode based on content entropy and risk patterns.
    /// 2. Compress using the pipeline (safe preset for Safe mode, default otherwise).
    /// 3. Verify invariants (error lines, JSON keys, diff hunks, etc.).
    /// 4. If verification confidence is low, fall back to safe mode and re-compress.
    pub fn compress(&self, input: &str) -> Result<CompressedContent> {
        let preset = self.preset.lock()
            .map_err(|_| SqzError::Other("preset lock poisoned".into()))?;
        let pipeline = self.pipeline.lock()
            .map_err(|_| SqzError::Other("pipeline lock poisoned".into()))?;
        let ctx = crate::pipeline::SessionContext {
            session_id: "engine".to_string(),
        };

        // Step 1: Route — check content risk before compressing
        let mode = self.confidence_router.route(input);

        // Step 2: If Safe mode, skip aggressive pipeline and go straight to safe compress
        if mode == crate::confidence_router::CompressionMode::Safe {
            // Silent by default: in the hook path every stderr line lands in
            // the agent's context, and a passthrough has nothing to report.
            if std::env::var_os("SQZ_VERBOSE").is_some() {
                eprintln!("[sqz] safe mode: high-risk content passed through unchanged");
            }
            return self.compress_safe(input, &pipeline, &ctx, true);
        }

        // Step 3: Compress with the configured pipeline
        let mut result = pipeline.compress_uncounted(input, &ctx, &preset)?;

        // Step 4: Verify invariants
        let verify = Verifier::verify(input, &result.data);
        let fallback = verify.fallback_triggered;
        result.verify = Some(verify);

        // Step 5: If verifier signals low confidence, re-compress with safe settings
        if fallback && result.data != input {
            if !Self::quiet_stderr() {
                eprintln!("[sqz] fallback: verifier confidence {:.2} below threshold — re-compressing in safe mode",
                    result.verify.as_ref().map(|v| v.confidence).unwrap_or(0.0));
            }
            let safe_result = self.compress_safe(input, &pipeline, &ctx, false)?;
            return Ok(safe_result);
        }

        pipeline.count_output(&mut result, &preset);
        Ok(result)
    }

    /// Compress with dedup cache lookup.
    ///
    /// Unlike [`compress`](Self::compress), this method consults the persistent cache
    /// (`~/.sqz/sessions.db`). On a hit with a fresh ref, it returns a
    /// 13-token `§ref:HASH§` token instead of re-compressing. On a
    /// near-duplicate, it returns a compact delta. On a cache miss, it
    /// runs the full pipeline, stores the result, and returns it.
    /// Content [`never_store`](Self::never_store) flags skips the cache.
    ///
    /// This is the path the MCP `sqz_read_file` / `sqz_grep` /
    /// `sqz_list_dir` tools use — repeat reads in the same session
    /// (and across sessions if the DB survives) collapse to 13 tokens.
    /// Reported on issue #12: the MCP file tools bypassed the cache so
    /// dedup never fired, and users were getting 30%-range pipeline
    /// compression instead of the 92% dedup path advertised in the
    /// README.
    pub fn compress_with_cache(&self, input: &str) -> Result<crate::cache_manager::CacheResult> {
        // The `path` argument is informational only — cache keys are
        // SHA-256 of the content bytes, not the path. Pass an empty
        // path so we don't partition the cache by accidental file
        // location differences (e.g. `./src/main.rs` vs
        // `/abs/path/to/src/main.rs`). Misses go through `compress`, so
        // the verifier's safe-mode fallback applies here too.
        if self.never_store(input) {
            return Ok(crate::cache_manager::CacheResult::Fresh { output: self.compress(input)? });
        }
        self.cache_manager.get_or_compress_with(
            std::path::Path::new(""),
            input.as_bytes(),
            &|text| self.compress(text),
        )
    }

    /// Lossless variant of [`compress_with_cache`](Self::compress_with_cache)
    /// for the file-read MCP tools
    /// (`sqz_read_file` / `sqz_grep` / `sqz_list_dir`): on a cache miss the
    /// content is returned faithfully — no entropy truncation, pruning, or
    /// re-encoding. The general `compress` tool keeps using
    /// [`compress_with_cache`](Self::compress_with_cache) for aggressive
    /// compression. See issue #32.
    pub fn compress_with_cache_lossless(&self, input: &str) -> Result<crate::cache_manager::CacheResult> {
        let pipeline = self.pipeline.lock()
            .map_err(|_| SqzError::Other("pipeline lock poisoned".into()))?;
        if self.never_store(input) {
            let ctx = crate::pipeline::SessionContext { session_id: "cache".to_string() };
            let output = pipeline.compress_lossless(input, &ctx, &Preset::default())?;
            return Ok(crate::cache_manager::CacheResult::Fresh { output });
        }
        self.cache_manager.get_or_compress_lossless(
            std::path::Path::new(""),
            input.as_bytes(),
            &pipeline,
        )
    }

    /// What [`Self::compress_with_cache_lossless`] returns on a miss, without the cache.
    pub fn compress_lossless(&self, input: &str) -> Result<CompressedContent> {
        let pipeline = self.pipeline.lock()
            .map_err(|_| SqzError::Other("pipeline lock poisoned".into()))?;
        let ctx = crate::pipeline::SessionContext { session_id: "cache".to_string() };
        pipeline.compress_lossless(input, &ctx, &Preset::default())
    }

    /// What `sqz-mcp proxy` does to one text result from an upstream MCP
    /// server. `None` keeps the original: too small to pay for the
    /// markers, high-risk content (credentials, stack traces,
    /// migrations), or a saving under the net-win gate. `sqz discover
    /// --replay` uses it too, so its MCP estimate is what the proxy does.
    pub fn compress_tool_result(&self, text: &str, use_cache: bool) -> Option<ToolResultCompression> {
        use crate::cache_manager::CacheResult;
        let tokens_original = chars_over_four(text);
        if tokens_original < TOOL_RESULT_MIN_SAVING * 2 {
            return None;
        }
        if self.route_compression_mode(text) == crate::confidence_router::CompressionMode::Safe {
            return None;
        }
        let use_cache = use_cache && !self.never_store(text);

        let result = if use_cache {
            self.compress_with_cache(text).ok()
        } else {
            self.compress(text).ok().map(|output| CacheResult::Fresh { output })
        }?;

        let (out, tokens_compressed, repeat) = match result {
            CacheResult::Dedup { inline_ref, token_cost } => (inline_ref, token_cost, true),
            CacheResult::Delta { delta_text, token_cost, .. } => (delta_text, token_cost, true),
            CacheResult::Fresh { output } => {
                // First pass is lossy, so point at the cached original.
                let mut data = output.data;
                if use_cache {
                    let hash = crate::cache_manager::CacheManager::sha256_hex(text.as_bytes());
                    data.push_str(&format!(
                        "\n[full output: call sqz_expand with ref \"{}\"]",
                        &hash[..16]
                    ));
                }
                let t = chars_over_four(&data);
                (data, t, false)
            }
        };
        if tokens_original.saturating_sub(tokens_compressed) < TOOL_RESULT_MIN_SAVING {
            return None;
        }
        Some(ToolResultCompression {
            text: out,
            tokens_original,
            tokens_compressed,
            repeat,
        })
    }

    /// Defensive compression: any input in, `CompressedContent` out, guaranteed.
    ///
    /// Unlike `compress()` which returns `Result`, this method never returns
    /// an error. On any internal failure it returns the original input
    /// unchanged with a 1.0 compression ratio. This makes it safe to call
    /// from contexts where error handling is impractical (e.g. shell hooks,
    /// browser extension bridges).
    pub fn compress_or_passthrough(&self, input: &str) -> CompressedContent {
        match self.compress(input) {
            Ok(result) => result,
            Err(_) => {
                let tokens = (input.len() as u32 + 3) / 4;
                CompressedContent {
                    data: input.to_string(),
                    tokens_compressed: tokens,
                    tokens_original: tokens,
                    stages_applied: vec![],
                    compression_ratio: 1.0,
                    provenance: crate::types::Provenance::default(),
                    verify: None,
                }
            }
        }
    }

    /// Compress with explicit mode override, bypassing the confidence router.
    ///
    /// - `CompressionMode::Safe` → safe pipeline only (ANSI strip + condense)
    /// - `CompressionMode::Default` → standard pipeline
    /// - `CompressionMode::Aggressive` → standard pipeline (aggressive preset TBD)
    pub fn compress_with_mode(&self, input: &str, mode: crate::confidence_router::CompressionMode) -> Result<CompressedContent> {
        let pipeline = self.pipeline.lock()
            .map_err(|_| SqzError::Other("pipeline lock poisoned".into()))?;
        let ctx = crate::pipeline::SessionContext {
            session_id: "engine".to_string(),
        };

        match mode {
            crate::confidence_router::CompressionMode::Safe => {
                let lossless = self.confidence_router.route(input)
                    == crate::confidence_router::CompressionMode::Safe;
                self.compress_safe(input, &pipeline, &ctx, lossless)
            }
            _ => {
                // Default and Aggressive: run normal pipeline + verify
                drop(pipeline); // release lock before calling compress()
                self.compress(input)
            }
        }
    }

    /// Safe-mode compression: minimal transforms only (ANSI strip + condense).
    fn compress_safe(
        &self,
        input: &str,
        pipeline: &crate::pipeline::CompressionPipeline,
        ctx: &crate::pipeline::SessionContext,
        lossless: bool,
    ) -> Result<CompressedContent> {
        use crate::preset::{
            CompressionConfig, CondenseConfig, CustomTransformsConfig, BudgetConfig,
            ModelConfig, PresetMeta, TerseModeConfig, TerseLevel, ToolSelectionConfig,
        };

        let safe_preset = Preset {
            preset: PresetMeta {
                name: "safe".to_string(),
                version: "1.0".to_string(),
                description: "Safe fallback — minimal compression".to_string(),
            },
            compression: CompressionConfig {
                stages: vec!["condense".to_string()],
                keep_fields: None,
                strip_fields: None,
                condense: Some(CondenseConfig { enabled: true, max_repeated_lines: 3 }),
                git_diff_fold: None,
                strip_nulls: None,
                flatten: None,
                truncate_strings: None,
                collapse_arrays: None,
                custom_transforms: Some(CustomTransformsConfig { enabled: false }),
            },
            tool_selection: ToolSelectionConfig {
                max_tools: 5,
                similarity_threshold: 0.7,
                default_tools: vec![],
            },
            budget: BudgetConfig {
                warning_threshold: 0.70,
                ceiling_threshold: 0.85,
                default_window_size: 200_000,
                agents: Default::default(),
            },
            terse_mode: TerseModeConfig { enabled: false, level: TerseLevel::Moderate },
            model: ModelConfig {
                family: "anthropic".to_string(),
                primary: String::new(),
                local: String::new(),
                complexity_threshold: 0.4,
                pricing: None,
            },
        };

        let mut result = if lossless {
            pipeline.compress_lossless(input, ctx, &safe_preset)?
        } else {
            pipeline.compress(input, ctx, &safe_preset)?
        };
        let verify = Verifier::verify(input, &result.data);
        result.verify = Some(verify);
        result.provenance = Provenance {
            label: Some("safe-fallback".to_string()),
            ..Default::default()
        };
        Ok(result)
    }

    /// Compress with explicit provenance metadata attached to the result.
    pub fn compress_with_provenance(
        &self,
        input: &str,
        provenance: Provenance,
    ) -> Result<CompressedContent> {
        let mut result = self.compress(input)?;
        result.provenance = provenance;
        Ok(result)
    }

    /// Export a session to CTX format.
    pub fn export_ctx(&self, session_id: &str) -> Result<String> {
        let session = self.session_store.load_session(session_id.to_string())?;
        CtxFormat::serialize(&session)
    }

    /// Import a CTX string and save as a new session.
    pub fn import_ctx(&self, ctx: &str) -> Result<SessionId> {
        let session = CtxFormat::deserialize(ctx)?;
        self.session_store.save_session(&session)
    }

    /// Pin a conversation turn.
    pub fn pin(&self, session_id: &str, turn_index: usize, reason: &str, tokens: u32) -> Result<PinEntry> {
        self.pin_manager.pin(session_id, turn_index, reason, tokens)
    }

    /// Unpin a conversation turn.
    pub fn unpin(&self, session_id: &str, turn_index: usize) -> Result<()> {
        self.pin_manager.unpin(session_id, turn_index)
    }

    /// Search sessions by keyword.
    pub fn search_sessions(&self, query: &str) -> Result<Vec<SessionSummary>> {
        self.session_store.search(query)
    }

    /// Get usage report for an agent.
    pub fn usage_report(&self, agent_id: &str) -> UsageReport {
        self.budget_tracker.usage_report(agent_id.to_string())
    }

    /// Get cost summary for a session.
    pub fn cost_summary(&self, session_id: &str) -> Result<SessionCostSummary> {
        let session = self.session_store.load_session(session_id.to_string())?;
        Ok(self.cost_calculator.session_summary(&session))
    }

    /// Reload the preset from a TOML string (hot-reload support).
    pub fn reload_preset(&mut self, toml: &str) -> Result<()> {
        let new_preset = PresetParser::parse(toml)?;
        if let Ok(mut pipeline) = self.pipeline.lock() {
            pipeline.reload_preset(&new_preset)?;
        }
        if let Ok(mut router) = self.model_router.lock() {
            *router = ModelRouter::new(&new_preset);
        }
        if let Ok(mut preset) = self.preset.lock() {
            *preset = new_preset;
        }
        Ok(())
    }

    /// Spawn a background thread that watches `path` for preset file changes.
    ///
    /// Only the preset, pipeline, and model_router are shared with the watcher
    /// thread (via `Arc<Mutex<>>`). All other engine state stays on the main thread.
    pub fn watch_preset_file(&self, path: &Path) -> Result<notify::RecommendedWatcher> {
        use notify::{Event, EventKind, RecursiveMode, Watcher};

        let preset_arc = Arc::clone(&self.preset);
        let pipeline_arc = Arc::clone(&self.pipeline);
        let router_arc = Arc::clone(&self.model_router);
        let watched_path = path.to_owned();

        let mut watcher = notify::recommended_watcher(move |res: notify::Result<Event>| {
            if let Ok(event) = res {
                if matches!(event.kind, EventKind::Modify(_) | EventKind::Create(_)) {
                    match std::fs::read_to_string(&watched_path) {
                        Ok(toml_str) => match PresetParser::parse(&toml_str) {
                            Ok(new_preset) => {
                                if let Ok(mut p) = pipeline_arc.lock() {
                                    let _ = p.reload_preset(&new_preset);
                                }
                                if let Ok(mut r) = router_arc.lock() {
                                    *r = ModelRouter::new(&new_preset);
                                }
                                if let Ok(mut pr) = preset_arc.lock() {
                                    *pr = new_preset;
                                }
                            }
                            Err(e) => eprintln!("[sqz] invalid preset: {e}"),
                        },
                        Err(e) => eprintln!("[sqz] preset read error: {e}"),
                    }
                }
            }
        })
        .map_err(|e| SqzError::Other(format!("watcher error: {e}")))?;

        watcher
            .watch(path, RecursiveMode::NonRecursive)
            .map_err(|e| SqzError::Other(format!("watch error: {e}")))?;

        Ok(watcher)
    }

    /// Access the underlying `SessionStore`.
    pub fn session_store(&self) -> &SessionStore {
        &self.session_store
    }

    /// Access the `CacheManager` for persistent dedup.
    pub fn cache_manager(&self) -> &CacheManager {
        &self.cache_manager
    }

    /// Access the `AstParser`.
    pub fn ast_parser(&self) -> &AstParser {
        &self.ast_parser
    }

    /// Access the `TerseMode` helper.
    pub fn terse_mode(&self) -> &TerseMode {
        &self.terse_mode
    }

    /// Reorder context sections using the LITM positioner to mitigate
    /// the "Lost In The Middle" attention bias in long-context models.
    ///
    /// Places highest-priority sections at the beginning and end of the
    /// context window, lowest-priority in the middle.
    pub fn reorder_context(
        &self,
        sections: &mut Vec<crate::litm_positioner::ContextSection>,
        strategy: crate::litm_positioner::LitmStrategy,
    ) {
        let positioner = crate::litm_positioner::LitmPositioner::new(strategy);
        positioner.reorder(sections);
    }

    /// Route content to the appropriate compression mode based on entropy
    /// and risk pattern analysis.
    pub fn route_compression_mode(&self, content: &str) -> crate::confidence_router::CompressionMode {
        self.confidence_router.route(content)
    }

    /// Whether content must stay out of the persistent cache.
    pub fn never_store(&self, content: &str) -> bool {
        self.confidence_router.never_store(content)
    }

    /// Count tokens the way `compress` counts `tokens_original`.
    pub fn count_tokens(&self, text: &str) -> u32 {
        let preset = self.preset.lock().unwrap();
        self.pipeline.lock().unwrap().count_tokens(text, &preset)
    }

    /// Clear the lock poison a caught panic left, so the engine stays usable.
    pub fn clear_poison(&self) {
        self.preset.clear_poison();
        self.pipeline.clear_poison();
        self.model_router.clear_poison();
    }
}

/// Smallest saving, in tokens, worth replacing an MCP tool result for: the
/// markers and the agent's attention cost about this much.
pub const TOOL_RESULT_MIN_SAVING: u32 = 16;

/// A compressed MCP tool result from [`SqzEngine::compress_tool_result`].
/// Token counts are chars/4 estimates, the counter the net-win gate uses.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolResultCompression {
    pub text: String,
    pub tokens_original: u32,
    pub tokens_compressed: u32,
    /// A ref or delta against content already sent: lossless. Otherwise a
    /// first-pass compression, which is lossy.
    pub repeat: bool,
}

fn chars_over_four(text: &str) -> u32 {
    ((text.len() as f64) / 4.0).ceil() as u32
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{BudgetState, CorrectionLog, ModelFamily, SessionState};
    use chrono::Utc;
    use std::path::PathBuf;

    fn make_session(id: &str) -> SessionState {
        let now = Utc::now();
        SessionState {
            id: id.to_string(),
            project_dir: PathBuf::from("/tmp/test"),
            conversation: vec![],
            corrections: CorrectionLog::default(),
            pins: vec![],
            learnings: vec![],
            compressed_summary: "test session".to_string(),
            budget: BudgetState {
                window_size: 200_000,
                consumed: 0,
                pinned: 0,
                model_family: ModelFamily::AnthropicClaude,
            },
            tool_usage: vec![],
            created_at: now,
            updated_at: now,
        }
    }

    #[test]
    fn test_engine_new() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db"));
        assert!(engine.is_ok(), "engine with the default preset should start");
    }

    #[test]
    fn test_compress_or_passthrough_returns_result_on_valid_input() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let result = engine.compress_or_passthrough("hello world");
        assert_eq!(result.data, "hello world");
        assert!(result.tokens_original > 0);
    }

    #[test]
    fn test_compress_or_passthrough_never_panics_on_empty() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let result = engine.compress_or_passthrough("");
        assert_eq!(result.data, "");
        assert_eq!(result.compression_ratio, 1.0);
    }

    #[test]
    fn test_compress_or_passthrough_handles_json() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let result = engine.compress_or_passthrough(r#"{"key":"value"}"#);
        // Should compress successfully — data may be TOON-encoded
        assert!(!result.data.is_empty());
    }

    #[test]
    fn test_compress_or_passthrough_handles_binary_garbage() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        // Feed it something weird — should never panic, always return something
        let garbage = "\x00\x01\x02\x7f invalid control chars \t\n\r";
        let result = engine.compress_or_passthrough(garbage);
        assert!(!result.data.is_empty());
    }

    #[test]
    fn test_compress_plain_text() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let result = engine.compress("hello world");
        assert!(result.is_ok());
        assert_eq!(result.unwrap().data, "hello world");
    }

    #[test]
    fn test_compress_json_applies_toon() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let result = engine.compress(r#"{"name":"Alice","age":30}"#).unwrap();
        assert!(result.data.starts_with("TOON:"), "JSON should be TOON-encoded");
    }

    #[test]
    fn compress_tool_result_gates_like_the_proxy() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();

        assert_eq!(engine.compress_tool_result("ok, 3 rows updated", true), None, "too small to pay off");

        let rows: String = (0..80)
            .map(|i| format!("issue row {i} status open priority medium assignee nobody\n"))
            .collect();
        let first = engine.compress_tool_result(&rows, true).expect("compresses");
        assert!(first.tokens_compressed + TOOL_RESULT_MIN_SAVING <= first.tokens_original);
        assert!(!first.text.starts_with("§ref:") && !first.repeat);
        // A lossy first pass always says how to get the original back.
        let hash = CacheManager::sha256_hex(rows.as_bytes());
        assert!(first.text.ends_with(&format!("[full output: call sqz_expand with ref \"{}\"]", &hash[..16])), "{}", first.text);
        let second = engine.compress_tool_result(&rows, true).expect("repeat");
        assert!(second.text.starts_with("§ref:") && second.repeat, "{}", second.text);
        let uncached = engine.compress_tool_result(&rows, false).expect("no cache");
        assert!(!uncached.text.starts_with("§ref:") && !uncached.text.contains("sqz_expand"));
    }

    #[test]
    fn clear_poison_after_a_caught_panic() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
            let _preset = engine.preset.lock().unwrap();
            let _pipeline = engine.pipeline.lock().unwrap();
            panic!("a bug while compressing");
        }));
        assert!(caught.is_err());
        assert!(engine.compress("hello world").is_err());

        engine.clear_poison();
        assert_eq!(engine.compress("hello world").unwrap().data, "hello world");
        assert!(engine.count_tokens("hello world") > 0);
    }

    fn high_risk_inputs() -> Vec<(&'static str, String)> {
        let frames: String = (0..40)
            .map(|i| format!("  {i:>2}: myapp::scheduler::step_{i}\n             at ./src/scheduler/mod_{}.rs:{}:5\n", i % 7, 40 + 7 * i))
            .collect();
        let tables: String = ["customers", "orders", "refunds", "payments", "shipments", "invoices"]
            .iter()
            .map(|t| format!("CREATE TABLE {t} (\n    id BIGSERIAL PRIMARY KEY,\n    status TEXT NOT NULL DEFAULT 'pending',\n    created_at TIMESTAMPTZ NOT NULL DEFAULT now()\n);\nCREATE INDEX idx_{t}_created_at ON {t} (created_at);\n"))
            .collect();
        let pem: String = (0..20).map(|i| CacheManager::sha256_hex(format!("pem {i}").as_bytes()) + "\n").collect();
        let creds: String = ["stripe", "sendgrid", "github", "datadog", "sentry"]
            .iter()
            .map(|s| format!("{s}_api_key: sk_test_{}\n{s}_endpoint: https://{s}.example.com/v1\n", &CacheManager::sha256_hex(s.as_bytes())[..32]))
            .collect();
        vec![
            ("panic", format!("thread 'main' panicked at src/scheduler/queue.rs:217:31:\ncalled `Option::unwrap()` on a `None` value\nstack backtrace:\n{frames}")),
            ("migration", format!("BEGIN;\n{tables}ALTER TABLE refunds ADD COLUMN reason TEXT;\nCOMMIT;\n")),
            ("keys", format!("-----BEGIN RSA PRIVATE KEY-----\n{pem}-----END RSA PRIVATE KEY-----\n{creds}")),
        ]
    }

    #[test]
    fn high_risk_content_passes_through_unchanged() {
        use crate::confidence_router::CompressionMode;
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let inputs = high_risk_inputs();
        for (name, input) in &inputs {
            assert!(input.len() > 500, "{name}");
            assert_eq!(engine.route_compression_mode(input), CompressionMode::Safe, "{name}");
            assert_eq!(&engine.compress(input).unwrap().data, input, "{name}");
            for mode in [CompressionMode::Safe, CompressionMode::Default, CompressionMode::Aggressive] {
                assert_eq!(&engine.compress_with_mode(input, mode).unwrap().data, input, "{name} {mode:?}");
            }
        }
        let colored = format!("\x1b[31m{}\x1b[0m", inputs[0].1);
        assert_eq!(engine.compress(&colored).unwrap().data, inputs[0].1);
    }

    #[test]
    fn verifier_fallback_still_compresses() {
        use crate::confidence_router::CompressionMode;
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let input = (0..24)
            .map(|i| {
                if i % 2 == 0 {
                    format!("SEG{i} /srv/app/mod_{i}.py llllllllllllllllllllllllllllllllllllllll")
                } else {
                    format!("SEG{i} the quick brown fox jumps over {i} lazy dogs by silver rivers")
                }
            })
            .collect::<Vec<_>>()
            .join("\n\n");
        assert_ne!(engine.route_compression_mode(&input), CompressionMode::Safe);
        let result = engine.compress(&input).unwrap();
        assert_eq!(result.provenance.label.as_deref(), Some("safe-fallback"));
        assert!(result.data.contains("low-information segments omitted"), "{}", result.data);
        assert_eq!(result.data, engine.compress_with_mode(&input, CompressionMode::Safe).unwrap().data);
    }

    #[test]
    fn json_values_survive_every_mode() {
        use crate::confidence_router::CompressionMode;
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let pretty = "[\n  0,\n  0,\n  0,\n  0,\n  0,\n  0,\n  1\n]\n";
        let expected = serde_json::json!([0, 0, 0, 0, 0, 0, 1]);
        let mut outputs = vec![engine.compress(pretty).unwrap().data];
        for mode in [CompressionMode::Safe, CompressionMode::Default, CompressionMode::Aggressive] {
            outputs.push(engine.compress_with_mode(pretty, mode).unwrap().data);
        }
        for data in outputs {
            assert_eq!(crate::toon::ToonEncoder.decode(&data).unwrap(), expected, "{data}");
        }
    }

    #[test]
    fn never_stored_content_skips_every_cache() {
        use crate::cache_manager::CacheResult;
        use crate::confidence_router::CompressionMode;
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let mut log: String = (0..120)
            .map(|i| format!("2026-10-05T12:{:02}:{:02}Z INFO  GET /healthz 200 1ms\n", i / 60, i % 60))
            .collect();
        log.push_str("2026-10-05T12:30:00Z DEBUG db config host=db.internal user=app password=hunter2\n");
        assert_ne!(engine.route_compression_mode(&log), CompressionMode::Safe);
        assert!(engine.never_store(&log));
        let stored = || engine.cache_manager().check_dedup_with_meta(log.as_bytes()).unwrap().is_some();

        let CacheResult::Fresh { output } = engine.compress_with_cache(&log).unwrap() else { panic!("not fresh") };
        assert!(output.data.len() < log.len());
        let CacheResult::Fresh { output } = engine.compress_with_cache_lossless(&log).unwrap() else { panic!("not fresh") };
        assert_eq!(output.data, log);
        let result = engine.compress_tool_result(&log, true).expect("compresses");
        assert!(!result.repeat && !result.text.contains("sqz_expand"), "{}", result.text);
        assert!(!stored());
        assert!(matches!(engine.compress_with_cache(&log).unwrap(), CacheResult::Fresh { .. }));
    }

    #[test]
    fn test_export_import_ctx_round_trip() {
        let dir = tempfile::tempdir().unwrap();
        let store_path = dir.path().join("store.db");
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &store_path).unwrap();

        let session = make_session("sess-rt");
        engine.session_store().save_session(&session).unwrap();

        let ctx = engine.export_ctx("sess-rt").unwrap();
        let imported_id = engine.import_ctx(&ctx).unwrap();
        assert_eq!(imported_id, "sess-rt");
    }

    #[test]
    fn test_search_sessions() {
        let dir = tempfile::tempdir().unwrap();
        let store_path = dir.path().join("store.db");
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &store_path).unwrap();

        let mut session = make_session("sess-search");
        session.compressed_summary = "authentication refactor".to_string();
        engine.session_store().save_session(&session).unwrap();

        let results = engine.search_sessions("authentication").unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, "sess-search");
    }

    #[test]
    fn test_usage_report_starts_at_zero() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let report = engine.usage_report("default");
        assert_eq!(report.consumed, 0);
        assert_eq!(report.available, report.allocated);
    }

    #[test]
    fn test_cost_summary() {
        let dir = tempfile::tempdir().unwrap();
        let store_path = dir.path().join("store.db");
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &store_path).unwrap();

        let session = make_session("sess-cost");
        engine.session_store().save_session(&session).unwrap();

        let summary = engine.cost_summary("sess-cost").unwrap();
        assert_eq!(summary.total_tokens, 0);
        assert!((summary.total_usd - 0.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_reload_preset_updates_state() {
        let dir = tempfile::tempdir().unwrap();
        let mut engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let toml = r#"
[preset]
name = "reloaded"
version = "2.0"

[compression]
stages = []

[tool_selection]
max_tools = 5
similarity_threshold = 0.7

[budget]
warning_threshold = 0.70
ceiling_threshold = 0.85
default_window_size = 200000

[terse_mode]
enabled = false
level = "moderate"

[model]
family = "anthropic"
primary = "claude-sonnet-4-20250514"
complexity_threshold = 0.4
"#;
        assert!(engine.reload_preset(toml).is_ok());
        // Verify the preset was actually updated
        let preset = engine.preset.lock().unwrap();
        assert_eq!(preset.preset.name, "reloaded");
    }

    #[test]
    fn test_reload_invalid_preset_returns_error() {
        let dir = tempfile::tempdir().unwrap();
        let mut engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let result = engine.reload_preset("not valid toml [[[");
        assert!(result.is_err(), "invalid TOML should return error");
    }

    #[test]
    fn test_export_nonexistent_session_returns_error() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let result = engine.export_ctx("does-not-exist");
        assert!(result.is_err());
    }

    #[test]
    fn test_import_invalid_ctx_returns_error() {
        let dir = tempfile::tempdir().unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &dir.path().join("store.db")).unwrap();
        let result = engine.import_ctx("not valid json {{{");
        assert!(result.is_err());
    }

    #[test]
    fn store_path_env_override_used_when_set() {
        let dir = tempfile::tempdir().unwrap();
        let target = dir.path().join("nested").join("project.db");
        let resolved =
            SqzEngine::store_path_from_env(Some(target.to_string_lossy().to_string()));
        assert_eq!(resolved, Some(target.clone()));
        // Parent directory is created so SQLite can open the file.
        assert!(target.parent().unwrap().exists());
    }

    #[test]
    fn store_path_env_ignores_empty_values() {
        assert_eq!(SqzEngine::store_path_from_env(None), None);
        assert_eq!(SqzEngine::store_path_from_env(Some(String::new())), None);
        assert_eq!(SqzEngine::store_path_from_env(Some("   ".to_string())), None);
    }

    #[test]
    fn store_path_env_leaves_existing_parent_permissions_alone() {
        let dir = tempfile::tempdir().unwrap();
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            std::fs::set_permissions(dir.path(), std::fs::Permissions::from_mode(0o755))
                .unwrap();
        }
        let target = dir.path().join("sessions.db");
        let resolved =
            SqzEngine::store_path_from_env(Some(target.to_string_lossy().to_string()));
        assert_eq!(resolved, Some(target));
        #[cfg(unix)]
        {
            use std::os::unix::fs::PermissionsExt;
            let mode = std::fs::metadata(dir.path()).unwrap().permissions().mode() & 0o777;
            assert_eq!(mode, 0o755, "pre-existing parent must not be chmodded");
        }
    }

    #[test]
    fn ref_ttl_env_parses_seconds_and_falls_back() {
        use std::time::Duration;
        // Unset, empty, whitespace, garbage → 30-minute default.
        assert_eq!(SqzEngine::ref_ttl_from_env(None), Duration::from_secs(1800));
        assert_eq!(
            SqzEngine::ref_ttl_from_env(Some(String::new())),
            Duration::from_secs(1800)
        );
        assert_eq!(
            SqzEngine::ref_ttl_from_env(Some("  ".into())),
            Duration::from_secs(1800)
        );
        assert_eq!(
            SqzEngine::ref_ttl_from_env(Some("ten".into())),
            Duration::from_secs(1800)
        );
        assert_eq!(
            SqzEngine::ref_ttl_from_env(Some("-5".into())),
            Duration::from_secs(1800)
        );
        // Valid values are honored, including 0 (disable refs entirely).
        assert_eq!(
            SqzEngine::ref_ttl_from_env(Some("300".into())),
            Duration::from_secs(300)
        );
        assert_eq!(
            SqzEngine::ref_ttl_from_env(Some(" 60 ".into())),
            Duration::from_secs(60)
        );
        assert_eq!(
            SqzEngine::ref_ttl_from_env(Some("0".into())),
            Duration::ZERO
        );
    }

    #[test]
    fn engine_with_custom_store_path_round_trips() {
        let dir = tempfile::tempdir().unwrap();
        let db = dir.path().join("proj").join("sessions.db");
        let resolved = SqzEngine::store_path_from_env(Some(db.to_string_lossy().to_string()))
            .unwrap();
        let engine = SqzEngine::with_preset_and_store(Preset::default(), &resolved).unwrap();
        engine
            .session_store()
            .log_compression(100, 40, &[], "test")
            .unwrap();
        assert!(db.exists(), "database file must land at the custom path");
        let stats = engine.session_store().compression_stats().unwrap();
        assert_eq!(stats.total_compressions, 1);
    }
}