noc/src/state.rs

use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use anyhow::Result;
use chrono::NaiveDate;
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::{error, info};

use crate::tools::SubAgent;

// ── persistent state ────────────────────────────────────────────────

#[derive(Serialize, Deserialize, Default)]
pub struct Persistent {
    pub authed: HashMap<i64, NaiveDate>,
    pub known_sessions: HashSet<String>,
}

#[derive(Serialize, Deserialize, Clone, Default)]
pub struct ConversationState {
    pub summary: String,
    pub messages: Vec<serde_json::Value>,
    pub total_messages: usize,
}

pub const MAX_WINDOW: usize = 100;
pub const SLIDE_SIZE: usize = 50;

pub struct AppState {
    pub persist: RwLock<Persistent>,
    pub state_path: PathBuf,
    pub db: tokio::sync::Mutex<rusqlite::Connection>,
    pub agents: RwLock<HashMap<String, Arc<SubAgent>>>,
}

impl AppState {
    pub fn load(path: PathBuf) -> Self {
        let persist = std::fs::read_to_string(&path)
            .ok()
            .and_then(|s| serde_json::from_str(&s).ok())
            .unwrap_or_default();
        info!("loaded state from {}", path.display());

        let db_path = path.parent().unwrap_or(Path::new(".")).join("noc.db");
        let conn = rusqlite::Connection::open(&db_path)
            .unwrap_or_else(|e| panic!("open {}: {e}", db_path.display()));
        conn.execute_batch(
            "CREATE TABLE IF NOT EXISTS conversations (
                session_id TEXT PRIMARY KEY,
                summary TEXT NOT NULL DEFAULT '',
                total_messages INTEGER NOT NULL DEFAULT 0
            );
            CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                session_id TEXT NOT NULL,
                role TEXT NOT NULL,
                content TEXT NOT NULL,
                created_at TEXT NOT NULL DEFAULT (datetime('now', 'localtime'))
            );
            CREATE INDEX IF NOT EXISTS idx_messages_session ON messages(session_id);
            CREATE TABLE IF NOT EXISTS scratch_area (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                content TEXT NOT NULL,
                created_at TEXT NOT NULL DEFAULT (datetime('now'))
            );
            CREATE TABLE IF NOT EXISTS config (
                key TEXT PRIMARY KEY,
                value TEXT NOT NULL DEFAULT '',
                create_time TEXT NOT NULL DEFAULT (datetime('now')),
                update_time TEXT NOT NULL DEFAULT (datetime('now'))
            );
            CREATE TABLE IF NOT EXISTS config_history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                key TEXT NOT NULL,
                value TEXT NOT NULL,
                create_time TEXT NOT NULL,
                update_time TEXT NOT NULL
            );
            CREATE TABLE IF NOT EXISTS memory_slots (
                slot_nr INTEGER PRIMARY KEY CHECK(slot_nr BETWEEN 0 AND 99),
                content TEXT NOT NULL DEFAULT ''
            );
            CREATE TABLE IF NOT EXISTS timers (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                chat_id INTEGER NOT NULL,
                label TEXT NOT NULL,
                schedule TEXT NOT NULL,
                next_fire TEXT NOT NULL,
                enabled INTEGER NOT NULL DEFAULT 1,
                created_at TEXT NOT NULL DEFAULT (datetime('now', 'localtime'))
            );
            CREATE TABLE IF NOT EXISTS inner_state (
                id INTEGER PRIMARY KEY CHECK(id = 1),
                content TEXT NOT NULL DEFAULT ''
            );
            INSERT OR IGNORE INTO inner_state (id, content) VALUES (1, '');",
        )
        .expect("init db schema");

        // migrations
        let _ = conn.execute(
            "ALTER TABLE messages ADD COLUMN created_at TEXT NOT NULL DEFAULT ''",
            [],
        );

        info!("opened db {}", db_path.display());

        Self {
            persist: RwLock::new(persist),
            state_path: path,
            db: tokio::sync::Mutex::new(conn),
            agents: RwLock::new(HashMap::new()),
        }
    }

    pub async fn save(&self) {
        let data = self.persist.read().await;
        if let Ok(json) = serde_json::to_string_pretty(&*data) {
            if let Err(e) = std::fs::write(&self.state_path, json) {
                error!("save state: {e}");
            }
        }
    }

    pub async fn load_conv(&self, sid: &str) -> ConversationState {
        let db = self.db.lock().await;
        let (summary, total) = db
            .query_row(
                "SELECT summary, total_messages FROM conversations WHERE session_id = ?1",
                [sid],
                |row| Ok((row.get::<_, String>(0)?, row.get::<_, usize>(1)?)),
            )
            .unwrap_or_default();

        let mut stmt = db
            .prepare("SELECT role, content, created_at FROM messages WHERE session_id = ?1 ORDER BY id")
            .unwrap();
        let messages: Vec<serde_json::Value> = stmt
            .query_map([sid], |row| {
                let role: String = row.get(0)?;
                let content: String = row.get(1)?;
                let ts: String = row.get(2)?;
                let tagged = if ts.is_empty() {
                    content
                } else {
                    format!("[{ts}] {content}")
                };
                Ok(serde_json::json!({"role": role, "content": tagged}))
            })
            .unwrap()
            .filter_map(|r| r.ok())
            .collect();

        ConversationState {
            summary,
            messages,
            total_messages: total,
        }
    }

    pub async fn push_message(&self, sid: &str, role: &str, content: &str) {
        let db = self.db.lock().await;
        let _ = db.execute(
            "INSERT OR IGNORE INTO conversations (session_id) VALUES (?1)",
            [sid],
        );
        let _ = db.execute(
            "INSERT INTO messages (session_id, role, content, created_at) VALUES (?1, ?2, ?3, datetime('now', 'localtime'))",
            rusqlite::params![sid, role, content],
        );
    }

    pub async fn message_count(&self, sid: &str) -> usize {
        let db = self.db.lock().await;
        db.query_row(
            "SELECT COUNT(*) FROM messages WHERE session_id = ?1",
            [sid],
            |row| row.get(0),
        )
        .unwrap_or(0)
    }

    pub async fn slide_window(&self, sid: &str, new_summary: &str, slide_size: usize) {
        let db = self.db.lock().await;
        let _ = db.execute(
            "DELETE FROM messages WHERE id IN (
                SELECT id FROM messages WHERE session_id = ?1 ORDER BY id LIMIT ?2
            )",
            rusqlite::params![sid, slide_size],
        );
        let _ = db.execute(
            "UPDATE conversations SET summary = ?1, total_messages = total_messages + ?2 \
             WHERE session_id = ?3",
            rusqlite::params![new_summary, slide_size, sid],
        );
    }

    pub async fn get_oldest_messages(&self, sid: &str, count: usize) -> Vec<serde_json::Value> {
        let db = self.db.lock().await;
        let mut stmt = db
            .prepare(
                "SELECT role, content FROM messages WHERE session_id = ?1 ORDER BY id LIMIT ?2",
            )
            .unwrap();
        stmt.query_map(rusqlite::params![sid, count], |row| {
            let role: String = row.get(0)?;
            let content: String = row.get(1)?;
            Ok(serde_json::json!({"role": role, "content": content}))
        })
        .unwrap()
        .filter_map(|r| r.ok())
        .collect()
    }

    pub async fn get_scratch(&self) -> String {
        let db = self.db.lock().await;
        db.query_row(
            "SELECT content FROM scratch_area ORDER BY id DESC LIMIT 1",
            [],
            |row| row.get(0),
        )
        .unwrap_or_default()
    }

    pub async fn push_scratch(&self, content: &str) {
        let db = self.db.lock().await;
        let _ = db.execute(
            "INSERT INTO scratch_area (content) VALUES (?1)",
            [content],
        );
    }

    pub async fn get_config(&self, key: &str) -> Option<String> {
        let db = self.db.lock().await;
        db.query_row(
            "SELECT value FROM config WHERE key = ?1",
            [key],
            |row| row.get(0),
        )
        .ok()
    }

    pub async fn get_inner_state(&self) -> String {
        let db = self.db.lock().await;
        db.query_row("SELECT content FROM inner_state WHERE id = 1", [], |row| row.get(0))
            .unwrap_or_default()
    }

    pub async fn set_inner_state(&self, content: &str) {
        let db = self.db.lock().await;
        let _ = db.execute(
            "UPDATE inner_state SET content = ?1 WHERE id = 1",
            [content],
        );
    }

    pub async fn add_timer(&self, chat_id: i64, label: &str, schedule: &str, next_fire: &str) -> i64 {
        let db = self.db.lock().await;
        db.execute(
            "INSERT INTO timers (chat_id, label, schedule, next_fire) VALUES (?1, ?2, ?3, ?4)",
            rusqlite::params![chat_id, label, schedule, next_fire],
        )
        .unwrap();
        db.last_insert_rowid()
    }

    pub async fn list_timers(&self, chat_id: Option<i64>) -> Vec<(i64, i64, String, String, String, bool)> {
        let db = self.db.lock().await;
        let (sql, params): (&str, Vec<Box<dyn rusqlite::types::ToSql>>) = match chat_id {
            Some(cid) => (
                "SELECT id, chat_id, label, schedule, next_fire, enabled FROM timers WHERE chat_id = ?1 ORDER BY next_fire",
                vec![Box::new(cid)],
            ),
            None => (
                "SELECT id, chat_id, label, schedule, next_fire, enabled FROM timers ORDER BY next_fire",
                vec![],
            ),
        };
        let mut stmt = db.prepare(sql).unwrap();
        stmt.query_map(rusqlite::params_from_iter(params), |row| {
            Ok((
                row.get(0)?,
                row.get(1)?,
                row.get::<_, String>(2)?,
                row.get::<_, String>(3)?,
                row.get::<_, String>(4)?,
                row.get::<_, bool>(5)?,
            ))
        })
        .unwrap()
        .filter_map(|r| r.ok())
        .collect()
    }

    pub async fn cancel_timer(&self, timer_id: i64) -> bool {
        let db = self.db.lock().await;
        db.execute("DELETE FROM timers WHERE id = ?1", [timer_id]).unwrap() > 0
    }

    pub async fn due_timers(&self) -> Vec<(i64, i64, String, String)> {
        let db = self.db.lock().await;
        let now = chrono::Local::now().format("%Y-%m-%d %H:%M:%S").to_string();
        let mut stmt = db
            .prepare(
                "SELECT id, chat_id, label, schedule FROM timers WHERE enabled = 1 AND next_fire <= ?1",
            )
            .unwrap();
        stmt.query_map([&now], |row| {
            Ok((
                row.get(0)?,
                row.get(1)?,
                row.get::<_, String>(2)?,
                row.get::<_, String>(3)?,
            ))
        })
        .unwrap()
        .filter_map(|r| r.ok())
        .collect()
    }

    pub async fn update_timer_next_fire(&self, timer_id: i64, next_fire: &str) {
        let db = self.db.lock().await;
        let _ = db.execute(
            "UPDATE timers SET next_fire = ?1 WHERE id = ?2",
            rusqlite::params![next_fire, timer_id],
        );
    }

    pub async fn get_memory_slots(&self) -> Vec<(i32, String)> {
        let db = self.db.lock().await;
        let mut stmt = db
            .prepare("SELECT slot_nr, content FROM memory_slots WHERE content != '' ORDER BY slot_nr")
            .unwrap();
        stmt.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))
            .unwrap()
            .filter_map(|r| r.ok())
            .collect()
    }

    pub async fn set_memory_slot(&self, slot_nr: i32, content: &str) -> Result<()> {
        if !(0..=99).contains(&slot_nr) {
            anyhow::bail!("slot_nr must be 0-99, got {slot_nr}");
        }
        if content.len() > 200 {
            anyhow::bail!("content too long: {} chars (max 200)", content.len());
        }
        let db = self.db.lock().await;
        db.execute(
            "INSERT INTO memory_slots (slot_nr, content) VALUES (?1, ?2) \
             ON CONFLICT(slot_nr) DO UPDATE SET content = ?2",
            rusqlite::params![slot_nr, content],
        )?;
        Ok(())
    }
}