PicoBot/src/gateway/execution.rs

use std::collections::HashMap;
use std::sync::Arc;

use crate::agent::{
    AgentError, AgentProcessResult, EmittedMessageHandler, PersistingEmittedMessageHandler,
    SystemPromptContext,
};
use crate::bus::message::ToolMessageState;
use crate::bus::{ChatMessage, MediaItem, OutboundMessage, SYSTEM_CONTEXT_SCHEDULED_PROMPT};
use crate::config::LLMProviderConfig;
use crate::storage::{ConversationRepository, persistent_session_id};
use async_trait::async_trait;
use tokio::sync::Mutex;

use super::compaction::schedule_background_history_compaction;
use super::message_prepare::enrich_user_content_with_media_refs;
use super::session::Session;

/// 空的 EmittedMessageHandler，不转发消息，仅配合 PersistingEmittedMessageHandler 做持久化。
struct NoOpEmittedMessageHandler;

#[async_trait]
impl EmittedMessageHandler for NoOpEmittedMessageHandler {
    async fn handle(&self, _message: ChatMessage) {}
}

const SCHEDULED_TASK_EXECUTION_SYSTEM_PROMPT: &str = "系统说明：当前输入来自一次已经触发的定时任务执行。你现在需要执行任务内容本身，而不是创建、修改、恢复、暂停或查询新的定时任务。除非当前任务内容明确要求管理调度器，否则不要调用任何定时任务管理工具；像“每小时”、“每天”、“cron”、“定时”等词，只应视为任务背景，不应再解释为新的建任务请求。";

pub(crate) fn compose_scheduled_task_system_prompt(system_prompt: Option<&str>) -> String {
    match system_prompt
        .map(str::trim)
        .filter(|value| !value.is_empty())
    {
        Some(system_prompt) => format!(
            "{}\n\n任务专属要求：{}",
            SCHEDULED_TASK_EXECUTION_SYSTEM_PROMPT, system_prompt
        ),
        None => SCHEDULED_TASK_EXECUTION_SYSTEM_PROMPT.to_string(),
    }
}

pub(crate) struct AgentExecutionService {
    show_tool_results: bool,
}

pub(crate) struct FinalizeAgentResultRequest<'a> {
    pub(crate) channel_name: &'a str,
    pub(crate) chat_id: &'a str,
    pub(crate) user_message: &'a ChatMessage,
    pub(crate) result: AgentProcessResult,
    pub(crate) metadata: &'a HashMap<String, String>,
    pub(crate) suppress_live_tool_calls: bool,
    pub(crate) execution_kind: &'a str,
    pub(crate) original_topic_id: Option<String>,
}

pub(crate) struct FinalizedAgentResult {
    pub(crate) outbound_messages: Vec<OutboundMessage>,
    pub(crate) should_schedule_compaction: bool,
}

pub(crate) struct MessageExecutionRequest<'a> {
    pub(crate) session: Arc<Mutex<Session>>,
    pub(crate) channel_name: &'a str,
    pub(crate) sender_id: &'a str,
    pub(crate) chat_id: &'a str,
    pub(crate) content: &'a str,
    pub(crate) media: Vec<MediaItem>,
    pub(crate) live_emitter: Option<Arc<dyn EmittedMessageHandler>>,
    /// 消息接收时捕获的 topic_id，全程显式传递避免从共享状态重复读取竞态
    pub(crate) topic_id: Option<String>,
}

pub(crate) struct ScheduledExecutionRequest<'a> {
    pub(crate) session: Arc<Mutex<Session>>,
    pub(crate) channel_name: &'a str,
    pub(crate) chat_id: &'a str,
    pub(crate) notification_chat_id: Option<&'a str>,
    pub(crate) prompt: &'a str,
    pub(crate) sender_id: &'a str,
    pub(crate) provider_config: LLMProviderConfig,
    pub(crate) system_prompt: Option<&'a str>,
    pub(crate) metadata: &'a HashMap<String, String>,
    pub(crate) fresh_session: bool,
}

impl AgentExecutionService {
    pub(crate) fn new(show_tool_results: bool) -> Self {
        Self { show_tool_results }
    }

    pub(crate) fn finalize_result(
        &self,
        session: &mut Session,
        request: FinalizeAgentResultRequest<'_>,
    ) -> Result<FinalizedAgentResult, AgentError> {
        // 判断是否是最新的用户回合
        // 直接比较 current_topic(chat_id) 与 original_topic_id
        // 这比"检查内存历史最新消息"更可靠，且天然处理"切走又切回"的 case
        let is_current_turn = match request.original_topic_id.as_deref() {
            Some(orig_tid) => session.current_topic(request.chat_id).as_deref() == Some(orig_tid),
            None => true, // 无 topic 时总是视为当前回合
        };

        if !is_current_turn {
            let (latest_user_id, latest_user_preview, compression_in_flight, history_len) = session
                .stale_result_diagnostics(
                    request
                        .original_topic_id
                        .as_deref()
                        .unwrap_or(request.chat_id),
                );
            tracing::info!(
                channel = %request.channel_name,
                chat_id = %request.chat_id,
                user_message_id = %request.user_message.id,
                latest_user_id,
                latest_user_preview,
                compression_in_flight,
                history_len,
                execution_kind = %request.execution_kind,
                original_topic_id = ?request.original_topic_id,
                "User switched topic during agent execution - saving result to original topic"
            );
        }

        // 确定保存消息的话题 ID
        // 始终使用执行开始时捕获的 original_topic_id，避免从共享状态重复读取竞态
        let target_topic_id = request.original_topic_id.as_deref();

        // 将结果消息保存到确定的话题
        if let Some(topic_id) = target_topic_id {
            if is_current_turn {
                // 话题未切换（current_topic == original_topic_id），安全更新内存历史
                if let Err(err) = session
                    .append_persisted_messages(topic_id, request.result.emitted_messages.clone())
                {
                    tracing::error!(
                        error = %err,
                        topic_id = %topic_id,
                        "Failed to append messages to session history"
                    );
                }
            } else {
                // 话题已切换，只写 DB 不更新内存（避免污染新话题的历史）
                if let Err(err) = session.append_messages_to_topic(
                    request.chat_id,
                    topic_id,
                    &request.result.emitted_messages,
                ) {
                    tracing::error!(
                        error = %err,
                        topic_id = %topic_id,
                        "Failed to append messages to topic"
                    );
                }
            }
        } else if is_current_turn {
            // 没有话题：直接更新内存历史（append_persisted_messages 会处理持久化）
            // 无 topic 场景用 chat_id 作为 topic_histories 的回退 key
            if let Err(err) = session
                .append_persisted_messages(request.chat_id, request.result.emitted_messages.clone())
            {
                tracing::error!(
                    error = %err,
                    chat_id = %request.chat_id,
                    "Failed to append messages to session history"
                );
            }
        }

        // 只有当是最新回合时才发送 outbound 消息给用户
        // 如果用户已经切换到其他话题，只保存结果，不发送消息（避免打扰）
        let outbound_messages = if is_current_turn {
            request
                .result
                .emitted_messages
                .iter()
                .filter(|message| {
                    // 当存在 live_emitter 时，所有消息已在 loop 中实时广播，不需要 post-loop 发送
                    !request.suppress_live_tool_calls
                        && should_display_message_to_user(self.show_tool_results, message)
                })
                .flat_map(|message| {
                    OutboundMessage::from_chat_message(
                        request.channel_name,
                        request.chat_id,
                        None, // session_id
                        None,
                        request.metadata,
                        message,
                    )
                })
                .collect()
        } else {
            Vec::new()
        };

        // 只有当是最新回合时才触发历史压缩
        let should_schedule_compaction = is_current_turn;

        Ok(FinalizedAgentResult {
            outbound_messages,
            should_schedule_compaction,
        })
    }

    pub(crate) async fn prepare_and_execute_message(
        &self,
        request: MessageExecutionRequest<'_>,
    ) -> Result<Vec<OutboundMessage>, AgentError> {
        // 获取该 topic 的串行锁（通过短暂获取 session 锁）
        // 同一 topic 的消息处理必须串行执行，防止并发 loop 操作同一历史的不同快照
        // 不同 topic 之间互不阻塞，支持多话题并发执行
        let serial_lock = {
            let mut session_guard = request.session.lock().await;
            let lock_key = request.topic_id.as_deref().unwrap_or(request.chat_id);
            session_guard.topic_serial_lock(lock_key)
        };

        // 等待该 topic 的前一条消息处理完成（含压缩）
        // await 串行锁时不持有 session 锁，其他 topic 的消息可以正常处理
        let _serial_guard = serial_lock.lock().await;

        let (history, agent, user_message, user_message_count, original_topic_id) = {
            let mut session_guard = request.session.lock().await;

            session_guard.ensure_persistent_session(request.chat_id)?;

            // 优先使用消息接收时捕获的 topic_id，消除 #1 与 #2 之间的竞态
            let original_topic_id = match &request.topic_id {
                Some(tid) => Some(tid.clone()),
                None => session_guard
                    .current_topic(request.chat_id)
                    .map(|s| s.to_string()),
            };

            session_guard.ensure_chat_loaded(request.chat_id, original_topic_id.as_deref())?;

            session_guard.ensure_agent_prompt_before_user_message(request.chat_id)?;

            let media_refs: Vec<String> = request
                .media
                .iter()
                .map(|media| media.path.clone())
                .collect();
            #[cfg(debug_assertions)]
            if !media_refs.is_empty() {
                tracing::debug!(media_count = %request.media.len(), media_refs = ?media_refs, "Adding user message with media");
            }
            let enriched_content =
                enrich_user_content_with_media_refs(request.content, &media_refs)?;

            // 先计算 user_message_count（在添加新消息之前）
            // 无 topic 时用 chat_id 作为 topic_histories 的回退 key
            let history_key = original_topic_id.as_deref().unwrap_or(request.chat_id);
            let history_before = session_guard.get_or_create_history(history_key).clone();
            let user_message_count = history_before.iter().filter(|m| m.role == "user").count();

            let user_message = session_guard.create_user_message(&enriched_content, media_refs);
            session_guard.append_persisted_message(
                request.chat_id,
                original_topic_id.as_deref(),
                user_message.clone(),
            )?;

            // 再获取包含新消息的完整历史记录
            let history = session_guard.get_or_create_history(history_key).clone();
            session_guard.record_skill_offer(request.chat_id)?;

            let mut agent = session_guard.create_agent(
                request.chat_id,
                Some(request.sender_id),
                Some(&user_message.id),
                original_topic_id.as_deref(),
            )?;
            if let Some(handler) = request.live_emitter.clone() {
                agent = agent.with_emitted_message_handler(handler);
            }

            (
                history,
                agent,
                user_message,
                user_message_count,
                original_topic_id,
            )
        };

        // 构建系统提示词上下文
        let system_prompt_context = SystemPromptContext {
            session_id: Some(persistent_session_id(request.channel_name, request.chat_id)),
            chat_id: request.chat_id.to_string(),
            user_message_count,
        };

        let result = agent.process(history, Some(&system_prompt_context)).await?;
        let mut metadata = HashMap::new();
        // 把用户消息的 UUID 回传给前端，前端用此更新本地消息 ID，使 todo 点击跳转能匹配
        metadata.insert("user_message_id".to_string(), user_message.id.clone());

        self.finalize_result_and_schedule_compaction(
            request.session.clone(),
            FinalizeAgentResultRequest {
                channel_name: request.channel_name,
                chat_id: request.chat_id,
                user_message: &user_message,
                result,
                metadata: &metadata,
                suppress_live_tool_calls: request.live_emitter.is_some(),
                execution_kind: "message",
                original_topic_id,
            },
        )
        .await
    }

    pub(crate) async fn prepare_and_execute_scheduled_task(
        &self,
        request: ScheduledExecutionRequest<'_>,
    ) -> Result<Vec<OutboundMessage>, AgentError> {
        // 获取该 topic 的串行锁（与普通消息路径共享，保证串行执行）
        // 定时任务由调度器触发，无用户消息竞态；在锁前一次性捕获 topic_id，
        // 锁后复用同一值作为 original_topic_id，保证锁键与写入目标一致。
        let (serial_lock, lock_time_topic_id) = {
            let mut session_guard = request.session.lock().await;
            let tid = session_guard
                .current_topic(request.chat_id)
                .map(|s| s.to_string());
            let lock_key = tid.as_deref().unwrap_or(request.chat_id);
            (session_guard.topic_serial_lock(lock_key), tid)
        };

        // 等待该 topic 的前一条消息处理完成（含压缩）
        let _serial_guard = serial_lock.lock().await;

        let (
            history,
            mut agent,
            user_message,
            user_message_count,
            original_topic_id,
            store,
            session_id,
        ) = {
            let mut session_guard = request.session.lock().await;

            session_guard.ensure_persistent_session(request.chat_id)?;

            // 复用锁前捕获的 topic_id，保证锁键与写入目标一致
            let original_topic_id = lock_time_topic_id.clone();

            // 如果 fresh_session 为 true，清理历史（内存 + 数据库）
            if request.fresh_session {
                session_guard.clear_chat_history(request.chat_id, original_topic_id.as_deref())?;
                tracing::info!(
                    chat_id = %request.chat_id,
                    "Fresh session enabled, history cleared"
                );
            }

            session_guard.ensure_chat_loaded(request.chat_id, original_topic_id.as_deref())?;
            session_guard.ensure_agent_prompt_before_user_message(request.chat_id)?;

            let scheduled_system_prompt =
                compose_scheduled_task_system_prompt(request.system_prompt);
            session_guard.append_persisted_message(
                request.chat_id,
                original_topic_id.as_deref(),
                ChatMessage::system_with_context(
                    &scheduled_system_prompt,
                    Some(SYSTEM_CONTEXT_SCHEDULED_PROMPT.to_string()),
                ),
            )?;

            // 先计算 user_message_count（在添加新消息之前）
            let history_key = original_topic_id.as_deref().unwrap_or(request.chat_id);
            let history_before = session_guard.get_or_create_history(history_key).clone();
            let user_message_count = history_before.iter().filter(|m| m.role == "user").count();

            let user_message = session_guard.create_user_message(request.prompt, Vec::new());
            session_guard.append_persisted_message(
                request.chat_id,
                original_topic_id.as_deref(),
                user_message.clone(),
            )?;

            // 再获取包含新消息的完整历史记录
            let history = session_guard.get_or_create_history(history_key).clone();
            session_guard.record_skill_offer(request.chat_id)?;

            let agent = session_guard.create_agent_with_provider_config(
                request.chat_id,
                request.notification_chat_id, // 传入真实 chat_id
                Some(request.sender_id),
                Some(&user_message.id),
                request.provider_config.clone(),
                original_topic_id.as_deref(),
            )?;

            // 获取 store 和 session_id，用于构造消息持久化 handler
            let store = session_guard.store();
            let session_id =
                crate::storage::persistent_session_id(request.channel_name, request.chat_id);

            (
                history,
                agent,
                user_message,
                user_message_count,
                original_topic_id,
                store,
                session_id,
            )
        };

        // 定时任务没有 live_emitter，需要 PersistingEmittedMessageHandler 来持久化消息
        {
            let persisting_handler = PersistingEmittedMessageHandler::new(
                NoOpEmittedMessageHandler,
                store as Arc<dyn ConversationRepository>,
                &session_id,
                None,
            );
            agent = agent.with_emitted_message_handler(Arc::new(persisting_handler));
        }

        // 构建系统提示词上下文
        let system_prompt_context = SystemPromptContext {
            session_id: Some(persistent_session_id(request.channel_name, request.chat_id)),
            chat_id: request.chat_id.to_string(),
            user_message_count,
        };

        let result = agent.process(history, Some(&system_prompt_context)).await?;

        let outbound_messages = self
            .finalize_result_and_schedule_compaction(
                request.session.clone(),
                FinalizeAgentResultRequest {
                    channel_name: request.channel_name,
                    chat_id: request.chat_id,
                    user_message: &user_message,
                    result,
                    metadata: request.metadata,
                    suppress_live_tool_calls: false,
                    execution_kind: "scheduled_task",
                    original_topic_id: original_topic_id.clone(),
                },
            )
            .await?;

        // 清理内存历史，释放内存（数据库历史保留）
        {
            let mut session_guard = request.session.lock().await;
            let history_key = original_topic_id.as_deref().unwrap_or(request.chat_id);
            session_guard.remove_history(history_key);
            tracing::info!(
                chat_id = %request.chat_id,
                "Scheduled task completed, memory history released"
            );
        }

        Ok(outbound_messages)
    }

    pub(crate) async fn finalize_result_and_schedule_compaction(
        &self,
        session: Arc<Mutex<Session>>,
        request: FinalizeAgentResultRequest<'_>,
    ) -> Result<Vec<OutboundMessage>, AgentError> {
        let channel_name = request.channel_name.to_string();
        let chat_id = request.chat_id.to_string();
        let execution_kind = request.execution_kind.to_string();
        let topic_id = request.original_topic_id.clone();

        let finalized_result = {
            let mut session_guard = session.lock().await;
            self.finalize_result(&mut session_guard, request)?
        };

        if finalized_result.should_schedule_compaction {
            let compaction_topic_id = topic_id.unwrap_or_else(|| chat_id.clone());
            if let Err(error) = schedule_background_history_compaction(
                session.clone(),
                chat_id.clone(),
                compaction_topic_id,
            )
            .await
            {
                tracing::warn!(
                    channel = %channel_name,
                    chat_id = %chat_id,
                    execution_kind = %execution_kind,
                    error = %error,
                    "Failed to schedule background history compaction"
                );
            }
        }

        Ok(finalized_result.outbound_messages)
    }
}

pub(crate) fn should_display_message_to_user(
    show_tool_results: bool,
    message: &ChatMessage,
) -> bool {
    if message.role != "tool" {
        return true;
    }

    show_tool_results
        || matches!(
            message
                .tool_state
                .as_ref()
                .unwrap_or(&ToolMessageState::Completed),
            ToolMessageState::PendingUserAction
        )
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::bus::ChatMessage;

    #[test]
    fn test_compose_scheduled_task_system_prompt_appends_task_specific_prompt() {
        let prompt = compose_scheduled_task_system_prompt(Some(" 只汇报异常 "));

        assert!(prompt.contains("当前输入来自一次已经触发的定时任务执行"));
        assert!(prompt.contains("任务专属要求：只汇报异常"));
    }

    #[test]
    fn test_compose_scheduled_task_system_prompt_ignores_blank_override() {
        let prompt = compose_scheduled_task_system_prompt(Some("   "));

        assert!(prompt.contains("当前输入来自一次已经触发的定时任务执行"));
        assert!(!prompt.contains("任务专属要求"));
    }

    #[test]
    fn test_should_display_message_to_user_keeps_pending_tool_action_visible() {
        let message = ChatMessage::tool_with_state(
            "call-1",
            "approval",
            "需要用户确认",
            ToolMessageState::PendingUserAction,
        );

        assert!(should_display_message_to_user(false, &message));
    }

    #[test]
    fn test_should_display_message_to_user_hides_completed_tool_when_disabled() {
        let message = ChatMessage::tool("call-1", "calculator", "2");

        assert!(!should_display_message_to_user(false, &message));
        assert!(should_display_message_to_user(true, &message));
    }

    /// 对抗性测试：同一 topic 的串行锁被持有时，第二次获取应阻塞
    #[tokio::test]
    async fn test_topic_serial_lock_blocks_concurrent_access() {
        let lock = std::sync::Arc::new(tokio::sync::Mutex::new(()));
        let _guard1 = lock.lock().await;

        // 第二次获取应阻塞，1ms 超时验证
        let result = tokio::time::timeout(std::time::Duration::from_millis(1), lock.lock()).await;

        assert!(result.is_err(), "第二次获取同一锁应阻塞");
    }

    /// 对抗性测试：不同 topic 的串行锁互不影响，可同时获取
    #[tokio::test]
    async fn test_different_topic_locks_independent() {
        let lock_a = std::sync::Arc::new(tokio::sync::Mutex::new(()));
        let lock_b = std::sync::Arc::new(tokio::sync::Mutex::new(()));

        let _guard_a = lock_a.lock().await;

        // 不同锁应立即可获取
        let result =
            tokio::time::timeout(std::time::Duration::from_millis(100), lock_b.lock()).await;

        assert!(result.is_ok(), "不同 topic 的锁应互不影响");
    }

    /// 对抗性测试：错误返回路径锁被正确释放（RAII 保证）
    #[tokio::test]
    async fn test_serial_lock_released_on_error() {
        let lock = std::sync::Arc::new(tokio::sync::Mutex::new(()));

        // 模拟 prepare_and_execute_message 的错误路径：
        // 获取锁 → 返回错误 → 锁应通过 RAII 释放
        {
            let _serial_guard = lock.lock().await;
            // 模拟错误返回（`?` 或 `Err` 分支）
            let _result: Result<(), AgentError> = Err(AgentError::Other("simulated".to_string()));
            // _serial_guard 在此块结束时 Drop，释放锁
        }

        // 锁应已释放，可再次获取
        let result = tokio::time::timeout(std::time::Duration::from_millis(100), lock.lock()).await;

        assert!(result.is_ok(), "错误返回后锁应已释放");
    }
}