PicoBot/src/command/handlers/get_current.rs

use crate::agent::context_compressor::estimate_tokens;
use crate::agent::{SystemPromptContext, SystemPromptProvider};
use crate::command::context::CommandContext;
use crate::command::handler::{CommandHandler, CommandMetadata};
use crate::command::handlers::get_messages_from_session;
use crate::command::response::{CommandError, CommandResponse, MessageKind};
use crate::command::Command;
use crate::gateway::session::SessionManager;
use crate::storage::SessionStore;
use async_trait::async_trait;
use std::sync::Arc;

/// 获取当前话题命令处理器
pub struct GetCurrentSessionCommandHandler {
    store: Arc<SessionStore>,
    session_manager: Option<SessionManager>,
    system_prompt_provider: Option<Arc<dyn SystemPromptProvider>>,
}

impl GetCurrentSessionCommandHandler {
    pub fn new(store: Arc<SessionStore>) -> Self {
        Self {
            store,
            session_manager: None,
            system_prompt_provider: None,
        }
    }

    pub fn with_session_manager(mut self, session_manager: SessionManager) -> Self {
        self.session_manager = Some(session_manager);
        self
    }

    pub fn with_system_prompt_provider(mut self, provider: Arc<dyn SystemPromptProvider>) -> Self {
        self.system_prompt_provider = Some(provider);
        self
    }
}

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

    fn metadata(&self) -> Option<CommandMetadata> {
        Some(CommandMetadata {
            name: "current",
            description: "获取当前话题信息",
            usage: "/current",
        })
    }

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

async fn handle_get_current_session(
    handler: &GetCurrentSessionCommandHandler,
    ctx: CommandContext,
) -> Result<CommandResponse, CommandError> {
    let topic_id = ctx.topic_id.as_deref()
        .ok_or_else(|| CommandError::new("NO_CURRENT_TOPIC", "No current topic"))?;

    let chat_id = ctx.chat_id.as_deref()
        .ok_or_else(|| CommandError::new("NO_CHAT_ID", "No chat id".to_string()))?;

    let topic = handler
        .store
        .get_topic(topic_id)
        .map_err(|e| CommandError::new("GET_TOPIC_ERROR", e.to_string()))?
        .ok_or_else(|| CommandError::new("TOPIC_NOT_FOUND", format!("Topic not found: {}", topic_id)))?;

    // Load messages from session memory
    let messages = get_messages_from_session(
        &handler.session_manager,
        &ctx.channel_name,
        chat_id,
    ).await?;

    let actual_message_count = messages.len();
    let message_tokens = estimate_tokens(&messages);

    // Calculate system prompt tokens if provider is available
    let system_prompt_tokens = if let Some(ref provider) = handler.system_prompt_provider {
        let user_message_count = messages.iter().filter(|m| m.role == "user").count();
        let system_prompt_context = SystemPromptContext {
            session_id: ctx.session_id.clone(),
            chat_id: chat_id.to_string(),
            user_message_count,
        };

        provider.build(&system_prompt_context)
            .map(|sp| {
                use crate::bus::ChatMessage;
                let system_msg = ChatMessage::system(&sp.content);
                estimate_tokens(&[system_msg])
            })
            .unwrap_or(0)
    } else {
        0
    };

    let total_tokens = system_prompt_tokens + message_tokens;

    let last_active = format_time_ago(topic.last_active_at);
    let created_at = format_time_ago(topic.created_at);

    let message = format!(
        "Current Topic:\n\n  Topic ID: {}\n  Title: {}\n  Messages: {}\n  Tokens: ~{} (系统提示词: ~{}, 用户消息: ~{})\n  Created: {}\n  Last Active: {}",
        topic.id,
        topic.title,
        actual_message_count,
        total_tokens,
        system_prompt_tokens,
        message_tokens,
        created_at,
        last_active
    );

    Ok(CommandResponse::success(ctx.request_id)
        .with_message(MessageKind::Notification, &message)
        .with_metadata("topic_id", &topic.id)
        .with_metadata("title", &topic.title)
        .with_metadata("message_count", &actual_message_count.to_string())
        .with_metadata("estimated_tokens", &total_tokens.to_string())
        .with_metadata("system_prompt_tokens", &system_prompt_tokens.to_string())
        .with_metadata("message_tokens", &message_tokens.to_string()))
}

fn format_time_ago(timestamp_ms: i64) -> String {
    let now = std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap()
        .as_millis() as i64;

    let diff_ms = now - timestamp_ms;
    let diff_secs = diff_ms / 1000;

    if diff_secs < 60 {
        "just now".to_string()
    } else if diff_secs < 3600 {
        format!("{} mins ago", diff_secs / 60)
    } else if diff_secs < 86400 {
        format!("{} hours ago", diff_secs / 3600)
    } else {
        format!("{} days ago", diff_secs / 86400)
    }
}