sqz_engine/

engine.rs

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

use crate::ast_parser::AstParser;
use crate::budget_tracker::{BudgetTracker, UsageReport};
use crate::cache_manager::CacheManager;
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, SessionId};

/// 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,
    #[allow(dead_code)] // used internally by compress pipeline; public API pending
    cache_manager: CacheManager,
    budget_tracker: BudgetTracker,
    cost_calculator: CostCalculator,
    ast_parser: AstParser,
    terse_mode: TerseMode,
    pin_manager: PinManager,
    _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/sessions.db`.
    /// Falls back to a temp-file path if home dir is unavailable.
    fn default_store_path() -> std::path::PathBuf {
        if let Some(home) = dirs_next::home_dir() {
            let sqz_dir = home.join(".sqz");
            if std::fs::create_dir_all(&sqz_dir).is_ok() {
                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)
        ))
    }

    /// 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::new(cache_store, 512 * 1024 * 1024),
            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),
            _plugin_loader: PluginLoader::new(Path::new("plugins")),
        })
    }

    /// Compress input text using the current preset.
    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(),
        };
        pipeline.compress(input, &ctx, &preset)
    }

    /// 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 `AstParser`.
    pub fn ast_parser(&self) -> &AstParser {
        &self.ast_parser
    }

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


#[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 engine = SqzEngine::new();
        assert!(engine.is_ok(), "SqzEngine::new() should succeed");
    }

    #[test]
    fn test_compress_plain_text() {
        let engine = SqzEngine::new().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 engine = SqzEngine::new().unwrap();
        let result = engine.compress(r#"{"name":"Alice","age":30}"#).unwrap();
        assert!(result.data.starts_with("TOON:"), "JSON should be TOON-encoded");
    }

    #[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 engine = SqzEngine::new().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 mut engine = SqzEngine::new().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 mut engine = SqzEngine::new().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 engine = SqzEngine::new().unwrap();
        let result = engine.export_ctx("does-not-exist");
        assert!(result.is_err());
    }

    #[test]
    fn test_import_invalid_ctx_returns_error() {
        let engine = SqzEngine::new().unwrap();
        let result = engine.import_ctx("not valid json {{{");
        assert!(result.is_err());
    }
}