agent-context 0.1.5

//! 集成测试：覆盖工具调用、reasoning_content 保留/剥离场景

use std::sync::Arc;
use std::sync::atomic::Ordering;

use agent_context::{
    AgentContext, ContextBackend, ContextBackendResponse, ContextMessage,
    RequestAppend, RequestMessages, RequestSend, RequestSendStream,
    ResponseType,
};
use agent_context::{AgentError, Role};
use kameo::prelude::*;
use serde::{Deserialize, Serialize};

// ---------------------------------------------------------------------------
// 辅助函数
// ---------------------------------------------------------------------------

fn rich_mock_opts(token: u64) -> RichMockOpts {
    RichMockOpts {
        common: agent_context::CommonOpts {
            model: "test".into(),
            context_window: 10000,
            max_tokens: 4096,
            auto_compress: false,
            scratch: None,
        },
        token,
    }
}

// ---------------------------------------------------------------------------
// RichMockMessage — 模拟真实 provider 消息（含 reasoning_content、tool_calls）
// ---------------------------------------------------------------------------

#[derive(Debug, Clone, Serialize, Deserialize)]
struct ToolCall {
    id: String,
    function_name: String,
    arguments: String,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
struct RichMockMessage {
    role: Role,
    content: String,
    reasoning_content: Option<String>,
    tool_calls: Option<Vec<ToolCall>>,
    tool_call_id: Option<String>,
}

impl RichMockMessage {
    fn user(content: impl Into<String>) -> Self {
        Self {
            role: Role::User,
            content: content.into(),
            reasoning_content: None,
            tool_calls: None,
            tool_call_id: None,
        }
    }

    fn system(content: impl Into<String>) -> Self {
        Self {
            role: Role::System,
            content: content.into(),
            reasoning_content: None,
            tool_calls: None,
            tool_call_id: None,
        }
    }

    fn assistant(content: impl Into<String>) -> Self {
        Self {
            role: Role::Assistant,
            content: content.into(),
            reasoning_content: None,
            tool_calls: None,
            tool_call_id: None,
        }
    }

    fn assistant_with_reasoning(content: impl Into<String>, reasoning: impl Into<String>) -> Self {
        Self {
            role: Role::Assistant,
            content: content.into(),
            reasoning_content: Some(reasoning.into()),
            tool_calls: None,
            tool_call_id: None,
        }
    }

    fn assistant_with_tool_calls(
        content: impl Into<String>,
        reasoning: impl Into<String>,
        tool_calls: Vec<ToolCall>,
    ) -> Self {
        Self {
            role: Role::Assistant,
            content: content.into(),
            reasoning_content: Some(reasoning.into()),
            tool_calls: Some(tool_calls),
            tool_call_id: None,
        }
    }

    fn tool_result(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
        Self {
            role: Role::Tool,
            content: content.into(),
            reasoning_content: None,
            tool_calls: None,
            tool_call_id: Some(tool_call_id.into()),
        }
    }
}

impl ContextMessage for RichMockMessage {
    fn role(&self) -> Role {
        self.role
    }

    /// 模拟 DeepSeek 规则：仅工具调用的 assistant 消息保留 reasoning_content
    fn preserve_reasoning(&self) -> bool {
        self.role == Role::Assistant && self.tool_calls.is_some()
    }

    fn without_reasoning(mut self) -> Self {
        self.reasoning_content = None;
        self
    }

    fn with_role(mut self, role: Role) -> Self {
        self.role = role;
        self
    }
}

// ---------------------------------------------------------------------------
// RichMockBackend — 模拟真实 provider 后端行为
// ---------------------------------------------------------------------------

#[derive(Clone)]
pub struct RichMockOpts {
    pub common: agent_context::CommonOpts,
    pub token: u64,
}

impl std::convert::AsRef<agent_context::CommonOpts> for RichMockOpts {
    fn as_ref(&self) -> &agent_context::CommonOpts {
        &self.common
    }
}

#[derive(Clone)]
struct RichMockBackend {
    send_responses: Vec<Vec<RichMockMessage>>,
    call_index: Arc<std::sync::atomic::AtomicUsize>,
    token_count: usize,
    last_opts_token: Arc<std::sync::atomic::AtomicU64>,
}

impl RichMockBackend {
    fn new(responses: Vec<Vec<RichMockMessage>>) -> Self {
        Self {
            send_responses: responses,
            call_index: Arc::new(std::sync::atomic::AtomicUsize::new(0)),
            token_count: 0,
            last_opts_token: Arc::new(std::sync::atomic::AtomicU64::new(0)),
        }
    }
}

#[derive(Debug, Clone)]
struct RichMockResponse(Vec<RichMockMessage>);

impl ContextBackendResponse for RichMockResponse {
    fn response_type(&self) -> ResponseType {
        let has_r = self
            .0
            .iter()
            .any(|m| m.reasoning_content.as_ref().is_some_and(|r| !r.is_empty()));
        let has_c = self.0.iter().any(|m| !m.content.is_empty());
        match (has_r, has_c) {
            (false, false) => ResponseType::Empty,
            (true, false) => ResponseType::Reasoning,
            (false, true) => ResponseType::Content,
            (true, true) => ResponseType::ReasoningAndContent,
        }
    }

    fn reasoning_content(&self) -> Option<String> {
        let text: String = self
            .0
            .iter()
            .filter_map(|m| m.reasoning_content.clone())
            .collect();
        if text.is_empty() { None } else { Some(text) }
    }

    fn content(&self) -> Option<String> {
        let text: String = self.0.iter().map(|m| m.content.as_str()).collect();
        if text.is_empty() { None } else { Some(text) }
    }

    fn tool_calls(&self) -> Vec<agent_context::ToolCallInfo> {
        self.0
            .iter()
            .filter_map(|m| m.tool_calls.as_ref())
            .flatten()
            .map(|tc| agent_context::ToolCallInfo {
                id: tc.id.clone(),
                name: tc.function_name.clone(),
                arguments: tc.arguments.clone(),
            })
            .collect()
    }
}

impl ContextBackend for RichMockBackend {
    type Message = RichMockMessage;
    type Opts = RichMockOpts;
    type Response = RichMockResponse;

    fn user_message(&self, content: impl Into<String> + Send) -> Self::Message {
        RichMockMessage::user(content)
    }

    fn system_message(&self, content: impl Into<String> + Send) -> Self::Message {
        RichMockMessage::system(content)
    }

    fn tool_message(
        &self,
        tool_call_id: impl Into<String> + Send,
        content: impl Into<String> + Send,
    ) -> Self::Message {
        RichMockMessage::tool_result(tool_call_id, content)
    }

    fn extract_messages(
        &self,
        responses: &[Self::Response],
    ) -> Result<Vec<Self::Message>, AgentError> {
        Ok(responses.iter().flat_map(|r| r.0.iter().cloned()).collect())
    }

    fn merge_chunks(&self, _responses: &[Self::Response]) -> Option<Self::Message> {
        None
    }

    async fn estimate_tokens(&self, _messages: &[Self::Message]) -> Result<usize, AgentError> {
        Ok(self.token_count)
    }

    async fn send(
        &self,
        _messages: &[Self::Message],
        opts: &Self::Opts,
    ) -> Result<Self::Response, AgentError> {
        self.last_opts_token
            .store(opts.token, std::sync::atomic::Ordering::SeqCst);
        let idx = self
            .call_index
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        let responses = &self.send_responses;
        if idx < responses.len() {
            Ok(RichMockResponse(responses[idx].clone()))
        } else {
            Ok(RichMockResponse(vec![]))
        }
    }

    fn send_stream(
        &self,
        _messages: Vec<Self::Message>,
        opts: Self::Opts,
    ) -> impl futures_core::Stream<Item = Result<Self::Response, AgentError>> + Send + 'static {
        self.last_opts_token
            .store(opts.token, std::sync::atomic::Ordering::SeqCst);
        let idx = self
            .call_index
            .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
        let responses = &self.send_responses;
        let chunks = if idx < responses.len() {
            responses[idx].clone()
        } else {
            vec![]
        };
        futures::stream::iter(chunks.into_iter().map(|m| Ok(RichMockResponse(vec![m]))))
    }
}

// ---------------------------------------------------------------------------
// 辅助函数
// ---------------------------------------------------------------------------

fn spawn_with_backend(backend: RichMockBackend) -> ActorRef<AgentContext<RichMockBackend>> {
    let ctx = AgentContext::new(backend, vec![]);
    AgentContext::spawn(ctx)
}

// ---------------------------------------------------------------------------
// 测试
// ---------------------------------------------------------------------------

#[tokio::test]
async fn basic_conversation() {
    let backend = RichMockBackend::new(vec![vec![RichMockMessage::assistant("Hello!")]]);
    let actor = spawn_with_backend(backend);

    actor
        .ask(RequestAppend {
            message: RichMockMessage::user("Hi"),
        })
        .await
        .unwrap();

    // ask() 自动展开 Result：handler 返回 Result<Vec<M>, AgentError>，
    // ask() 返回 Result<Vec<M>, SendError>，unwrap() 得 Vec<M>
    let raw: RichMockResponse = actor
        .ask(RequestSend {
            opts: rich_mock_opts(0),
        })
        .await
        .unwrap();
    assert_eq!(raw.0[0].content, "Hello!");

    let all: Vec<RichMockMessage> = actor.ask(RequestMessages).await.unwrap();
    assert_eq!(all.len(), 2);
    assert_eq!(all[0].role, Role::User);
    assert_eq!(all[1].role, Role::Assistant);
    assert_eq!(all[1].content, "Hello!");

    actor.stop_gracefully().await.unwrap();
}

#[tokio::test]
async fn reasoning_stripped_without_tool_call() {
    // backend 返回带 reasoning_content 的普通 assistant 消息（无 tool_calls）
    let backend = RichMockBackend::new(vec![vec![RichMockMessage::assistant_with_reasoning(
        "答案是42",
        "让我想想...",
    )]]);
    let actor = spawn_with_backend(backend);

    actor
        .ask(RequestAppend {
            message: RichMockMessage::user("问题"),
        })
        .await
        .unwrap();

    // 调用方拿到原始响应（含 reasoning_content）
    let raw: RichMockResponse = actor
        .ask(RequestSend {
            opts: rich_mock_opts(0),
        })
        .await
        .unwrap();
    assert!(
        raw.0[0].reasoning_content.is_some(),
        "原始响应应含 reasoning"
    );

    // 存储的消息应已剥离 reasoning_content
    let all: Vec<RichMockMessage> = actor.ask(RequestMessages).await.unwrap();
    let stored_assistant = &all[1];
    assert_eq!(stored_assistant.role, Role::Assistant);
    assert!(
        stored_assistant.reasoning_content.is_none(),
        "非工具调用场景应剥离 reasoning_content"
    );

    actor.stop_gracefully().await.unwrap();
}

#[tokio::test]
async fn reasoning_preserved_with_tool_call() {
    // backend 返回带 tool_calls + reasoning_content 的 assistant 消息
    let backend = RichMockBackend::new(vec![vec![RichMockMessage::assistant_with_tool_calls(
        "",
        "需要调用工具",
        vec![ToolCall {
            id: "call_1".into(),
            function_name: "get_weather".into(),
            arguments: r#"{"city":"Beijing"}"#.into(),
        }],
    )]]);
    let actor = spawn_with_backend(backend);

    actor
        .ask(RequestAppend {
            message: RichMockMessage::user("北京天气"),
        })
        .await
        .unwrap();

    let raw: RichMockResponse = actor
        .ask(RequestSend {
            opts: rich_mock_opts(0),
        })
        .await
        .unwrap();
    assert_eq!(raw.0.len(), 1, "Should have one response message");

    // 存储的 assistant 消息应保留 reasoning_content（因为有 tool_calls）
    let all: Vec<RichMockMessage> = actor.ask(RequestMessages).await.unwrap();
    let stored_assistant = &all[1];
    assert_eq!(stored_assistant.role, Role::Assistant);
    assert!(
        stored_assistant.reasoning_content.is_some(),
        "工具调用场景应保留 reasoning_content"
    );
    assert!(stored_assistant.tool_calls.is_some());

    actor.stop_gracefully().await.unwrap();
}

#[tokio::test]
async fn full_tool_call_cycle() {
    // 模拟完整工具调用周期：
    // 1. user → assistant(tool_calls + reasoning)
    // 2. tool → assistant(final answer, 无 tool_calls, reasoning → 被剥离)
    let backend = RichMockBackend::new(vec![
        // 第一次 send：返回工具调用
        vec![RichMockMessage::assistant_with_tool_calls(
            "",
            "需要查询天气",
            vec![ToolCall {
                id: "call_1".into(),
                function_name: "get_weather".into(),
                arguments: r#"{"city":"Beijing"}"#.into(),
            }],
        )],
        // 第二次 send：返回最终回答（有 reasoning 无 tool_calls → reasoning 被剥离）
        vec![RichMockMessage::assistant_with_reasoning(
            "北京今天晴，25°C",
            "根据工具返回的数据分析",
        )],
    ]);
    let actor = spawn_with_backend(backend);

    // Round 1: user 发问
    actor
        .ask(RequestAppend {
            message: RichMockMessage::user("北京天气怎么样？"),
        })
        .await
        .unwrap();

    // Send → assistant 返回 tool_calls
    let raw1: RichMockResponse = actor
        .ask(RequestSend {
            opts: rich_mock_opts(0),
        })
        .await
        .unwrap();
    assert!(raw1.0[0].tool_calls.is_some());

    // 存储的消息保留 reasoning_content（有 tool_calls）
    let after_r1: Vec<RichMockMessage> = actor.ask(RequestMessages).await.unwrap();
    assert_eq!(after_r1.len(), 2); // user + assistant
    assert!(
        after_r1[1].reasoning_content.is_some(),
        "工具调用 assistant 保留 reasoning"
    );

    // 追加工具结果
    actor
        .ask(RequestAppend {
            message: RichMockMessage::tool_result("call_1", r#"{"temp":25,"condition":"sunny"}"#),
        })
        .await
        .unwrap();

    // Round 2: Send → assistant 返回最终回答（无 tool_calls）
    let raw2: RichMockResponse = actor
        .ask(RequestSend {
            opts: rich_mock_opts(0),
        })
        .await
        .unwrap();
    assert!(
        raw2.0[0].reasoning_content.is_some(),
        "原始响应含 reasoning"
    );

    // 验证最终上下文：
    // user → assistant(tool_calls, reasoning保留) → tool → assistant(reasoning被剥离)
    let final_msgs: Vec<RichMockMessage> = actor.ask(RequestMessages).await.unwrap();
    assert_eq!(final_msgs.len(), 4);

    assert_eq!(final_msgs[0].role, Role::User);
    assert_eq!(final_msgs[1].role, Role::Assistant);
    assert!(final_msgs[1].tool_calls.is_some());
    assert!(
        final_msgs[1].reasoning_content.is_some(),
        "工具调用 assistant 保留 reasoning"
    );

    assert_eq!(final_msgs[2].role, Role::Tool);
    assert_eq!(final_msgs[2].tool_call_id.as_deref(), Some("call_1"));

    assert_eq!(final_msgs[3].role, Role::Assistant);
    assert!(
        final_msgs[3].reasoning_content.is_none(),
        "非工具调用 assistant 应剥离 reasoning"
    );
    assert_eq!(final_msgs[3].content, "北京今天晴，25°C");

    actor.stop_gracefully().await.unwrap();
}

#[tokio::test]
async fn send_returns_raw_to_caller() {
    let backend = RichMockBackend::new(vec![vec![RichMockMessage::assistant_with_reasoning(
        "答案",
        "思考过程",
    )]]);
    let actor = spawn_with_backend(backend);

    actor
        .ask(RequestAppend {
            message: RichMockMessage::user("问"),
        })
        .await
        .unwrap();

    // Send 返回原始响应——调用方拿到完整 reasoning_content
    let raw: RichMockResponse = actor
        .ask(RequestSend {
            opts: rich_mock_opts(0),
        })
        .await
        .unwrap();
    assert_eq!(raw.0[0].content, "答案");
    assert!(
        raw.0[0].reasoning_content.is_some(),
        "调用方应拿到原始响应（含 reasoning）"
    );

    // AC 内部存储的是转换后的版本（reasoning 被剥离，因为无 tool_calls）
    let stored: Vec<RichMockMessage> = actor.ask(RequestMessages).await.unwrap();
    assert!(stored[1].reasoning_content.is_none(), "AC 内部存储应已剥离");

    actor.stop_gracefully().await.unwrap();
}

#[tokio::test]
async fn send_stream_tool_call() {
    let backend = RichMockBackend::new(vec![vec![RichMockMessage::assistant_with_tool_calls(
        "",
        "调用工具",
        vec![ToolCall {
            id: "call_s".into(),
            function_name: "search".into(),
            arguments: r#"{"q":"test"}"#.into(),
        }],
    )]]);
    let actor = spawn_with_backend(backend);

    actor
        .ask(RequestAppend {
            message: RichMockMessage::user("搜索"),
        })
        .await
        .unwrap();

    // SendStream 返回 AgentSendStream，消费流以触发 Drop 存储
    {
        use futures::StreamExt;
        let mut stream = actor
            .ask(RequestSendStream {
                opts: rich_mock_opts(0),
            })
            .await
            .unwrap();
        let mut count = 0usize;
        while stream.next().await.is_some() {
            count += 1;
        }
        assert_eq!(count, 1, "Should have drained one stream item");
        for msg in stream.take_chunks().into_iter().flat_map(|r| r.0) {
            actor.ask(RequestAppend { message: msg }).await.unwrap();
        }
    }

    // 流式工具调用：reasoning_content 应保留（有 tool_calls）
    let stored: Vec<RichMockMessage> = actor.ask(RequestMessages).await.unwrap();
    assert_eq!(stored.len(), 2);
    assert!(stored[1].tool_calls.is_some());
    assert!(
        stored[1].reasoning_content.is_some(),
        "流式工具调用应保留 reasoning_content"
    );

    actor.stop_gracefully().await.unwrap();
}

#[tokio::test]
async fn opts_passthrough_to_backend() {
    let last_opts_token = Arc::new(std::sync::atomic::AtomicU64::new(0));
    let mut backend = RichMockBackend::new(vec![
        vec![RichMockMessage::assistant("R1")],
        vec![RichMockMessage::assistant("R2")],
    ]);
    backend.last_opts_token = last_opts_token.clone();

    let actor = spawn_with_backend(backend);

    actor
        .ask(RequestAppend {
            message: RichMockMessage::user("Q1"),
        })
        .await
        .unwrap();

    // Send：传入 token=100，验证后端收到
    let send_response = actor
        .ask(RequestSend {
            opts: rich_mock_opts(100),
        })
        .await
        .unwrap();
    assert!(!send_response.0.is_empty(), "Response should not be empty");
    assert_eq!(
        last_opts_token.load(Ordering::SeqCst),
        100,
        "Send 的 opts 应透传到 backend.send()"
    );

    // SendStream：传入 token=200，验证后端收到
    actor
        .ask(RequestSendStream {
            opts: rich_mock_opts(200),
        })
        .await
        .unwrap();
    tokio::time::sleep(std::time::Duration::from_millis(200)).await;
    assert_eq!(
        last_opts_token.load(Ordering::SeqCst),
        200,
        "SendStream 的 opts 应透传到 backend.send_stream()"
    );

    actor.stop_gracefully().await.unwrap();
}