rucora 0.1.5

//! ReflectAgent - 反思迭代 Agent
//!
//! # 概述
//!
//! ReflectAgent 实现反思迭代循环：
//! 1. **Generate**（生成）：生成初始版本
//! 2. **Reflect**（反思）：自我批评、分析问题
//! 3. **Improve**（改进）：根据反思改进
//! 4. 循环直到达到质量阈值或最大迭代次数
//!
//! # 适用场景
//!
//! - 代码生成（需要高质量代码）
//! - 文档写作
//! - 方案设计
//! - 任何需要迭代改进的任务
//!
//! # 使用示例
//!
//! ```rust,no_run
//! use rucora::agent::ReflectAgent;
//! use rucora::provider::OpenAiProvider;
//! use rucora::tools::FileWriteTool;
//!
//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
//! let provider = OpenAiProvider::from_env()?;
//!
//! let agent = ReflectAgent::builder()
//!     .provider(provider)
//!     .model("gpt-4o-mini")
//!     .system_prompt("你是一个追求卓越的程序员")
//!     .tool(FileWriteTool)
//!     .max_iterations(3)
//!     .quality_threshold(0.9)
//!     .try_build()?;
//!
//! let output = agent.run("帮我写一个快速排序算法，要求有详细注释").await?;
//! # Ok(())
//! # }
//! ```

use async_trait::async_trait;
use rucora_core::agent::{Agent, AgentContext, AgentDecision, AgentError, AgentInput, AgentOutput};
use rucora_core::provider::LlmProvider;
use rucora_core::provider::types::{ChatMessage, ChatRequest, LlmParams, Role};
use rucora_core::tool::Tool;
use serde_json::{Value, json};
use std::sync::Arc;
use tokio::sync::Mutex;

use crate::agent::ToolRegistry;
use crate::agent::execution::DefaultExecution;
use crate::conversation::ConversationManager;

/// ReflectAgent - 反思迭代 Agent
///
/// 特点：
/// - 生成 - 反思 - 改进循环
/// - 自我批评、持续改进
/// - 适合需要高质量输出的任务
pub struct ReflectAgent<P> {
    /// LLM Provider
    #[allow(dead_code)]
    provider: Arc<P>,
    /// 默认使用的模型
    #[allow(dead_code)]
    model: String,
    /// 系统提示词
    #[allow(dead_code)]
    system_prompt: Option<String>,
    /// 工具注册表
    #[allow(dead_code)]
    tools: ToolRegistry,
    /// 最大迭代次数
    #[allow(dead_code)]
    max_iterations: usize,
    /// 质量阈值（0.0-1.0）
    #[allow(dead_code)]
    quality_threshold: f32,
    /// 对话管理器（可选）
    #[allow(dead_code)]
    conversation_manager: Option<Arc<Mutex<ConversationManager>>>,
    /// LLM 请求参数
    llm_params: LlmParams,
    /// 执行能力（内聚）
    execution: DefaultExecution,
}

#[async_trait]
impl<P> Agent for ReflectAgent<P>
where
    P: LlmProvider + Send + Sync + 'static,
{
    async fn think(&self, context: &AgentContext) -> AgentDecision {
        let iteration = context.step / 2; // 每 2 步为一次迭代（生成 + 反思）

        if iteration == 0 {
            // 第一次：生成初始版本
            AgentDecision::Chat {
                request: Box::new(self._build_generate_request(context)),
            }
        } else if iteration >= self.max_iterations {
            // 达到最大迭代次数，返回当前最佳
            AgentDecision::Return(self._build_final_result(context))
        } else {
            // 奇数步：反思
            if context.step % 2 == 1 {
                AgentDecision::Chat {
                    request: Box::new(self._build_reflect_request(context, iteration)),
                }
            } else {
                // 偶数步：根据反思改进
                AgentDecision::Chat {
                    request: Box::new(self._build_improve_request(context, iteration)),
                }
            }
        }
    }

    fn name(&self) -> &str {
        "reflect_agent"
    }

    fn description(&self) -> Option<&str> {
        Some("反思迭代 Agent，生成 - 反思 - 改进循环")
    }

    /// 运行 Agent（覆盖默认实现，使用 DefaultExecution）
    async fn run(&self, input: AgentInput) -> Result<AgentOutput, rucora_core::agent::AgentError> {
        self.execution.run(self, input).await
    }

    /// 流式运行
    fn run_stream(
        &self,
        input: AgentInput,
    ) -> futures_util::stream::BoxStream<
        'static,
        Result<rucora_core::channel::types::ChannelEvent, rucora_core::agent::AgentError>,
    > {
        self.execution.run_stream_simple(input)
    }
}

impl<P> ReflectAgent<P>
where
    P: LlmProvider + Send + Sync + 'static,
{
    /// 流式运行并返回拼接后的最终文本。
    pub async fn run_stream_text(
        &self,
        input: impl Into<AgentInput>,
    ) -> Result<String, rucora_core::agent::AgentError> {
        self.execution.run_stream_text(input.into()).await
    }
}

impl<P> ReflectAgent<P>
where
    P: LlmProvider,
{
    /// 创建新的构建器
    pub fn builder() -> ReflectAgentBuilder<P> {
        ReflectAgentBuilder::new()
    }

    /// 构建生成请求
    fn _build_generate_request(&self, context: &AgentContext) -> ChatRequest {
        let prompt = format!(
            "请生成初始版本：{}\n\
             \n\
             要求：\n\
             1. 完整实现功能\n\
             2. 保证正确性\n\
             3. 尽可能详细\n\
             \n\
             稍后我会进行自我反思和改进。",
            context.input.text()
        );

        self._build_request(context, prompt)
    }

    /// 构建反思请求
    fn _build_reflect_request(&self, context: &AgentContext, iteration: usize) -> ChatRequest {
        let prompt = format!(
            "请反思第 {iteration} 版本的质量：\n\
             \n\
             反思维度：\n\
             1. **正确性**：是否有错误或遗漏？\n\
             2. **完整性**：是否覆盖所有需求？\n\
             3. **清晰度**：是否易于理解？\n\
             4. **优化空间**：哪些地方可以改进？\n\
             \n\
             请详细列出问题和改进建议。"
        );

        self._build_request(context, prompt)
    }

    /// 构建改进请求
    fn _build_improve_request(&self, context: &AgentContext, iteration: usize) -> ChatRequest {
        let prompt = format!(
            "请根据反思改进第 {} 版本：\n\
             \n\
             改进要求：\n\
             1. 修复所有发现的问题\n\
             2. 采纳所有合理的改进建议\n\
             3. 保持原有优点\n\
             4. 生成更高质量的版本\n\
             \n\
             目标质量阈值：{}",
            iteration, self.quality_threshold
        );

        self._build_request(context, prompt)
    }

    /// 构建最终结果
    fn _build_final_result(&self, context: &AgentContext) -> Value {
        // 从上下文中提取最新版本
        if let Some(last_msg) = context.messages.last() {
            json!({
                "content": last_msg.content,
                "iterations": self.max_iterations,
                "completed": true
            })
        } else {
            json!({
                "content": "未能生成结果",
                "iterations": self.max_iterations,
                "completed": false
            })
        }
    }

    /// 构建通用请求
    fn _build_request(&self, context: &AgentContext, prompt: String) -> ChatRequest {
        let mut messages = context.messages.clone();

        // 添加系统提示词
        if let Some(ref sys_prompt) = self.system_prompt
            && (messages.is_empty() || messages.first().map(|m| &m.role) != Some(&Role::System))
        {
            messages.insert(0, ChatMessage::system(sys_prompt.clone()));
        }

        // 添加当前提示词
        messages.push(ChatMessage::user(prompt));

        let mut request = ChatRequest {
            messages,
            model: Some(self.model.clone()),
            tools: if !self.tools.definitions().is_empty() {
                Some(self.tools.definitions())
            } else {
                None
            },
            ..Default::default()
        };
        self.llm_params.apply_to(&mut request);
        request
    }

    /// 获取工具列表
    pub fn tools(&self) -> Vec<&str> {
        self.tools
            .tool_names()
            .into_iter()
            .map(|s| s.as_str())
            .collect()
    }
}

/// ReflectAgent 构建器
pub struct ReflectAgentBuilder<P> {
    provider: Option<P>,
    system_prompt: Option<String>,
    model: Option<String>,
    tools: ToolRegistry,
    max_iterations: usize,
    quality_threshold: f32,
    with_conversation: bool,
    middleware_chain: crate::middleware::MiddlewareChain,
    llm_params: LlmParams,
}

impl<P> ReflectAgentBuilder<P> {
    /// 创建新的构建器
    pub fn new() -> Self {
        Self {
            provider: None,
            system_prompt: None,
            model: None,
            tools: ToolRegistry::new(),
            max_iterations: 3,
            quality_threshold: 0.9,
            with_conversation: false,
            middleware_chain: crate::middleware::MiddlewareChain::new(),
            llm_params: LlmParams::default(),
        }
    }
}

impl<P> ReflectAgentBuilder<P>
where
    P: LlmProvider + Send + Sync + 'static,
{
    /// 设置 Provider（必需）
    pub fn provider(mut self, provider: P) -> Self {
        self.provider = Some(provider);
        self
    }

    /// 设置系统提示词
    pub fn system_prompt(mut self, prompt: impl Into<String>) -> Self {
        self.system_prompt = Some(prompt.into());
        self
    }

    /// 设置默认模型（必需）
    pub fn model(mut self, model: impl Into<String>) -> Self {
        self.model = Some(model.into());
        self
    }

    /// 注册工具
    pub fn tool(mut self, tool: impl Tool + 'static) -> Self {
        self.tools = self.tools.register(tool);
        self
    }

    /// 注册多个工具
    pub fn tools<I, T>(mut self, tools: I) -> Self
    where
        I: IntoIterator<Item = T>,
        T: Tool + 'static,
    {
        for tool in tools {
            self.tools = self.tools.register(tool);
        }
        self
    }

    /// 设置最大迭代次数
    pub fn max_iterations(mut self, max: usize) -> Self {
        self.max_iterations = max;
        self
    }

    /// 设置质量阈值（0.0-1.0）
    ///
    /// 当自我评估质量达到此阈值时，提前停止迭代
    pub fn quality_threshold(mut self, threshold: f32) -> Self {
        self.quality_threshold = threshold.clamp(0.0, 1.0);
        self
    }

    /// 设置温度参数（控制随机性，0.0-1.0）
    pub fn temperature(mut self, value: f32) -> Self {
        self.llm_params.temperature = Some(value);
        self
    }

    /// 设置 top_p
    pub fn top_p(mut self, value: f32) -> Self {
        self.llm_params.top_p = Some(value);
        self
    }

    /// 设置 top_k
    pub fn top_k(mut self, value: u32) -> Self {
        self.llm_params.top_k = Some(value);
        self
    }

    /// 设置 max_tokens
    pub fn max_tokens(mut self, value: u32) -> Self {
        self.llm_params.max_tokens = Some(value);
        self
    }

    /// 设置 frequency_penalty
    pub fn frequency_penalty(mut self, value: f32) -> Self {
        self.llm_params.frequency_penalty = Some(value);
        self
    }

    /// 设置 presence_penalty
    pub fn presence_penalty(mut self, value: f32) -> Self {
        self.llm_params.presence_penalty = Some(value);
        self
    }

    /// 设置 stop 序列
    pub fn stop(mut self, value: Vec<String>) -> Self {
        self.llm_params.stop = Some(value);
        self
    }

    /// 设置额外参数（provider 特定）
    pub fn extra_params(mut self, value: serde_json::Value) -> Self {
        self.llm_params.extra = Some(value);
        self
    }

    /// 设置 LLM 请求参数
    pub fn llm_params(mut self, params: LlmParams) -> Self {
        self.llm_params = params;
        self
    }

    /// 启用对话历史管理
    pub fn with_conversation(mut self, enabled: bool) -> Self {
        self.with_conversation = enabled;
        self
    }

    /// 设置中间件链
    pub fn with_middleware_chain(
        mut self,
        middleware_chain: crate::middleware::MiddlewareChain,
    ) -> Self {
        self.middleware_chain = middleware_chain;
        self
    }

    /// 添加中间件
    pub fn with_middleware<M: crate::middleware::Middleware + 'static>(
        mut self,
        middleware: M,
    ) -> Self {
        self.middleware_chain = self.middleware_chain.with(middleware);
        self
    }

    /// 尝试构建 Agent。
    pub fn try_build(self) -> Result<ReflectAgent<P>, AgentError> {
        let provider = self.provider.ok_or_else(|| {
            AgentError::Message("构建 ReflectAgent 失败：缺少 provider".to_string())
        })?;
        let model = self
            .model
            .ok_or_else(|| AgentError::Message("构建 ReflectAgent 失败：缺少 model".to_string()))?;
        let conversation_manager = if self.with_conversation {
            let mut conv = ConversationManager::new();
            if let Some(ref prompt) = self.system_prompt {
                conv = conv.with_system_prompt(prompt.clone());
            }
            Some(Arc::new(Mutex::new(conv)))
        } else {
            None
        };

        // 创建执行能力
        let provider_arc = Arc::new(provider);
        let execution =
            DefaultExecution::new(provider_arc.clone(), model.clone(), self.tools.clone())
                .with_system_prompt_opt(self.system_prompt.clone())
                .with_max_steps(self.max_iterations * 2) // 每次迭代需要 2 步
                .with_conversation_manager(conversation_manager.clone())
                .with_middleware_chain(self.middleware_chain)
                .with_llm_params(self.llm_params.clone());

        Ok(ReflectAgent {
            provider: provider_arc,
            model,
            system_prompt: self.system_prompt,
            tools: self.tools,
            max_iterations: self.max_iterations,
            quality_threshold: self.quality_threshold,
            conversation_manager,
            llm_params: self.llm_params,
            execution,
        })
    }

    /// 构建 Agent。
    ///
    /// 推荐优先使用 [`Self::try_build`] 处理配置错误。
    /// 此方法保留为便捷入口，内部仍会在配置缺失时 panic。
    pub fn build(self) -> ReflectAgent<P> {
        self.try_build()
            .unwrap_or_else(|err| panic!("ReflectAgentBuilder::build 失败：{err}"))
    }
}

impl<P> Default for ReflectAgentBuilder<P> {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use futures_util::stream;
    use futures_util::stream::BoxStream;
    use rucora_core::error::ProviderError;
    use rucora_core::provider::types::{ChatResponse, ChatStreamChunk};

    struct MockProvider;

    #[async_trait]
    impl LlmProvider for MockProvider {
        async fn chat(&self, _request: ChatRequest) -> Result<ChatResponse, ProviderError> {
            Ok(ChatResponse {
                message: ChatMessage {
                    role: Role::Assistant,
                    content: "Mock response".to_string(),
                    name: None,
                },
                tool_calls: vec![],
                usage: None,
                finish_reason: None,
            })
        }

        fn stream_chat(
            &self,
            _request: ChatRequest,
        ) -> Result<BoxStream<'static, Result<ChatStreamChunk, ProviderError>>, ProviderError>
        {
            Ok(Box::pin(stream::empty()))
        }
    }

    #[test]
    fn test_reflect_agent_builder() {
        let _agent = ReflectAgentBuilder::<MockProvider>::new()
            .provider(MockProvider)
            .model("gpt-4o-mini")
            .max_iterations(3)
            .quality_threshold(0.9)
            .build();
    }
}