sol/src/persistence.rs

use rusqlite::{params, Connection, Result as SqlResult};
use std::path::Path;
use std::sync::Mutex;
use tracing::{info, warn};

/// SQLite-backed persistent state for Sol.
///
/// Stores:
/// - Conversation registry: room_id → Mistral conversation_id + token estimates
/// - Agent registry: agent_name → Mistral agent_id
///
/// ## Kubernetes mount
///
/// The database file should be on a persistent volume. Recommended setup:
///
/// ```yaml
/// # PersistentVolumeClaim (Longhorn)
/// apiVersion: v1
/// kind: PersistentVolumeClaim
/// metadata:
///   name: sol-data
///   namespace: matrix
/// spec:
///   accessModes: [ReadWriteOnce]
///   storageClassName: longhorn
///   resources:
///     requests:
///       storage: 1Gi
///
/// # Deployment volume mount
/// volumes:
///   - name: sol-data
///     persistentVolumeClaim:
///       claimName: sol-data
/// containers:
///   - name: sol
///     volumeMounts:
///       - name: sol-data
///         mountPath: /data
/// ```
///
/// Default path: `/data/sol.db` (configurable via `matrix.db_path` in sol.toml).
/// The `/data` mount also holds the Matrix SDK state store at `/data/matrix-state`.
pub struct Store {
    conn: Mutex<Connection>,
}

impl Store {
    /// Open or create the database at the given path.
    /// Creates tables if they don't exist.
    pub fn open(path: &str) -> anyhow::Result<Self> {
        // Ensure parent directory exists
        if let Some(parent) = Path::new(path).parent() {
            std::fs::create_dir_all(parent)?;
        }

        let conn = Connection::open(path)?;

        // Enable WAL mode for better concurrent read performance
        conn.execute_batch("PRAGMA journal_mode=WAL;")?;

        conn.execute_batch(
            "CREATE TABLE IF NOT EXISTS conversations (
                room_id         TEXT PRIMARY KEY,
                conversation_id TEXT NOT NULL,
                estimated_tokens INTEGER NOT NULL DEFAULT 0,
                created_at      TEXT NOT NULL DEFAULT (datetime('now'))
            );

            CREATE TABLE IF NOT EXISTS agents (
                name       TEXT PRIMARY KEY,
                agent_id   TEXT NOT NULL,
                model      TEXT NOT NULL DEFAULT '',
                created_at TEXT NOT NULL DEFAULT (datetime('now'))
            );",
        )?;

        info!(path, "Opened Sol state database");
        Ok(Self {
            conn: Mutex::new(conn),
        })
    }

    /// Open an in-memory database (for tests).
    pub fn open_memory() -> anyhow::Result<Self> {
        Self::open(":memory:")
    }

    // =========================================================================
    // Conversations
    // =========================================================================

    /// Get the conversation_id for a room, if one exists.
    pub fn get_conversation(&self, room_id: &str) -> Option<(String, u32)> {
        let conn = self.conn.lock().unwrap();
        conn.query_row(
            "SELECT conversation_id, estimated_tokens FROM conversations WHERE room_id = ?1",
            params![room_id],
            |row| Ok((row.get::<_, String>(0)?, row.get::<_, u32>(1)?)),
        )
        .ok()
    }

    /// Store or update a conversation mapping.
    pub fn upsert_conversation(
        &self,
        room_id: &str,
        conversation_id: &str,
        estimated_tokens: u32,
    ) {
        let conn = self.conn.lock().unwrap();
        if let Err(e) = conn.execute(
            "INSERT INTO conversations (room_id, conversation_id, estimated_tokens)
             VALUES (?1, ?2, ?3)
             ON CONFLICT(room_id) DO UPDATE SET
                conversation_id = excluded.conversation_id,
                estimated_tokens = excluded.estimated_tokens",
            params![room_id, conversation_id, estimated_tokens],
        ) {
            warn!("Failed to upsert conversation: {e}");
        }
    }

    /// Update token estimate for a conversation.
    pub fn update_tokens(&self, room_id: &str, estimated_tokens: u32) {
        let conn = self.conn.lock().unwrap();
        let _ = conn.execute(
            "UPDATE conversations SET estimated_tokens = ?1 WHERE room_id = ?2",
            params![estimated_tokens, room_id],
        );
    }

    /// Remove a conversation mapping (e.g., after compaction).
    pub fn delete_conversation(&self, room_id: &str) {
        let conn = self.conn.lock().unwrap();
        let _ = conn.execute(
            "DELETE FROM conversations WHERE room_id = ?1",
            params![room_id],
        );
    }

    /// Load all conversation mappings (for startup recovery).
    pub fn load_all_conversations(&self) -> Vec<(String, String, u32)> {
        let conn = self.conn.lock().unwrap();
        let mut stmt = match conn.prepare(
            "SELECT room_id, conversation_id, estimated_tokens FROM conversations",
        ) {
            Ok(s) => s,
            Err(_) => return Vec::new(),
        };

        stmt.query_map([], |row| {
            Ok((
                row.get::<_, String>(0)?,
                row.get::<_, String>(1)?,
                row.get::<_, u32>(2)?,
            ))
        })
        .ok()
        .map(|rows| rows.filter_map(|r| r.ok()).collect())
        .unwrap_or_default()
    }

    // =========================================================================
    // Agents
    // =========================================================================

    /// Get the agent_id for a named agent.
    pub fn get_agent(&self, name: &str) -> Option<String> {
        let conn = self.conn.lock().unwrap();
        conn.query_row(
            "SELECT agent_id FROM agents WHERE name = ?1",
            params![name],
            |row| row.get(0),
        )
        .ok()
    }

    /// Store or update an agent mapping.
    pub fn upsert_agent(&self, name: &str, agent_id: &str, model: &str) {
        let conn = self.conn.lock().unwrap();
        if let Err(e) = conn.execute(
            "INSERT INTO agents (name, agent_id, model)
             VALUES (?1, ?2, ?3)
             ON CONFLICT(name) DO UPDATE SET
                agent_id = excluded.agent_id,
                model = excluded.model",
            params![name, agent_id, model],
        ) {
            warn!("Failed to upsert agent: {e}");
        }
    }

    /// Remove an agent mapping.
    pub fn delete_agent(&self, name: &str) {
        let conn = self.conn.lock().unwrap();
        let _ = conn.execute("DELETE FROM agents WHERE name = ?1", params![name]);
    }

    /// Load all agent mappings (for startup recovery).
    pub fn load_all_agents(&self) -> Vec<(String, String)> {
        let conn = self.conn.lock().unwrap();
        let mut stmt = match conn.prepare("SELECT name, agent_id FROM agents") {
            Ok(s) => s,
            Err(_) => return Vec::new(),
        };

        stmt.query_map([], |row| {
            Ok((row.get::<_, String>(0)?, row.get::<_, String>(1)?))
        })
        .ok()
        .map(|rows| rows.filter_map(|r| r.ok()).collect())
        .unwrap_or_default()
    }
}

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

    #[test]
    fn test_open_memory_db() {
        let store = Store::open_memory().unwrap();
        assert!(store.load_all_conversations().is_empty());
        assert!(store.load_all_agents().is_empty());
    }

    #[test]
    fn test_conversation_crud() {
        let store = Store::open_memory().unwrap();

        // Insert
        store.upsert_conversation("!room:x", "conv_abc", 100);
        let (conv_id, tokens) = store.get_conversation("!room:x").unwrap();
        assert_eq!(conv_id, "conv_abc");
        assert_eq!(tokens, 100);

        // Update tokens
        store.update_tokens("!room:x", 500);
        let (_, tokens) = store.get_conversation("!room:x").unwrap();
        assert_eq!(tokens, 500);

        // Upsert (replace conversation_id)
        store.upsert_conversation("!room:x", "conv_def", 0);
        let (conv_id, tokens) = store.get_conversation("!room:x").unwrap();
        assert_eq!(conv_id, "conv_def");
        assert_eq!(tokens, 0);

        // Delete
        store.delete_conversation("!room:x");
        assert!(store.get_conversation("!room:x").is_none());
    }

    #[test]
    fn test_agent_crud() {
        let store = Store::open_memory().unwrap();

        store.upsert_agent("sol-orchestrator", "ag_123", "mistral-medium-latest");
        assert_eq!(
            store.get_agent("sol-orchestrator").unwrap(),
            "ag_123"
        );

        // Update
        store.upsert_agent("sol-orchestrator", "ag_456", "mistral-medium-latest");
        assert_eq!(
            store.get_agent("sol-orchestrator").unwrap(),
            "ag_456"
        );

        // Delete
        store.delete_agent("sol-orchestrator");
        assert!(store.get_agent("sol-orchestrator").is_none());
    }

    #[test]
    fn test_load_all_conversations() {
        let store = Store::open_memory().unwrap();
        store.upsert_conversation("!a:x", "conv_1", 10);
        store.upsert_conversation("!b:x", "conv_2", 20);
        store.upsert_conversation("!c:x", "conv_3", 30);

        let all = store.load_all_conversations();
        assert_eq!(all.len(), 3);
    }

    #[test]
    fn test_load_all_agents() {
        let store = Store::open_memory().unwrap();
        store.upsert_agent("orch", "ag_1", "medium");
        store.upsert_agent("obs", "ag_2", "medium");

        let all = store.load_all_agents();
        assert_eq!(all.len(), 2);
    }

    #[test]
    fn test_nonexistent_keys_return_none() {
        let store = Store::open_memory().unwrap();
        assert!(store.get_conversation("!nope:x").is_none());
        assert!(store.get_agent("nope").is_none());
    }
}