otherone-agent 0.1.1

// 作用：Agent 循环驱动 — 整个框架的核心，驱动 AI 对话流程
// 关联：调用 ai、context、tools、storage 等所有子模块
// 预期结果：根据 input 和 ai 配置，执行完整的 Agent 循环，返回最终响应

pub mod error;
pub mod response_parser;
pub mod types;

use std::collections::HashMap;
use std::time::Duration;
use tokio::sync::mpsc;

use crate::error::AgentError;
use crate::response_parser::parse_ai_response;
use crate::types::{
    AiOptions, ContextLoadType as AgentContextLoadType, InputOptions,
    StorageType as AgentStorageType,
};
use otherone_storage::types::{StorageType, WriteEntryOptions};

/// 流式 Agent 事件类型
/// 作用：定义流式 Agent 循环中 yield 给调用者的事件
/// 关联：被 invoke_agent_stream 使用
/// 预期结果：调用者可以根据 event_type 区分不同的事件
#[derive(Debug, Clone)]
pub struct StreamAgentEvent {
    /// 事件类型: "chunk" | "thinking" | "tool_calls" | "complete" | "error"
    pub event_type: String,
    /// 事件内容
    pub content: String,
    /// 原始 chunk 数据（当 event_type 为 "chunk" 时）
    pub raw_chunk: Option<serde_json::Value>,
    /// 错误信息（当 event_type 为 "error" 时）
    pub error: Option<String>,
}

/// 流式 Agent 调用
/// 作用：启动流式 Agent 循环，通过 mpsc channel 实时发送事件
/// 关联：被 otherone 主 crate 的 Otherone::invoke_agent_stream 调用
/// 预期结果：返回 mpsc Receiver，调用者可以异步迭代接收事件
pub async fn invoke_agent_stream(
    input: InputOptions,
    mut ai: AiOptions,
) -> Result<mpsc::Receiver<StreamAgentEvent>, AgentError> {
    let (tx, rx) = mpsc::channel::<StreamAgentEvent>(256);

    tokio::spawn(async move {
        if let Err(e) = run_stream_loop(input, &mut ai, &tx).await {
            let _ = tx
                .send(StreamAgentEvent {
                    event_type: "error".to_string(),
                    content: format!("[error:{}]", e),
                    raw_chunk: None,
                    error: Some(e.to_string()),
                })
                .await;
        }
    });

    Ok(rx)
}

/// 流式循环实际逻辑
/// 作用：在后台 task 中执行完整的流式 Agent 循环
/// 关联：被 invoke_agent_stream 调用
/// 预期结果：通过 tx channel 发送所有事件，循环结束后返回
async fn run_stream_loop(
    input: InputOptions,
    ai: &mut AiOptions,
    tx: &mpsc::Sender<StreamAgentEvent>,
) -> Result<(), AgentError> {
    let storage_type: StorageType = match &input.storage_type {
        Some(AgentStorageType::LocalFile) => StorageType::LocalFile,
        Some(AgentStorageType::Database) => StorageType::Database,
        None => StorageType::LocalFile,
    };

    // 存储用户消息
    if let Some(ref user_prompt) = ai.user_prompt {
        otherone_storage::write_entry(&WriteEntryOptions {
            storage_type: storage_type.clone(),
            session_id: input.session_id.clone(),
            role: "user".to_string(),
            content: user_prompt.clone(),
            tools: None,
            token_consumption: None,
            create_at: None,
            database_config: input.database_config.clone(),
        })
        .await
        .map_err(|e| AgentError::ContextError(e.to_string()))?;
    }

    let max_iterations = input.max_iterations.unwrap_or(999999);
    let mut iteration: u32 = 0;

    while iteration < max_iterations {
        iteration += 1;

        // 组合 tools 配置
        let tools = otherone_tools::combine_tools(ai.tools.clone());
        ai.tools = tools;

        // 加载上下文
        let context_load_type = match &input.context_load_type {
            AgentContextLoadType::LocalFile => otherone_context::types::ContextLoadType::LocalFile,
            AgentContextLoadType::Database => otherone_context::types::ContextLoadType::Database,
        };

        let messages =
            otherone_context::combine_context(&otherone_context::types::CombineContextOptions {
                session_id: input.session_id.clone(),
                load_type: context_load_type,
                provider: ai.provider.clone(),
                context_window: input.context_window,
                threshold_percentage: input.threshold_percentage,
                ai: ai.other.clone(),
                system_prompt: ai.system_prompt.clone(),
                tools: ai.tools.clone(),
                database_config: input.database_config.clone(),
            })
            .await
            .map_err(|e| AgentError::ContextError(e.to_string()))?;

        // 构建 AI 配置，强制流式
        let mut ai_config = build_ai_config(ai, &messages);
        ai_config["stream"] = serde_json::json!(true);

        // 调用流式 AI 模型
        let mut stream = otherone_ai::invoke_model_stream(
            ai.provider.clone(),
            &ai.api_key,
            &ai.base_url,
            ai_config,
        )
        .await?;

        // 累积变量
        let mut full_content = String::new();
        let mut role = "assistant".to_string();
        let mut stream_tool_calls: Vec<otherone_ai::types::ToolCall> = Vec::new();
        let mut token_consumption: u32 = 0;

        // 遍历流，实时发送给调用者
        use futures::StreamExt;
        while let Some(chunk_result) = stream.next().await {
            let chunk = match chunk_result {
                Ok(c) => c,
                Err(e) => {
                    let _ = tx
                        .send(StreamAgentEvent {
                            event_type: "error".to_string(),
                            content: format!("[error:{}]", e),
                            raw_chunk: None,
                            error: Some(e.to_string()),
                        })
                        .await;
                    return Err(AgentError::AiError(e));
                }
            };

            // 发送原始 chunk
            let raw_json = serde_json::to_value(&chunk).unwrap_or_default();
            let _ = tx
                .send(StreamAgentEvent {
                    event_type: "chunk".to_string(),
                    content: String::new(),
                    raw_chunk: Some(raw_json),
                    error: None,
                })
                .await;

            // 提取 delta 信息
            if let Some(choice) = chunk.choices.first() {
                if let Some(ref delta) = choice.delta {
                    if let Some(ref r) = delta.role {
                        role = r.clone();
                    }
                    if let Some(ref c) = delta.content {
                        full_content.push_str(c);
                    }
                    if let Some(thinking) = delta_thinking_content(delta) {
                        let _ = tx
                            .send(StreamAgentEvent {
                                event_type: "thinking".to_string(),
                                content: thinking.to_string(),
                                raw_chunk: None,
                                error: None,
                            })
                            .await;
                    }
                }

                if let Some(ref delta) = choice.delta {
                    // 累积 tool_calls
                    if let Some(ref tc) = delta.tool_calls {
                        for tool_call in tc {
                            // 按 index 或 id 匹配合并
                            match stream_tool_calls.iter_mut().find(|existing| {
                                !tool_call.id.is_empty() && existing.id == tool_call.id
                            }) {
                                Some(existing) => {
                                    if !tool_call.function.name.is_empty() {
                                        existing.function.name.push_str(&tool_call.function.name);
                                    }
                                    if !tool_call.function.arguments.is_empty() {
                                        existing
                                            .function
                                            .arguments
                                            .push_str(&tool_call.function.arguments);
                                    }
                                }
                                None => {
                                    stream_tool_calls.push(tool_call.clone());
                                }
                            }
                        }
                    }
                }

                if let Some(ref usage) = chunk.usage {
                    token_consumption = usage.total_tokens.unwrap_or(0);
                }
            }
        }

        // 流结束，构建完整响应
        let tools_wrapper = if stream_tool_calls.is_empty() {
            None
        } else {
            Some(otherone_ai::types::ToolCallsWrapper {
                tool_calls: stream_tool_calls.clone(),
            })
        };

        // 存储 AI 响应
        otherone_storage::write_entry(&WriteEntryOptions {
            storage_type: storage_type.clone(),
            session_id: input.session_id.clone(),
            role: role.clone(),
            content: full_content.clone(),
            tools: None,
            token_consumption: Some(token_consumption),
            create_at: None,
            database_config: input.database_config.clone(),
        })
        .await
        .map_err(|e| AgentError::ContextError(e.to_string()))?;

        // 检查是否有 tool 调用
        if let Some(ref tw) = tools_wrapper {
            if !tw.tool_calls.is_empty() {
                let tool_calls_info: Vec<String> = tw
                    .tool_calls
                    .iter()
                    .map(|tc| format!("{}({})", tc.function.name, tc.function.arguments))
                    .collect();

                let _ = tx
                    .send(StreamAgentEvent {
                        event_type: "tool_calls".to_string(),
                        content: format!("[tool_calls:{}]", tool_calls_info.join(", ")),
                        raw_chunk: None,
                        error: None,
                    })
                    .await;

                // 处理 tool 调用
                let default_tools_realize: HashMap<
                    String,
                    Box<dyn Fn(Vec<String>) -> String + Send + Sync>,
                > = HashMap::new();
                let tools_realize = ai.tools_realize.as_ref().unwrap_or(&default_tools_realize);

                let tool_results = otherone_tools::process_tools(&tw.tool_calls, tools_realize)
                    .map_err(|e| AgentError::ToolError(e))?;

                for tool_result in &tool_results {
                    let tool_content = serde_json::to_string(
                        tool_result
                            .result
                            .as_ref()
                            .unwrap_or(&serde_json::Value::Null),
                    )
                    .unwrap_or_default();

                    let tools_value = serde_json::json!({
                        "tool_call_id": tool_result.tool_call_id,
                        "function_name": tool_result.function_name,
                        "result": tool_result.result,
                        "error": tool_result.error,
                    });

                    otherone_storage::write_entry(&WriteEntryOptions {
                        storage_type: storage_type.clone(),
                        session_id: input.session_id.clone(),
                        role: "tool".to_string(),
                        content: tool_content,
                        tools: Some(tools_value),
                        token_consumption: None,
                        create_at: None,
                        database_config: input.database_config.clone(),
                    })
                    .await
                    .map_err(|e| AgentError::ContextError(e.to_string()))?;
                }

                tokio::time::sleep(Duration::from_millis(1500)).await;
                continue;
            }
        }

        // 发送完成事件
        let _ = tx
            .send(StreamAgentEvent {
                event_type: "complete".to_string(),
                content: full_content,
                raw_chunk: None,
                error: None,
            })
            .await;

        return Ok(());
    }

    let err_msg = format!(
        "Agent循环次数超过限制({}次)，可能陷入无限循环",
        max_iterations
    );
    let _ = tx
        .send(StreamAgentEvent {
            event_type: "error".to_string(),
            content: format!("[error:{}]", err_msg),
            raw_chunk: None,
            error: Some(err_msg.clone()),
        })
        .await;
    Err(AgentError::MaxIterationsExceeded(max_iterations))
}

fn delta_thinking_content(delta: &otherone_ai::types::ResponseDelta) -> Option<&str> {
    delta
        .reasoning_content
        .as_deref()
        .or(delta.reasoning.as_deref())
        .or(delta.thinking.as_deref())
        .or(delta.thought.as_deref())
        .filter(|content| !content.is_empty())
}

/// 调用 Agent — 核心驱动方法
/// 作用：驱动完整的 AI 对话流程，支持非流式模式
/// 关联：调用所有子模块
/// 预期结果：返回最终的 ParsedResponse
pub async fn invoke_agent(
    input: &InputOptions,
    ai: &mut AiOptions,
) -> Result<otherone_ai::types::ParsedResponse, AgentError> {
    let storage_type: StorageType = match &input.storage_type {
        Some(AgentStorageType::LocalFile) => StorageType::LocalFile,
        Some(AgentStorageType::Database) => StorageType::Database,
        None => StorageType::LocalFile,
    };

    // 如果有 user_prompt，先存储用户消息
    if let Some(ref user_prompt) = ai.user_prompt {
        otherone_storage::write_entry(&WriteEntryOptions {
            storage_type: storage_type.clone(),
            session_id: input.session_id.clone(),
            role: "user".to_string(),
            content: user_prompt.clone(),
            tools: None,
            token_consumption: None,
            create_at: None,
            database_config: input.database_config.clone(),
        })
        .await
        .map_err(|e| AgentError::ContextError(e.to_string()))?;
    }

    // 循环次数限制
    let max_iterations = input.max_iterations.unwrap_or(999999);
    let mut iteration: u32 = 0;

    while iteration < max_iterations {
        iteration += 1;

        // 组合 tools 配置
        let tools = otherone_tools::combine_tools(ai.tools.clone());
        ai.tools = tools;

        // 组合 context 配置，加载历史消息
        let context_load_type = match &input.context_load_type {
            AgentContextLoadType::LocalFile => otherone_context::types::ContextLoadType::LocalFile,
            AgentContextLoadType::Database => otherone_context::types::ContextLoadType::Database,
        };

        let messages =
            otherone_context::combine_context(&otherone_context::types::CombineContextOptions {
                session_id: input.session_id.clone(),
                load_type: context_load_type,
                provider: ai.provider.clone(),
                context_window: input.context_window,
                threshold_percentage: input.threshold_percentage,
                ai: ai.other.clone(),
                system_prompt: ai.system_prompt.clone(),
                tools: ai.tools.clone(),
                database_config: input.database_config.clone(),
            })
            .await
            .map_err(|e| AgentError::ContextError(e.to_string()))?;

        // 构建 AI 配置 JSON
        let ai_config = build_ai_config(ai, &messages);

        // 调用 AI 模型
        let response =
            otherone_ai::invoke_model(ai.provider.clone(), &ai.api_key, &ai.base_url, ai_config)
                .await?;

        // 解析响应
        let mut parsed =
            parse_ai_response(&response, &ai.provider).map_err(|e| AgentError::ContextError(e))?;

        // 存储 AI 响应到 storage
        otherone_storage::write_entry(&WriteEntryOptions {
            storage_type: storage_type.clone(),
            session_id: input.session_id.clone(),
            role: parsed.role.clone(),
            content: parsed.content.clone(),
            tools: None,
            token_consumption: Some(parsed.token_consumption),
            create_at: None,
            database_config: input.database_config.clone(),
        })
        .await
        .map_err(|e| AgentError::ContextError(e.to_string()))?;

        // 检查是否有 tool 调用
        if let Some(ref tools_wrapper) = parsed.tools {
            if !tools_wrapper.tool_calls.is_empty() {
                let tool_calls_info: Vec<String> = tools_wrapper
                    .tool_calls
                    .iter()
                    .map(|tc| format!("{}({})", tc.function.name, tc.function.arguments))
                    .collect();
                parsed.content = format!(
                    "[tool_calls:{}]\n\n{}",
                    tool_calls_info.join(", "),
                    parsed.content
                );

                // 使用用户传入的 tools_realize，若未提供则使用空 HashMap
                let default_tools_realize: HashMap<
                    String,
                    Box<dyn Fn(Vec<String>) -> String + Send + Sync>,
                > = HashMap::new();
                let tools_realize = ai.tools_realize.as_ref().unwrap_or(&default_tools_realize);

                let tool_results =
                    otherone_tools::process_tools(&tools_wrapper.tool_calls, tools_realize)
                        .map_err(|e| AgentError::ToolError(e))?;

                for tool_result in &tool_results {
                    let tool_content = serde_json::to_string(
                        tool_result
                            .result
                            .as_ref()
                            .unwrap_or(&serde_json::Value::Null),
                    )
                    .unwrap_or_default();

                    let tools_value = serde_json::json!({
                        "tool_call_id": tool_result.tool_call_id,
                        "function_name": tool_result.function_name,
                        "result": tool_result.result,
                        "error": tool_result.error,
                    });

                    otherone_storage::write_entry(&WriteEntryOptions {
                        storage_type: storage_type.clone(),
                        session_id: input.session_id.clone(),
                        role: "tool".to_string(),
                        content: tool_content,
                        tools: Some(tools_value),
                        token_consumption: None,
                        create_at: None,
                        database_config: input.database_config.clone(),
                    })
                    .await
                    .map_err(|e| AgentError::ContextError(e.to_string()))?;
                }

                tokio::time::sleep(Duration::from_millis(1500)).await;
                continue;
            }
        }

        return Ok(parsed);
    }

    Err(AgentError::MaxIterationsExceeded(max_iterations))
}

/// 构建 AI 调用的配置 JSON
fn build_ai_config(ai: &AiOptions, messages: &[otherone_ai::types::Message]) -> serde_json::Value {
    let mut config = serde_json::json!({
        "model": ai.model,
        "messages": messages,
    });

    if let Some(context_length) = ai.context_length {
        config["contextLength"] = serde_json::json!(context_length);
    }
    if let Some(temperature) = ai.temperature {
        config["temperature"] = serde_json::json!(temperature);
    }
    if let Some(top_p) = ai.top_p {
        config["topP"] = serde_json::json!(top_p);
    }
    if let Some(ref tools) = ai.tools {
        config["tools"] = serde_json::to_value(tools).unwrap();
    }
    if let Some(ref tool_choice) = ai.tool_choice {
        config["toolChoice"] = serde_json::to_value(tool_choice).unwrap();
    }
    if let Some(parallel_tool_calls) = ai.parallel_tool_calls {
        config["parallelToolCalls"] = serde_json::json!(parallel_tool_calls);
    }
    if let Some(stream) = ai.stream {
        config["stream"] = serde_json::json!(stream);
    }
    if let Some(ref other) = ai.other {
        if let serde_json::Value::Object(ref obj) = other {
            for (key, value) in obj {
                config[key] = value.clone();
            }
        }
    }

    config
}

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

    #[test]
    fn test_build_ai_config_basic() {
        let ai = AiOptions {
            provider: otherone_ai::types::ProviderType::OpenAI,
            api_key: "test-key".to_string(),
            base_url: "https://api.openai.com/v1".to_string(),
            model: "gpt-4".to_string(),
            user_prompt: None,
            system_prompt: None,
            messages: None,
            context_length: Some(4096),
            temperature: Some(0.7),
            top_p: None,
            tools: None,
            tools_realize: None,
            tool_choice: None,
            parallel_tool_calls: None,
            stream: None,
            other: None,
        };

        let messages = vec![];
        let config = build_ai_config(&ai, &messages);

        assert_eq!(config["model"], "gpt-4");
        assert_eq!(config["contextLength"], 4096);
        assert!((config["temperature"].as_f64().unwrap() - 0.7).abs() < 0.001);
    }
}