PicoBot/src/gateway/session_history.rs

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

use crate::agent::AgentError;
use crate::bus::ChatMessage;
use crate::storage::{
    ConversationRepository, SessionRecord, SkillEventRepository, persistent_session_id,
};

fn preview_text(content: &str, max_chars: usize) -> String {
    let mut preview = content.chars().take(max_chars).collect::<String>();
    if content.chars().count() > max_chars {
        preview.push_str("...");
    }
    preview.replace('\n', "\\n")
}

pub(crate) struct SessionHistory {
    channel_name: String,
    /// 按 topic_id 键化的内存历史缓存。
    /// 不同 topic 的历史独立存储，互不干扰，支持多话题并发执行。
    topic_histories: HashMap<String, Vec<ChatMessage>>,
    /// UI 状态：每个 chat 当前活跃的 topic（按 chat_id 键）。
    chat_topic_ids: HashMap<String, String>,
    /// 正在压缩中的 topic_id 集合
    compression_in_flight: HashSet<String>,
    /// 按 topic_id 的串行化锁。
    /// 同一 topic 的消息处理（agent loop + 压缩）必须串行执行，
    /// 防止并发 loop 操作同一历史的不同快照产生交错序列。
    /// 不同 topic 之间互不阻塞，支持多话题并发执行。
    topic_serial_locks: HashMap<String, Arc<tokio::sync::Mutex<()>>>,
    conversations: Arc<dyn ConversationRepository>,
    skill_events: Arc<dyn SkillEventRepository>,
}

impl SessionHistory {
    pub(crate) fn new(
        channel_name: impl Into<String>,
        conversations: Arc<dyn ConversationRepository>,
        skill_events: Arc<dyn SkillEventRepository>,
    ) -> Self {
        Self {
            channel_name: channel_name.into(),
            topic_histories: HashMap::new(),
            chat_topic_ids: HashMap::new(),
            compression_in_flight: HashSet::new(),
            topic_serial_locks: HashMap::new(),
            conversations,
            skill_events,
        }
    }

    /// 获取或创建该 topic 的串行化锁。
    /// 同一 topic 的所有消息处理共享同一个锁，保证串行执行；
    /// 不同 topic 之间互不阻塞，支持多话题并发执行。
    pub(crate) fn topic_serial_lock(&mut self, topic_id: &str) -> Arc<tokio::sync::Mutex<()>> {
        self.topic_serial_locks
            .entry(topic_id.to_string())
            .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(())))
            .clone()
    }

    pub(crate) fn persistent_session_id(&self, chat_id: &str) -> String {
        persistent_session_id(&self.channel_name, chat_id)
    }

    pub(crate) fn ensure_persistent_session(
        &self,
        chat_id: &str,
    ) -> Result<SessionRecord, AgentError> {
        self.conversations
            .ensure_channel_session(&self.channel_name, chat_id)
            .map_err(|err| AgentError::Other(format!("session persistence error: {}", err)))
    }

    /// 确保指定 topic 的历史已加载到内存。
    /// 按 topic_id 键化查找，如果已存在则直接返回，否则从 DB 加载。
    pub(crate) fn ensure_chat_loaded(
        &mut self,
        chat_id: &str,
        topic_id: Option<&str>,
    ) -> Result<(), AgentError> {
        let Some(tid) = topic_id else {
            return Ok(());
        };
        if self.topic_histories.contains_key(tid) {
            return Ok(());
        }

        let sid = self.persistent_session_id(chat_id);
        let mut history = self
            .conversations
            .load_messages_for_topic(tid, Some(&sid))
            .map_err(|err| AgentError::Other(format!("session history load error: {}", err)))?;

        let removed = crate::bus::message::sanitize_incomplete_tool_call_sequences(&mut history);
        if removed > 0 {
            tracing::warn!(
                topic_id = %tid,
                removed_count = removed,
                "Sanitized incomplete tool_call sequences on history load"
            );
        }

        self.topic_histories.insert(tid.to_string(), history);
        Ok(())
    }

    pub(crate) fn ensure_agent_prompt_before_user_message(
        &mut self,
        _chat_id: &str,
    ) -> Result<(), AgentError> {
        // 提示词现在由 AgentPromptProvider 统一处理，不需要在此处注入
        Ok(())
    }

    pub(crate) fn get_or_create_history(&mut self, topic_id: &str) -> &mut Vec<ChatMessage> {
        self.topic_histories
            .entry(topic_id.to_string())
            .or_default()
    }

    pub(crate) fn get_history(&self, topic_id: &str) -> Option<&Vec<ChatMessage>> {
        self.topic_histories.get(topic_id)
    }

    pub(crate) fn set_history(&mut self, topic_id: &str, history: Vec<ChatMessage>) {
        self.topic_histories.insert(topic_id.to_string(), history);
    }

    /// 设置指定 chat 的当前 topic（UI 状态）
    pub(crate) fn set_chat_topic(&mut self, chat_id: &str, topic_id: String) {
        self.chat_topic_ids.insert(chat_id.to_string(), topic_id);
    }

    /// 获取指定 chat 的当前 topic（UI 状态）
    pub(crate) fn chat_topic(&self, chat_id: &str) -> Option<&str> {
        self.chat_topic_ids.get(chat_id).map(|s| s.as_str())
    }

    /// 清除指定 chat 的 topic（UI 状态）
    pub(crate) fn clear_chat_topic(&mut self, chat_id: &str) {
        self.chat_topic_ids.remove(chat_id);
    }

    pub(crate) fn add_message(&mut self, topic_id: &str, message: ChatMessage) {
        self.get_or_create_history(topic_id).push(message);
    }

    pub(crate) fn remove_history(&mut self, topic_id: &str) {
        self.topic_histories.remove(topic_id);
        self.compression_in_flight.remove(topic_id);
    }

    /// 清空指定 chat/topic 的内存历史和 DB 消息。
    /// 内存按 topic_id 清，DB 按 session_id 清（保留原行为以兼容无 topic 场景）。
    pub(crate) fn clear_chat_history(
        &mut self,
        chat_id: &str,
        topic_id: Option<&str>,
    ) -> Result<(), AgentError> {
        if let Some(tid) = topic_id {
            if let Some(history) = self.topic_histories.get_mut(tid) {
                let len = history.len();
                history.clear();
                #[cfg(debug_assertions)]
                tracing::debug!(topic_id = %tid, previous_len = len, "Topic history cleared");
            }
        }

        self.conversations
            .clear_messages(&self.persistent_session_id(chat_id))
            .map_err(|err| AgentError::Other(format!("clear history persistence error: {}", err)))
    }

    pub(crate) fn append_persisted_messages<I>(
        &mut self,
        topic_id: &str,
        messages: I,
    ) -> Result<(), AgentError>
    where
        I: IntoIterator<Item = ChatMessage>,
    {
        let messages: Vec<ChatMessage> = messages.into_iter().collect();
        if messages.is_empty() {
            return Ok(());
        }

        if let Some(history) = self.topic_histories.get_mut(topic_id) {
            let removed = crate::bus::message::sanitize_incomplete_tool_call_sequences(history);
            if removed > 0 {
                tracing::warn!(
                    topic_id = %topic_id,
                    removed_count = removed,
                    "Sanitized in-memory history before appending persisted messages"
                );
            }
        }

        for message in messages {
            self.add_message(topic_id, message);
        }
        Ok(())
    }

    /// 将消息保存到指定话题
    /// 每条消息已通过 PersistingEmittedMessageHandler 逐条持久化，此处仅保留接口兼容
    pub(crate) fn append_to_topic(
        &self,
        _chat_id: &str,
        _topic_id: &str,
        messages: &[ChatMessage],
    ) -> Result<(), AgentError> {
        if messages.is_empty() {
            return Ok(());
        }
        Ok(())
    }

    pub(crate) fn latest_user_message(&self, topic_id: &str) -> Option<&ChatMessage> {
        self.get_history(topic_id)
            .and_then(|history| history.iter().rev().find(|message| message.role == "user"))
    }

    #[cfg(test)]
    pub(crate) fn matches_current_user_turn(&self, topic_id: &str, message: &ChatMessage) -> bool {
        self.latest_user_message(topic_id)
            .map(|current| {
                current.id == message.id
                    || (current.content == message.content
                        && current.timestamp == message.timestamp
                        && current.media_refs == message.media_refs)
            })
            .unwrap_or(false)
    }

    pub(crate) fn stale_result_diagnostics(
        &self,
        topic_id: &str,
    ) -> (Option<&str>, Option<String>, bool, usize) {
        let latest_user = self.latest_user_message(topic_id);
        let latest_user_id = latest_user.map(|message| message.id.as_str());
        let latest_user_preview = latest_user.map(|message| preview_text(&message.content, 80));
        let compression_in_flight = self.compression_in_flight.contains(topic_id);
        let history_len = self
            .get_history(topic_id)
            .map(|history| history.len())
            .unwrap_or(0);

        (
            latest_user_id,
            latest_user_preview,
            compression_in_flight,
            history_len,
        )
    }

    /// 清空所有内存历史（主要用于测试全局重置）。
    /// 不遍历清 DB，生产环境如需清 DB 应由调用方显式调用。
    pub(crate) fn clear_all_history(&mut self) -> Result<(), AgentError> {
        let total: usize = self.topic_histories.values().map(|h| h.len()).sum();
        self.topic_histories.clear();
        self.compression_in_flight.clear();
        #[cfg(debug_assertions)]
        tracing::debug!(previous_total = total, "All topic histories cleared");
        Ok(())
    }

    /// 按 topic_id 从 DB 重新加载历史到内存
    pub(crate) fn reload_topic_history(
        &mut self,
        chat_id: &str,
        topic_id: &str,
    ) -> Result<(), AgentError> {
        let sid = self.persistent_session_id(chat_id);
        let history = self
            .conversations
            .load_messages_for_topic(topic_id, Some(&sid))
            .map_err(|err| AgentError::Other(format!("session history reload error: {}", err)))?;
        self.topic_histories.insert(topic_id.to_string(), history);
        Ok(())
    }

    pub(crate) fn conversations(&self) -> Arc<dyn ConversationRepository> {
        self.conversations.clone()
    }

    pub(crate) fn append_skill_event(
        &self,
        chat_id: &str,
        event_type: &str,
        skill_name: Option<&str>,
        payload: &serde_json::Value,
    ) -> Result<(), AgentError> {
        self.skill_events
            .append_skill_event(
                Some(&self.persistent_session_id(chat_id)),
                event_type,
                skill_name,
                payload,
            )
            .map_err(|err| AgentError::Other(format!("append skill event error: {}", err)))
    }
}