PicoBot/src/command/handlers/list_topics.rs

use crate::command::Command;
use crate::command::context::CommandContext;
use crate::command::handler::{CommandHandler, CommandMetadata};
use crate::command::response::{CommandError, CommandResponse, MessageKind};
use crate::storage::{SessionStore, SessionTokenStats, TopicRecord};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;

/// Topic 维度的 token 统计（cost 累计 + context 瞬时）。
///
/// - `prompt_tokens` / `completion_tokens` / `total_tokens`：累计求和（cost 维度）
/// - `last_prompt_tokens`：最后一条 assistant 消息的 prompt_tokens（context 占用瞬时值）
/// - `context_window_tokens`：当前 session 使用的模型上下文窗口上限（来自配置）；
///   为 0 表示未配置，前端不显示百分比。
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TopicTokenStats {
    pub prompt_tokens: u64,
    pub completion_tokens: u64,
    pub total_tokens: u64,
    /// 累计缓存命中的输入 tokens 数（老数据为 0）
    pub cached_tokens: u64,
    pub last_prompt_tokens: Option<u32>,
    pub context_window_tokens: u32,
}

/// Topic 摘要信息
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TopicSummary {
    pub topic_id: String,
    pub session_id: String,
    pub title: String,
    pub description: Option<String>,
    pub message_count: i64,
    pub created_at: i64,
    pub last_active_at: i64,
    /// Token 用量统计。老 topic 或无 LLM 调用时为 None，前端不显示 token 标签。
    #[serde(skip_serializing_if = "Option::is_none")]
    pub token_stats: Option<TopicTokenStats>,
}

/// 构建 TopicSummary 列表的公共函数。
///
/// 一次批量查询所有 topic 的 token 统计，按 `topic_id` 聚合，
/// 避免按 session_id 聚合时同 session 下多个 topic 共享同一总和。
/// 子代理天然隔离：子代理消息的 topic_id 属于子代理自身的 topic，
/// 不在主 topic 列表中。
///
/// `context_window_tokens` 来自最新 assistant 消息记录（LLM 调用时持久化），
/// 无需从配置链路注入，保持存储层与配置解耦。
pub fn build_topic_summaries(
    store: &SessionStore,
    topics: Vec<TopicRecord>,
) -> Result<Vec<TopicSummary>, CommandError> {
    if topics.is_empty() {
        return Ok(Vec::new());
    }

    // 收集所有 topic_id（去重）
    let mut topic_ids: Vec<String> = Vec::new();
    for t in &topics {
        if !topic_ids.contains(&t.id) {
            topic_ids.push(t.id.clone());
        }
    }
    let topic_id_refs: Vec<&str> = topic_ids.iter().map(|s| s.as_str()).collect();

    // 一次批量查询 token 统计（按 topic_id 聚合）
    let stats_map: HashMap<String, SessionTokenStats> = store
        .batch_topic_token_stats(&topic_id_refs)
        .map_err(|e| CommandError::new("TOKEN_STATS_ERROR", e.to_string()))?;

    let summaries = topics
        .into_iter()
        .map(|t| {
            let token_stats = stats_map.get(&t.id).map(|s| TopicTokenStats {
                prompt_tokens: s.prompt_tokens,
                completion_tokens: s.completion_tokens,
                total_tokens: s.total_tokens,
                cached_tokens: s.cached_tokens,
                last_prompt_tokens: s.last_prompt_tokens,
                context_window_tokens: s.context_window_tokens.unwrap_or(0),
            });
            TopicSummary {
                topic_id: t.id,
                session_id: t.session_id,
                title: t.title,
                description: t.description.filter(|d| !d.is_empty()),
                message_count: t.message_count,
                created_at: t.created_at,
                last_active_at: t.last_active_at,
                token_stats,
            }
        })
        .collect();
    Ok(summaries)
}

/// 在 blocking 线程池中执行 list_topics + build_topic_summaries。
///
/// 同步 rusqlite 查询不得直接跑在 tokio worker 上，否则大库查询会饿死
/// 同运行时上的其他任务。list / create / delete / rename 四个话题命令
/// 都返回完整的 TopicSummary 列表供前端刷新侧边栏，统一走本 helper。
pub async fn list_topic_summaries_blocking(
    store: Arc<SessionStore>,
    session_id: &str,
) -> Result<Vec<TopicSummary>, CommandError> {
    let session_id_bg = session_id.to_string();
    tokio::task::spawn_blocking(move || -> Result<Vec<TopicSummary>, CommandError> {
        let topics = store
            .list_topics(&session_id_bg)
            .map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?;
        build_topic_summaries(store.as_ref(), topics)
    })
    .await
    .map_err(|e| CommandError::new("LIST_TOPICS_ERROR", e.to_string()))?
}

/// 列出 Session 的 Topics 命令处理器
pub struct ListTopicsCommandHandler {
    store: Arc<SessionStore>,
}

impl ListTopicsCommandHandler {
    pub fn new(store: Arc<SessionStore>) -> Self {
        Self { store }
    }
}

#[async_trait]
impl CommandHandler for ListTopicsCommandHandler {
    fn can_handle(&self, cmd: &Command) -> bool {
        matches!(cmd, Command::ListTopics { .. })
    }

    fn metadata(&self) -> Option<CommandMetadata> {
        Some(CommandMetadata {
            name: "topics",
            description: "列出 Session 的所有 Topics",
            usage: "/topics <session_id>",
        })
    }

    async fn handle(
        &self,
        cmd: Command,
        ctx: CommandContext,
    ) -> Result<CommandResponse, CommandError> {
        match cmd {
            Command::ListTopics { session_id } => handle_list_topics(self, session_id, ctx).await,
            _ => unreachable!(),
        }
    }
}

async fn handle_list_topics(
    handler: &ListTopicsCommandHandler,
    session_id: String,
    ctx: CommandContext,
) -> Result<CommandResponse, CommandError> {
    let summaries = list_topic_summaries_blocking(handler.store.clone(), &session_id).await?;

    let topics_json = serde_json::to_string(&summaries)
        .map_err(|e| CommandError::new("SERIALIZE_ERROR", e.to_string()))?;

    let message = format!(
        "Found {} topic(s) in session '{}'",
        summaries.len(),
        session_id
    );

    Ok(CommandResponse::success(ctx.request_id)
        .with_message(MessageKind::Notification, &message)
        .with_metadata("topics", &topics_json)
        .with_metadata("session_id", &session_id)
        .with_metadata("count", summaries.len().to_string()))
}