import os
import sys
import json
from pathlib import Path
from unittest.mock import Mock, patch, MagicMock
from langchain_core.messages import AIMessage, SystemMessage, HumanMessage
from dotenv import load_dotenv

sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src"))
import pytest

# 加载 .env 文件
load_dotenv()

from alayaflow.component.llm_node import LLMComponent, ResponseFormat
from alayaflow.component.chat_model import mk_chat_model_deepseek


class TestLLMComponent:
    """测试LLMComponent类的测试类"""

    @pytest.fixture
    def mock_chat_openai(self):
        """创建一个模拟的ChatOpenAI模型"""
        with patch('alayaflow.component.llm_node.ChatOpenAI') as mock_chat:
            mock_instance = Mock()
            mock_chat.return_value = mock_instance
            yield mock_instance

    @pytest.fixture
    def test_data(self):
        """测试数据集"""
        return {
            "text_response": {
                "system_prompt": "你是一位精通 python 语言的资深架构师。",
                "prompt": "你是谁？",
                "response_format": ResponseFormat.TEXT,
                "expected_content": "我是一位精通 Python 语言的资深架构师。",
            },
            "json_response": {
                "system_prompt": "你是一个JSON输出助手。",
                "prompt": "返回一个包含name和age的JSON对象",
                "response_format": ResponseFormat.JSON,
                "expected_content": '{"name": "张三", "age": 25}',
            },
            "json_response_retry": {
                "system_prompt": "你是一个JSON输出助手。",
                "prompt": "返回一个包含name和age的JSON对象",
                "response_format": ResponseFormat.JSON,
                "first_response": "```json\n{\"name\": \"张三\", \"age\": 25}\n```",
                "expected_content": '{"name": "张三", "age": 25}',
            },
            "with_sampling_params": {
                "system_prompt": "你是一个助手。",
                "prompt": "你好",
                "response_format": ResponseFormat.TEXT,
                "temperature": 0.7,
                "top_p": 0.9,
                "max_tokens": 100,
                "expected_content": "你好！",
            },
        }

    def test_llm_component_text_response(self, mock_chat_openai, test_data):
        """测试TEXT格式响应"""
        scenario = test_data["text_response"]
        
        # 设置mock响应
        mock_llm_instance = Mock()
        mock_llm_instance.invoke = Mock(return_value=AIMessage(content=scenario["expected_content"]))
        mock_chat_openai.return_value = mock_llm_instance
        
        # 创建组件并调用
        component = LLMComponent(
            model_name="deepseek-chat",
            system_prompt=scenario["system_prompt"],
            prompt=scenario["prompt"],
            response_format=scenario["response_format"],
        )
        
        # 直接mock _get_llm方法返回mock实例
        component._get_llm = Mock(return_value=mock_llm_instance)
        
        result = component()
        
        # 验证结果
        assert isinstance(result, AIMessage), "result should be an instance of AIMessage"
        assert result.content == scenario["expected_content"]
        
        # 验证调用参数
        mock_llm_instance.invoke.assert_called_once()
        call_args = mock_llm_instance.invoke.call_args[0][0]
        assert len(call_args) == 2, "messages should have 2 elements"
        assert isinstance(call_args[0], SystemMessage), "First message should be SystemMessage"
        assert isinstance(call_args[1], HumanMessage), "Second message should be HumanMessage"
        assert call_args[0].content == scenario["system_prompt"]
        assert call_args[1].content == scenario["prompt"]

    def test_llm_component_json_response(self, mock_chat_openai, test_data):
        """测试JSON格式响应（成功）"""
        scenario = test_data["json_response"]
        
        # 设置mock响应
        mock_llm_instance = Mock()
        mock_llm_instance.invoke = Mock(return_value=AIMessage(content=scenario["expected_content"]))
        mock_chat_openai.return_value = mock_llm_instance
        
        # 创建组件并调用
        component = LLMComponent(
            model_name="deepseek-chat",
            system_prompt=scenario["system_prompt"],
            prompt=scenario["prompt"],
            response_format=scenario["response_format"],
        )
        
        # 直接mock _get_llm方法返回mock实例
        component._get_llm = Mock(return_value=mock_llm_instance)
        
        result = component()
        
        # 验证结果
        assert isinstance(result, AIMessage), "result should be an instance of AIMessage"
        assert result.content == scenario["expected_content"]
        
        # 验证JSON格式有效
        json.loads(result.content)
        
        # 验证调用参数 - JSON格式会在system_prompt后添加格式要求
        mock_llm_instance.invoke.assert_called_once()
        call_args = mock_llm_instance.invoke.call_args[0][0]
        assert len(call_args) == 2, "messages should have 2 elements"
        assert isinstance(call_args[0], SystemMessage), "First message should be SystemMessage"
        assert "JSON" in call_args[0].content, "System prompt should contain JSON format requirement"

    def test_llm_component_json_response_retry(self, mock_chat_openai, test_data):
        """测试JSON格式响应（需要重试）"""
        scenario = test_data["json_response_retry"]
        
        # 设置mock响应 - 第一次返回无效JSON，第二次返回有效JSON
        mock_llm_instance = Mock()
        call_count = [0]
        
        def mock_invoke(messages):
            call_count[0] += 1
            if call_count[0] == 1:
                return AIMessage(content=scenario["first_response"])
            else:
                return AIMessage(content=scenario["expected_content"])
        
        mock_llm_instance.invoke = Mock(side_effect=mock_invoke)
        mock_chat_openai.return_value = mock_llm_instance
        
        # 创建组件并调用
        component = LLMComponent(
            model_name="deepseek-chat",
            system_prompt=scenario["system_prompt"],
            prompt=scenario["prompt"],
            response_format=scenario["response_format"],
            retry_json_once=True,
        )
        
        # 直接mock _get_llm方法返回mock实例
        component._get_llm = Mock(return_value=mock_llm_instance)
        
        result = component()
        
        # 验证结果
        assert isinstance(result, AIMessage), "result should be an instance of AIMessage"
        assert result.content == scenario["expected_content"]
        
        # 验证JSON格式有效
        json.loads(result.content)
        
        # 验证调用了两次（第一次失败，第二次重试）
        assert mock_llm_instance.invoke.call_count == 2, "should be called twice for retry"
        
        # 验证第二次调用的system_prompt包含更严格的要求
        second_call_args = mock_llm_instance.invoke.call_args_list[1][0][0]
        assert "严格" in second_call_args[0].content or "strict" in second_call_args[0].content.lower(), \
            "Second call should have stricter JSON format requirement"

    def test_llm_component_with_sampling_params(self, mock_chat_openai, test_data):
        """测试带采样参数的调用"""
        scenario = test_data["with_sampling_params"]
        
        # 设置mock响应 - 模拟bind方法返回的对象
        mock_bind_instance = Mock()
        mock_bind_instance.invoke = Mock(return_value=AIMessage(content=scenario["expected_content"]))
        mock_llm_instance = Mock()
        mock_llm_instance.bind = Mock(return_value=mock_bind_instance)
        mock_chat_openai.return_value = mock_llm_instance
        
        # 创建组件并调用
        component = LLMComponent(
            model_name="deepseek-chat",
            system_prompt=scenario["system_prompt"],
            prompt=scenario["prompt"],
            response_format=scenario["response_format"],
            temperature=scenario["temperature"],
            top_p=scenario["top_p"],
            max_tokens=scenario["max_tokens"],
        )
        
        # 直接mock _get_llm方法，模拟bind的调用
        def mock_get_llm():
            bound_llm = mock_llm_instance.bind(
                temperature=scenario["temperature"],
                top_p=scenario["top_p"],
                max_tokens=scenario["max_tokens"]
            )
            return bound_llm
        
        component._get_llm = mock_get_llm
        
        result = component()
        
        # 验证结果
        assert isinstance(result, AIMessage), "result should be an instance of AIMessage"
        assert result.content == scenario["expected_content"]
        
        # 验证bind被调用，且参数正确
        mock_llm_instance.bind.assert_called_once()
        bind_kwargs = mock_llm_instance.bind.call_args[1]
        assert bind_kwargs["temperature"] == scenario["temperature"]
        assert bind_kwargs["top_p"] == scenario["top_p"]
        assert bind_kwargs["max_tokens"] == scenario["max_tokens"]
        
        # 验证invoke被调用
        mock_bind_instance.invoke.assert_called_once()

    def test_llm_component_json_no_retry(self, mock_chat_openai, test_data):
        """测试JSON格式响应（不重试）"""
        scenario = test_data["json_response_retry"]
        
        # 设置mock响应 - 返回无效JSON
        mock_llm_instance = Mock()
        mock_llm_instance.invoke = Mock(return_value=AIMessage(content=scenario["first_response"]))
        mock_chat_openai.return_value = mock_llm_instance
        
        # 创建组件并调用（retry_json_once=False）
        component = LLMComponent(
            model_name="deepseek-chat",
            system_prompt=scenario["system_prompt"],
            prompt=scenario["prompt"],
            response_format=scenario["response_format"],
            retry_json_once=False,
        )
        
        # 直接mock _get_llm方法返回mock实例
        component._get_llm = Mock(return_value=mock_llm_instance)
        
        result = component()
        
        # 验证结果 - 应该返回第一次的响应（即使不是有效JSON）
        assert isinstance(result, AIMessage), "result should be an instance of AIMessage"
        assert result.content == scenario["first_response"]
        
        # 验证只调用了一次（不重试）
        assert mock_llm_instance.invoke.call_count == 1, "should be called only once when retry is disabled"

    @pytest.mark.skipif(
        not os.getenv("ENABLE_REAL_LLM_TEST"),
        reason="未开启真实LLM测试开关（设置ENABLE_REAL_LLM_TEST=true可运行）"
    )
    def test_llm_component_real_llm_text(self, test_data):
        """测试真实LLM调用（TEXT格式）"""
        scenario = test_data["text_response"]
        
        # 创建组件并调用
        component = LLMComponent(
            model_name="deepseek-chat",
            system_prompt=scenario["system_prompt"],
            prompt=scenario["prompt"],
            response_format=scenario["response_format"],
        )
        
        result = component()
        
        # 验证结果
        assert isinstance(result, AIMessage), "result should be an instance of AIMessage"
        assert len(result.content) > 0, "response content should not be empty"

    @pytest.mark.skipif(
        not os.getenv("ENABLE_REAL_LLM_TEST"),
        reason="未开启真实LLM测试开关（设置ENABLE_REAL_LLM_TEST=true可运行）"
    )
    def test_llm_component_real_llm_json(self, test_data):
        """测试真实LLM调用（JSON格式）"""
        scenario = test_data["json_response"]
        
        # 创建组件并调用
        component = LLMComponent(
            model_name="deepseek-chat",
            system_prompt=scenario["system_prompt"],
            prompt=scenario["prompt"],
            response_format=scenario["response_format"],
        )
        
        result = component()
        
        # 验证结果
        assert isinstance(result, AIMessage), "result should be an instance of AIMessage"
        assert len(result.content) > 0, "response content should not be empty"
        
        # 验证JSON格式有效
        try:
            json.loads(result.content)
        except json.JSONDecodeError:
            pytest.fail("Response should be valid JSON")


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
