crabmate 0.5.0

//! Web `conversation_id` 持久化：可选 **SQLite**（进程重启后续聊）；未配置路径时仍用内存 `HashMap`（`web::app_state`）。

use std::path::Path;
use std::time::{SystemTime, UNIX_EPOCH};

use rusqlite::{Connection, OptionalExtension, params};
use serde_json::json;

use crate::cm_api_contract::chat::ConversationLayoutMeta;
use crate::types::Message;

const TABLE: &str = "crabmate_conversations";

/// Web 会话保存结果（内存与 SQLite 共用）。
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum SaveConversationOutcome {
    Saved,
    Conflict,
}

/// 与内存态 `app_state` 一致：24h TTL、最多条数（仅 SQLite 路径下 `prune` 使用）。
pub const CONVERSATION_STORE_TTL_SECS: u64 = 24 * 3600;
pub const CONVERSATION_STORE_MAX_ENTRIES: usize = 512;

fn now_unix() -> i64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs() as i64
}

/// 建表（幂等）。
pub fn migrate(conn: &Connection) -> Result<(), rusqlite::Error> {
    conn.execute_batch(&format!(
        r#"
        CREATE TABLE IF NOT EXISTS {TABLE} (
            id TEXT NOT NULL PRIMARY KEY,
            messages_json TEXT NOT NULL,
            revision INTEGER NOT NULL,
            updated_at_unix INTEGER NOT NULL
        );
        CREATE INDEX IF NOT EXISTS idx_{TABLE}_updated ON {TABLE}(updated_at_unix);
        "#
    ))?;
    ensure_active_agent_role_column(conn)?;
    ensure_active_session_mode_column(conn)?;
    ensure_layout_meta_json_column(conn)?;
    Ok(())
}

fn ensure_layout_meta_json_column(conn: &Connection) -> Result<(), rusqlite::Error> {
    ensure_text_column(conn, "layout_meta_json")
}

fn ensure_active_agent_role_column(conn: &Connection) -> Result<(), rusqlite::Error> {
    ensure_text_column(conn, "active_agent_role")
}

fn ensure_active_session_mode_column(conn: &Connection) -> Result<(), rusqlite::Error> {
    ensure_text_column(conn, "active_session_mode")
}

fn ensure_text_column(conn: &Connection, column: &str) -> Result<(), rusqlite::Error> {
    let mut stmt = conn.prepare(&format!("PRAGMA table_info({TABLE})"))?;
    let mut has = false;
    let mut rows = stmt.query([])?;
    while let Some(row) = rows.next()? {
        let name: String = row.get(1)?;
        if name == column {
            has = true;
            break;
        }
    }
    if !has {
        conn.execute(
            &format!("ALTER TABLE {TABLE} ADD COLUMN {column} TEXT NOT NULL DEFAULT ''"),
            [],
        )?;
    }
    Ok(())
}

fn layout_meta_json_for_messages(conversation_id: &str, messages: &[Message]) -> String {
    match serde_json::to_string(&crate::cm_turn_layout::layout_meta_from_messages(messages)) {
        Ok(s) => s,
        Err(e) => {
            log::error!(
                target: "crabmate",
                "会话 layout 元数据序列化失败 conversation_id={} error={e}",
                conversation_id
            );
            String::new()
        }
    }
}

/// 打开库文件并迁移（父目录不存在则创建）。
pub fn open_file(path: &Path) -> Result<Connection, Box<dyn std::error::Error + Send + Sync>> {
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)
            .map_err(|e| format!("无法创建会话库目录 {}: {}", parent.display(), e))?;
    }
    let conn = Connection::open(path)
        .map_err(|e| format!("无法打开会话 SQLite {}: {}", path.display(), e))?;
    migrate(&conn).map_err(|e| format!("会话库 schema 初始化失败 {}: {}", path.display(), e))?;
    Ok(conn)
}

fn messages_to_json(messages: &[Message]) -> Result<String, serde_json::Error> {
    serde_json::to_string(&json!({ "version": 1, "messages": messages }))
}

fn messages_from_json(s: &str) -> Result<Vec<Message>, String> {
    let v: serde_json::Value =
        serde_json::from_str(s).map_err(|e| format!("会话 messages_json 解析失败: {e}"))?;
    let arr = v
        .get("messages")
        .and_then(|m| m.as_array())
        .ok_or_else(|| "会话 JSON 缺少 messages 数组".to_string())?;
    let mut out = Vec::with_capacity(arr.len());
    for item in arr {
        let m: Message = serde_json::from_value(item.clone())
            .map_err(|e| format!("单条消息反序列化失败: {e}"))?;
        out.push(m);
    }
    Ok(out)
}

/// SQLite / 内存会话读出形状。
#[derive(Debug, Clone)]
pub struct LoadedConversation {
    pub messages: Vec<Message>,
    pub revision: u64,
    pub active_agent_role: String,
    pub active_session_mode: String,
    pub layout: Option<ConversationLayoutMeta>,
}

struct ConversationFetchRow {
    json: String,
    revision: i64,
    updated: i64,
    active_role: String,
    active_mode: String,
    layout_meta_json: String,
}

fn layout_meta_from_column(conversation_id: &str, raw: &str) -> Option<ConversationLayoutMeta> {
    let t = raw.trim();
    if t.is_empty() {
        return None;
    }
    match serde_json::from_str(t) {
        Ok(meta) => Some(meta),
        Err(e) => {
            log::warn!(
                target: "crabmate",
                "会话 layout_meta_json 解析失败，已忽略 conversation_id={} error={e}",
                conversation_id
            );
            None
        }
    }
}

/// 读取会话；不存在、或超过 TTL 视为无（并删除过期行）。
/// `active_*` 空串视为未设置；`layout` 仅在 `layout_meta_json` 非空且可解析时存在。
pub fn load(
    conn: &Connection,
    id: &str,
    ttl_secs: u64,
) -> Result<Option<LoadedConversation>, rusqlite::Error> {
    let now = now_unix();
    let row: Option<ConversationFetchRow> = conn
        .query_row(
            &format!(
                "SELECT messages_json, revision, updated_at_unix, active_agent_role, active_session_mode, layout_meta_json FROM {TABLE} WHERE id = ?1"
            ),
            params![id],
            |r| {
                Ok(ConversationFetchRow {
                    json: r.get(0)?,
                    revision: r.get(1)?,
                    updated: r.get(2)?,
                    active_role: r.get(3)?,
                    active_mode: r.get(4)?,
                    layout_meta_json: r.get(5)?,
                })
            },
        )
        .optional()?;
    let Some(row) = row else {
        return Ok(None);
    };
    load_row_after_fetch(conn, id, ttl_secs, now, row)
}

fn load_row_after_fetch(
    conn: &Connection,
    id: &str,
    ttl_secs: u64,
    now: i64,
    row: ConversationFetchRow,
) -> Result<Option<LoadedConversation>, rusqlite::Error> {
    if ttl_secs > 0 && now.saturating_sub(row.updated) > ttl_secs as i64 {
        conn.execute(&format!("DELETE FROM {TABLE} WHERE id = ?1"), params![id])?;
        return Ok(None);
    }
    let messages = match messages_from_json(&row.json) {
        Ok(m) => m,
        Err(e) => {
            log::warn!(
                target: "crabmate",
                "会话 {} 消息 JSON 损坏，已删除该行 error={}",
                id,
                e
            );
            conn.execute(&format!("DELETE FROM {TABLE} WHERE id = ?1"), params![id])?;
            return Ok(None);
        }
    };
    let rev = u64::try_from(row.revision).unwrap_or(0);
    conn.execute(
        &format!("UPDATE {TABLE} SET updated_at_unix = ?1 WHERE id = ?2"),
        params![now, id],
    )?;
    Ok(Some(LoadedConversation {
        messages,
        revision: rev,
        active_agent_role: row.active_role,
        active_session_mode: row.active_mode,
        layout: layout_meta_from_column(id, &row.layout_meta_json),
    }))
}

/// 与 `AppState::save_conversation_messages_if_revision` 语义一致。
pub fn save_if_revision(
    conn: &Connection,
    id: &str,
    messages: Vec<Message>,
    active_agent_role: Option<&str>,
    active_session_mode: Option<&str>,
    expected_revision: Option<u64>,
) -> Result<SaveConversationOutcome, rusqlite::Error> {
    let now = now_unix();
    let active_col = active_agent_role
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .unwrap_or("");
    let mode_col = active_session_mode
        .map(str::trim)
        .filter(|s| !s.is_empty())
        .unwrap_or("");
    let json = match messages_to_json(&messages) {
        Ok(j) => j,
        Err(e) => {
            log::error!(
                target: "crabmate",
                "会话 {} 序列化失败（不应发生）: {}",
                id,
                e
            );
            return Ok(SaveConversationOutcome::Conflict);
        }
    };
    let layout_json = layout_meta_json_for_messages(id, &messages);

    if let Some(exp) = expected_revision {
        let n = conn.execute(
            &format!(
                "UPDATE {TABLE} SET messages_json = ?1, active_agent_role = ?2, active_session_mode = ?3, layout_meta_json = ?4, revision = revision + 1, updated_at_unix = ?5 WHERE id = ?6 AND revision = ?7"
            ),
            params![json, active_col, mode_col, layout_json, now, id, exp as i64],
        )?;
        if n == 0 {
            return Ok(SaveConversationOutcome::Conflict);
        }
    } else {
        let exists: i64 = conn.query_row(
            &format!("SELECT COUNT(*) FROM {TABLE} WHERE id = ?1"),
            params![id],
            |r| r.get(0),
        )?;
        if exists > 0 {
            return Ok(SaveConversationOutcome::Conflict);
        }
        conn.execute(
            &format!(
                "INSERT INTO {TABLE} (id, messages_json, active_agent_role, active_session_mode, layout_meta_json, revision, updated_at_unix) VALUES (?1, ?2, ?3, ?4, ?5, 1, ?6)"
            ),
            params![id, json, active_col, mode_col, layout_json, now],
        )?;
    }
    prune(
        conn,
        CONVERSATION_STORE_TTL_SECS,
        CONVERSATION_STORE_MAX_ENTRIES,
    )?;
    Ok(SaveConversationOutcome::Saved)
}

/// 按 id 删除一行（E2E 夹具 `replace` 等；不存在时静默成功）。
pub fn delete_by_id(conn: &Connection, id: &str) -> Result<(), rusqlite::Error> {
    conn.execute(&format!("DELETE FROM {TABLE} WHERE id = ?1"), params![id])?;
    Ok(())
}

/// 删除过期行并按条数淘汰最旧记录。
pub fn prune(conn: &Connection, ttl_secs: u64, max_entries: usize) -> Result<(), rusqlite::Error> {
    let now = now_unix();
    if ttl_secs > 0 {
        let cutoff = now - ttl_secs as i64;
        conn.execute(
            &format!("DELETE FROM {TABLE} WHERE updated_at_unix < ?1"),
            params![cutoff],
        )?;
    }
    if max_entries == 0 {
        return Ok(());
    }
    let count: i64 = conn.query_row(&format!("SELECT COUNT(*) FROM {TABLE}"), [], |r| r.get(0))?;
    if count <= max_entries as i64 {
        return Ok(());
    }
    let to_drop = count - max_entries as i64;
    conn.execute(
        &format!(
            "DELETE FROM {TABLE} WHERE id IN (SELECT id FROM {TABLE} ORDER BY updated_at_unix ASC LIMIT ?1)"
        ),
        params![to_drop],
    )?;
    Ok(())
}

/// 当前库中的会话条数。
pub fn count(conn: &Connection) -> Result<usize, rusqlite::Error> {
    let n: i64 = conn.query_row(&format!("SELECT COUNT(*) FROM {TABLE}"), [], |r| r.get(0))?;
    Ok(n as usize)
}

/// 全部会话的 `messages_json` 原文（供扫描 `/uploads/` 引用，避免清掉仍在用的附图）。
pub fn list_all_messages_json(conn: &Connection) -> Result<Vec<String>, rusqlite::Error> {
    let mut stmt = conn.prepare(&format!("SELECT messages_json FROM {TABLE}"))?;
    let rows = stmt.query_map([], |r| r.get::<_, String>(0))?;
    let mut out = Vec::new();
    for row in rows {
        out.push(row?);
    }
    Ok(out)
}

/// 与 Web `normalize_client_conversation_id` 字符集一致。
pub const CONVERSATION_ID_MAX_LEN: usize = 128;

fn load_messages_at_revision(
    conn: &Connection,
    id: &str,
    expected_revision: u64,
    corrupt_log: &'static str,
) -> Result<Option<Vec<Message>>, rusqlite::Error> {
    let row: Option<(String, i64)> = conn
        .query_row(
            &format!("SELECT messages_json, revision FROM {TABLE} WHERE id = ?1"),
            params![id],
            |r| Ok((r.get(0)?, r.get(1)?)),
        )
        .optional()?;
    let Some((json, revision)) = row else {
        return Ok(None);
    };
    let rev = u64::try_from(revision).unwrap_or(0);
    if rev != expected_revision {
        return Ok(None);
    }
    match messages_from_json(&json) {
        Ok(m) => Ok(Some(m)),
        Err(e) => {
            log::warn!(
                target: "crabmate",
                "{} id={} error={}",
                corrupt_log,
                id,
                e
            );
            Ok(None)
        }
    }
}

fn persist_messages_bump_revision(
    conn: &Connection,
    id: &str,
    expected_revision: u64,
    messages: &[Message],
    serialize_fail_log: &'static str,
) -> Result<SaveConversationOutcome, rusqlite::Error> {
    let new_json = match messages_to_json(messages) {
        Ok(j) => j,
        Err(e) => {
            log::error!(
                target: "crabmate",
                "{} id={} error={}",
                serialize_fail_log,
                id,
                e
            );
            return Ok(SaveConversationOutcome::Conflict);
        }
    };
    let now = now_unix();
    let layout_json = layout_meta_json_for_messages(id, messages);
    let n = conn.execute(
        &format!(
            "UPDATE {TABLE} SET messages_json = ?1, layout_meta_json = ?2, revision = revision + 1, updated_at_unix = ?3 WHERE id = ?4 AND revision = ?5"
        ),
        params![new_json, layout_json, now, id, expected_revision as i64],
    )?;
    if n == 0 {
        return Ok(SaveConversationOutcome::Conflict);
    }
    prune(
        conn,
        CONVERSATION_STORE_TTL_SECS,
        CONVERSATION_STORE_MAX_ENTRIES,
    )?;
    Ok(SaveConversationOutcome::Saved)
}

fn update_messages_json_if_revision(
    conn: &Connection,
    id: &str,
    expected_revision: u64,
    corrupt_log: &'static str,
    serialize_fail_log: &'static str,
    mut mutate: impl FnMut(&mut Vec<Message>) -> bool,
) -> Result<SaveConversationOutcome, rusqlite::Error> {
    let Some(mut messages) = load_messages_at_revision(conn, id, expected_revision, corrupt_log)?
    else {
        return Ok(SaveConversationOutcome::Conflict);
    };
    if !mutate(&mut messages) {
        return Ok(SaveConversationOutcome::Saved);
    }
    persist_messages_bump_revision(conn, id, expected_revision, &messages, serialize_fail_log)
}

/// 截断到「第 `ordinal` 条用户消息」之前（`ordinal` 为 0-based：0 表示删掉从首条用户起的尾部，仅保留 system 等）。
/// 计数不含长期记忆 / 变更集 / 首轮工作区画像 / 上下文摘要等注入条（见 [`crate::types::user_message_counts_for_branch_truncation`]）。
pub fn truncate_before_user_ordinal_if_revision(
    conn: &Connection,
    id: &str,
    user_ordinal: usize,
    expected_revision: u64,
) -> Result<SaveConversationOutcome, rusqlite::Error> {
    update_messages_json_if_revision(
        conn,
        id,
        expected_revision,
        "truncate_before_user 会话 JSON 损坏",
        "truncate_before_user 序列化失败",
        |messages| {
            let mut u = 0usize;
            let mut cut = messages.len();
            for (i, m) in messages.iter().enumerate() {
                if crate::types::user_message_counts_for_branch_truncation(m) {
                    if u == user_ordinal {
                        cut = i;
                        break;
                    }
                    u += 1;
                }
            }
            if cut >= messages.len() {
                return false;
            }
            messages.truncate(cut);
            true
        },
    )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::{Message, message_content_as_str};

    #[test]
    fn truncate_before_user_ordinal_skips_first_turn_workspace_injection() {
        let conn = Connection::open_in_memory().unwrap();
        migrate(&conn).unwrap();
        let msgs = vec![
            Message::system_only("s".to_string()),
            Message::user_first_turn_workspace_context("ctx".to_string()),
            Message::user_only("hi".to_string()),
        ];
        assert_eq!(
            save_if_revision(&conn, "c1", msgs, None, None, None).unwrap(),
            SaveConversationOutcome::Saved
        );
        let loaded = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert_eq!(loaded.messages.len(), 3);
        assert_eq!(
            truncate_before_user_ordinal_if_revision(&conn, "c1", 0, loaded.revision).unwrap(),
            SaveConversationOutcome::Saved
        );
        let after = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert_eq!(after.messages.len(), 2);
        assert_eq!(after.messages[0].role, "system");
        assert_eq!(
            message_content_as_str(&after.messages[1].content),
            Some("ctx")
        );
        assert_eq!(
            after.messages[1].name.as_deref(),
            Some(crate::types::CRABMATE_FIRST_TURN_WORKSPACE_CONTEXT_NAME)
        );
    }

    #[test]
    fn save_load_roundtrip() {
        let conn = Connection::open_in_memory().unwrap();
        migrate(&conn).unwrap();
        let msgs = vec![
            Message::system_only("s".to_string()),
            Message::user_only("hi".to_string()),
        ];
        assert_eq!(
            save_if_revision(&conn, "c1", msgs.clone(), None, None, None).unwrap(),
            SaveConversationOutcome::Saved
        );
        let loaded = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert_eq!(loaded.revision, 1);
        assert_eq!(loaded.messages.len(), 2);
        assert_eq!(
            loaded.layout.as_ref().map(|m| m.layout_schema_version),
            Some(2)
        );
        assert_eq!(
            save_if_revision(&conn, "c1", msgs.clone(), None, None, Some(1)).unwrap(),
            SaveConversationOutcome::Saved
        );
        let loaded2 = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert_eq!(loaded2.revision, 2);
    }

    #[test]
    fn save_load_active_agent_role_roundtrip() {
        let conn = Connection::open_in_memory().unwrap();
        migrate(&conn).unwrap();
        let msgs = vec![Message::system_only("s".to_string())];
        assert_eq!(
            save_if_revision(&conn, "c1", msgs.clone(), None, None, None).unwrap(),
            SaveConversationOutcome::Saved
        );
        let loaded = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert_eq!(loaded.active_agent_role, "");

        assert_eq!(
            save_if_revision(&conn, "c1", msgs.clone(), Some("reviewer"), None, Some(1)).unwrap(),
            SaveConversationOutcome::Saved
        );
        let loaded2 = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert_eq!(loaded2.active_agent_role, "reviewer");

        assert_eq!(
            save_if_revision(&conn, "c1", msgs.clone(), None, None, Some(2)).unwrap(),
            SaveConversationOutcome::Saved
        );
        let loaded3 = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert_eq!(loaded3.active_agent_role, "");
    }

    #[test]
    fn save_load_active_session_mode_roundtrip() {
        let conn = Connection::open_in_memory().unwrap();
        migrate(&conn).unwrap();
        let msgs = vec![Message::system_only("s".to_string())];
        assert_eq!(
            save_if_revision(&conn, "c1", msgs.clone(), None, Some("ask"), None).unwrap(),
            SaveConversationOutcome::Saved
        );
        let loaded = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert_eq!(loaded.active_session_mode, "ask");

        assert_eq!(
            save_if_revision(&conn, "c1", msgs.clone(), None, Some("plan"), Some(1)).unwrap(),
            SaveConversationOutcome::Saved
        );
        let loaded2 = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert_eq!(loaded2.active_session_mode, "plan");

        assert_eq!(
            save_if_revision(&conn, "c1", msgs.clone(), None, None, Some(2)).unwrap(),
            SaveConversationOutcome::Saved
        );
        let loaded3 = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert_eq!(loaded3.active_session_mode, "");
    }

    #[test]
    fn load_optional_layout_meta_json_without_touching_save_path() {
        let conn = Connection::open_in_memory().unwrap();
        migrate(&conn).unwrap();
        let msgs = vec![Message::user_only("hi".to_string())];
        assert_eq!(
            save_if_revision(&conn, "c1", msgs, None, None, None).unwrap(),
            SaveConversationOutcome::Saved
        );
        let meta = ConversationLayoutMeta {
            layout_schema_version:
                crate::cm_api_contract::chat::CONVERSATION_LAYOUT_SCHEMA_VERSION_V2,
            projection_hash: Some("abc".into()),
            segments: vec![],
        };
        let json = serde_json::to_string(&meta).unwrap();
        conn.execute(
            "UPDATE crabmate_conversations SET layout_meta_json = ?1 WHERE id = ?2",
            params![json, "c1"],
        )
        .unwrap();
        let loaded = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert_eq!(loaded.layout.as_ref(), Some(&meta));

        conn.execute(
            "UPDATE crabmate_conversations SET layout_meta_json = ?1 WHERE id = ?2",
            params!["{not-json", "c1"],
        )
        .unwrap();
        let loaded_bad = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert!(loaded_bad.layout.is_none());
        assert_eq!(loaded_bad.messages.len(), 1);
    }

    #[test]
    fn save_if_revision_rewrites_layout_from_messages() {
        let conn = Connection::open_in_memory().unwrap();
        migrate(&conn).unwrap();
        let msgs = vec![
            Message::user_only("hi".to_string()),
            Message::assistant_only("ok".to_string()),
        ];
        assert_eq!(
            save_if_revision(&conn, "c1", msgs.clone(), None, None, None).unwrap(),
            SaveConversationOutcome::Saved
        );
        conn.execute(
            "UPDATE crabmate_conversations SET layout_meta_json = ?1 WHERE id = ?2",
            params![
                r#"{"layout_schema_version":2,"projection_hash":"stale"}"#,
                "c1"
            ],
        )
        .unwrap();
        assert_eq!(
            save_if_revision(&conn, "c1", msgs.clone(), None, None, Some(1)).unwrap(),
            SaveConversationOutcome::Saved
        );
        let loaded = load(&conn, "c1", 3600).unwrap().expect("exists");
        let expected = crate::cm_turn_layout::layout_meta_from_messages(&msgs);
        assert_eq!(loaded.layout.as_ref(), Some(&expected));
        assert_ne!(
            loaded
                .layout
                .as_ref()
                .and_then(|m| m.projection_hash.as_deref()),
            Some("stale")
        );
    }

    #[test]
    fn truncate_before_user_rewrites_layout_meta() {
        let conn = Connection::open_in_memory().unwrap();
        migrate(&conn).unwrap();
        let msgs = vec![
            Message::user_only("hi".to_string()),
            Message::assistant_only("ok".to_string()),
        ];
        assert_eq!(
            save_if_revision(&conn, "c1", msgs, None, None, None).unwrap(),
            SaveConversationOutcome::Saved
        );
        let loaded = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert!(!loaded.layout.as_ref().unwrap().segments.is_empty());
        assert_eq!(
            truncate_before_user_ordinal_if_revision(&conn, "c1", 0, loaded.revision).unwrap(),
            SaveConversationOutcome::Saved
        );
        let after = load(&conn, "c1", 3600).unwrap().expect("exists");
        assert!(after.messages.is_empty());
        assert!(after.layout.as_ref().unwrap().segments.is_empty());
    }
}