//! Mock LLM Provider — 用于测试 Agent 子系统，不需要真实 API

use crate::agent::llm::{ChatMessage, LlmClient, LlmError};
use crate::types::{LlmResponse, LlmToolCall, ToolDefinition};
use serde_json::json;

/// Mock 响应配置
#[derive(Clone, Debug)]
pub struct MockResponse {
    pub content: Option<String>,
    pub tool_calls: Vec<MockToolCall>,
}

#[derive(Clone, Debug)]
pub struct MockToolCall {
    pub name: String,
    pub arguments: serde_json::Value,
}

/// Mock LLM 客户端 — 返回预设的结构化输出
pub struct MockLlmClient {
    responses: Vec<MockResponse>,
    call_count: std::cell::Cell<usize>,
}

impl MockLlmClient {
    pub fn new(responses: Vec<MockResponse>) -> Self {
        Self {
            responses,
            call_count: std::cell::Cell::new(0),
        }
    }

    /// 创建一个 Reason 模式的 mock — 返回 Intent 建议
    pub fn reason_mock() -> Self {
        Self::new(vec![
            MockResponse {
                content: Some(
                    r#"我分析了当前攻击路径图，发现了几个关键的探索方向：

1. 目标机器 (192.168.1.100) 可能存在 SMB 服务漏洞
2. 从 Web 服务器到目标机器的路径尚未建立
3. 需要验证域控制器上的凭证有效性

```json
{"intents": [
  {"direction": "扫描目标机器 SMB 服务", "target_node": "target-node-id", "priority": "high"},
  {"direction": "建立 Web 服务器到目标的横向路径", "target_node": "web-server-id", "priority": "medium"}
]}
```"#
                    .to_string(),
                ),
                tool_calls: vec![],
            },
        ])
    }

    /// 创建一个 Explore 模式的 mock — 返回执行结果
    pub fn explore_mock() -> Self {
        Self::new(vec![
            MockResponse {
                content: Some(
                    r#"我执行了 SMB 服务扫描，发现了以下信息：

目标机器 192.168.1.100 开放了以下端口：
- 445/tcp: SMBv1/v2/v3
- 139/tcp: NetBIOS

发现了一个可能的漏洞：MS17-010 (EternalBlue)

```json
{"status": "completed", "summary": "发现目标机器存在 MS17-010 漏洞", "nodes_created": [], "edges_created": []}
```"#
                    .to_string(),
                ),
                tool_calls: vec![],
            },
        ])
    }

    /// 创建一个 Tell 模式的 mock — 返回节点和边创建
    pub fn tell_mock() -> Self {
        Self::new(vec![
            MockResponse {
                content: None,
                tool_calls: vec![
                    MockToolCall {
                        name: "add_node".to_string(),
                        arguments: json!({
                            "label": "SMB 服务",
                            "node_type": "process",
                            "properties": {"port": 445, "protocol": "tcp"}
                        }),
                    },
                    MockToolCall {
                        name: "add_edge".to_string(),
                        arguments: json!({
                            "source_node_id": "origin",
                            "target_node_id": "smb-node",
                            "label": "连接"
                        }),
                    },
                ],
            },
            MockResponse {
                content: Some("已添加 SMB 服务节点和连接。".to_string()),
                tool_calls: vec![],
            },
        ])
    }

    pub fn call_count(&self) -> usize {
        self.call_count.get()
    }
}

/// 将 MockLlmClient 适配为 LlmClient 接口
/// 注意：这是通过直接构造 LlmResponse 来绕过实际 HTTP 调用
impl MockLlmClient {
    pub async fn mock_chat(
        &self,
        _messages: Vec<ChatMessage>,
        _tools: Option<Vec<ToolDefinition>>,
        _temperature: f64,
        _max_tokens: u32,
    ) -> Result<LlmResponse, LlmError> {
        let count = self.call_count.get();
        if count >= self.responses.len() {
            // 超出预设响应数量，返回空响应
            return Ok(LlmResponse {
                content: Some("Mock: 无更多预设响应".to_string()),
                tool_calls: vec![],
                finish_reason: "stop".to_string(),
            });
        }

        let response = &self.responses[count];
        self.call_count.set(count + 1);

        let tool_calls: Vec<LlmToolCall> = response
            .tool_calls
            .iter()
            .enumerate()
            .map(|(i, tc)| LlmToolCall {
                id: format!("mock_call_{}", i),
                name: tc.name.clone(),
                arguments: tc.arguments.clone(),
            })
            .collect();

        Ok(LlmResponse {
            content: response.content.clone(),
            tool_calls,
            finish_reason: "stop".to_string(),
        })
    }
}

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

    #[tokio::test]
    async fn test_mock_reason_returns_intents() {
        let mock = MockLlmClient::reason_mock();
        let response = mock
            .mock_chat(vec![], None, 0.7, 4096)
            .await
            .unwrap();

        assert!(response.content.is_some());
        let content = response.content.unwrap();
        assert!(content.contains("intents"));
        assert!(content.contains("SMB"));
        assert_eq!(mock.call_count(), 1);
    }

    #[tokio::test]
    async fn test_mock_explore_returns_status() {
        let mock = MockLlmClient::explore_mock();
        let response = mock
            .mock_chat(vec![], None, 0.7, 4096)
            .await
            .unwrap();

        assert!(response.content.is_some());
        let content = response.content.unwrap();
        assert!(content.contains("completed"));
        assert!(content.contains("MS17-010"));
    }

    #[tokio::test]
    async fn test_mock_tell_returns_tool_calls() {
        let mock = MockLlmClient::tell_mock();
        let response = mock
            .mock_chat(vec![], None, 0.7, 4096)
            .await
            .unwrap();

        assert!(!response.tool_calls.is_empty());
        assert_eq!(response.tool_calls[0].name, "add_node");
        assert_eq!(response.tool_calls[1].name, "add_edge");

        // 第二次调用返回文本
        let response2 = mock
            .mock_chat(vec![], None, 0.7, 4096)
            .await
            .unwrap();
        assert!(response2.content.is_some());
        assert!(response2.tool_calls.is_empty());
    }

    #[tokio::test]
    async fn test_mock_tracks_call_count() {
        let mock = MockLlmClient::new(vec![
            MockResponse {
                content: Some("first".to_string()),
                tool_calls: vec![],
            },
            MockResponse {
                content: Some("second".to_string()),
                tool_calls: vec![],
            },
        ]);

        assert_eq!(mock.call_count(), 0);
        let _ = mock.mock_chat(vec![], None, 0.7, 4096).await;
        assert_eq!(mock.call_count(), 1);
        let _ = mock.mock_chat(vec![], None, 0.7, 4096).await;
        assert_eq!(mock.call_count(), 2);
    }
}
