iAs/src/context/tools.rs

use sqlx::PgPool;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;

/// ## 长期记忆管理器（内存 + PostgreSQL 双写）
///
/// 为每个用户维护独立的长期记忆列表，支持读写操作。
///
/// ### 存储策略
/// * **内存缓存** — `HashMap<String, Vec<String>>`，按 user_id 隔离
/// * **数据库持久化** — 可选的 `PgPool`，写入时同时写到 PostgreSQL
///
/// ### 数据流
/// ```text
/// 启动时 load(user_id) → 从数据库加载到缓存
/// 对话中 read_for() → 从缓存读取，返回 "1. ...\n2. ..." 格式
/// 对话中 write_for() → 写入缓存 + 写入数据库
/// ```
pub struct MemoryStore {
    pool: Option<Arc<PgPool>>,
    cache: Arc<Mutex<HashMap<String, Vec<String>>>>,
}

impl MemoryStore {
    pub fn new(pool: Option<Arc<PgPool>>) -> Self {
        Self {
            pool,
            cache: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    pub async fn load(&self, user_id: &str) {
        let pool = match self.pool.as_ref() {
            Some(p) => p.clone(),
            None => return,
        };
        let uid = if user_id.is_empty() { "default" } else { user_id };
        if let Ok(rows) = sqlx::query_as::<_, (String,)>(
            "SELECT content FROM user_memories WHERE user_id = $1 ORDER BY id",
        )
        .bind(uid)
        .fetch_all(pool.as_ref())
        .await
        {
            let mems: Vec<String> = rows.into_iter().map(|(c,)| c).collect();
            self.cache.lock().await.insert(uid.to_string(), mems);
        }
    }

    pub async fn read_for(&self, user_id: &str) -> String {
        let uid = if user_id.is_empty() { "default" } else { user_id };
        let cache = self.cache.lock().await;
        match cache.get(uid) {
            Some(m) if !m.is_empty() => m.iter().enumerate()
                .map(|(i, v)| format!("{}. {}", i + 1, v))
                .collect::<Vec<_>>()
                .join("\n"),
            _ => "暂无长期记忆".to_string(),
        }
    }

    pub async fn write_for(&self, user_id: &str, content: &str) -> String {
        let uid = if user_id.is_empty() { "default" } else { user_id };
        self.cache.lock().await.entry(uid.to_string()).or_default().push(content.to_string());
        if let Some(ref pool) = self.pool {
            let _ = sqlx::query("INSERT INTO user_memories (user_id, content) VALUES ($1, $2)")
                .bind(uid)
                .bind(content)
                .execute(pool.as_ref())
                .await;
        }
        "已添加长期记忆".to_string()
    }
}

/// ## `read_summaries` 工具 —— 读取历史摘要（内存 + 数据库）
///
/// 返回当前 session 中的摘要列表 + 数据库中存档的摘要。
/// 去重：数据库中的摘要如果与 session 中的摘要前 30 个字符相同，则跳过。
pub async fn read_summaries(
    session: &Arc<Mutex<super::types::ChatSession>>,
) -> String {
    let s = session.lock().await;
    let mut lines: Vec<String> = Vec::new();

    for sum in &s.summaries {
        let reason = match sum.reason {
            super::types::SummaryReason::Overflow => "上下文压缩",
            super::types::SummaryReason::Timeout => "空闲超时",
        };
        lines.push(format!(
            "[{}] {} (原因: {})",
            sum.created_at.format("%Y-%m-%d %H:%M"),
            sum.text,
            reason,
        ));
    }

    if let Some(pool) = &s.db_pool {
        let uid = if s.current_user_id.is_empty() { "default" } else { &s.current_user_id };
        if let Ok(db_summaries) = crate::db::models::load_summaries(pool, uid, 5).await {
            for text in db_summaries {
                if !lines.iter().any(|l| l.contains(&text.chars().take(30).collect::<String>())) {
                    lines.push(format!("[数据库存档] {}", text));
                }
            }
        }
    }

    if lines.is_empty() { "暂无历史摘要".to_string() } else { lines.join("\n---\n") }
}