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

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

# 加载环境变量
load_dotenv()

from alayaflow.component.llm_node import LLMComponent, ResponseFormat
# 导入真实的 ModelManager 类，以便获取单例实例
from alayaflow.component.model import ModelManager 


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

    @pytest.fixture
    def mock_llm_runnable(self):
        return Mock()

    @pytest.fixture
    def mock_get_model(self, mock_llm_runnable):
        # 获取真实的单例实例
        manager_instance = ModelManager()
        # Patch 该实例上的 get_model 方法
        with patch.object(manager_instance, 'get_model', return_value=mock_llm_runnable) as mock_method:
            yield mock_method

    @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_get_model, mock_llm_runnable, test_data):
        """测试 TEXT 格式响应"""
        scenario = test_data["text_response"]
        
        mock_llm_runnable.invoke.return_value = AIMessage(content=scenario["expected_content"])
        
        component = LLMComponent(
            model_id="deepseek-chat",
            system_prompt=scenario["system_prompt"],
            prompt=scenario["prompt"],
            response_format=scenario["response_format"],
        )
        
        result = component()
        
        assert result.content == scenario["expected_content"]
        
        mock_get_model.assert_called_with("deepseek-chat", runtime_config={})
    
        args = mock_llm_runnable.invoke.call_args[0][0]
        assert args[0].content == scenario["system_prompt"]
        assert args[1].content == scenario["prompt"]

    def test_llm_component_json_response(self, mock_get_model, mock_llm_runnable, test_data):
        """测试 JSON 格式"""
        scenario = test_data["json_response"]
        mock_llm_runnable.invoke.return_value = AIMessage(content=scenario["expected_content"])
        
        component = LLMComponent(
            model_id="deepseek-chat",
            system_prompt=scenario["system_prompt"],
            prompt=scenario["prompt"],
            response_format=scenario["response_format"],
        )
        
        result = component()
        assert result.content == scenario["expected_content"]
        
        args = mock_llm_runnable.invoke.call_args[0][0]
        assert "JSON" in args[0].content

    def test_llm_component_json_retry(self, mock_get_model, mock_llm_runnable, test_data):
        """测试 JSON 重试机制"""
        scenario = test_data["json_response_retry"]
        
        # 模拟两次调用
        mock_llm_runnable.invoke.side_effect = [
            AIMessage(content=scenario["first_response"]),
            AIMessage(content=scenario["expected_content"])
        ]
        
        component = LLMComponent(
            model_id="deepseek-chat",
            system_prompt=scenario["system_prompt"],
            prompt=scenario["prompt"],
            response_format=scenario["response_format"],
            retry_json_once=True
        )
        
        result = component()
        assert result.content == scenario["expected_content"]
        assert mock_llm_runnable.invoke.call_count == 2

    def test_llm_component_with_params(self, mock_get_model, mock_llm_runnable, test_data):
        """测试参数传递 (verify runtime_config)"""
        scenario = test_data["with_sampling_params"]
        mock_llm_runnable.invoke.return_value = AIMessage(content=scenario["expected_content"])
        
        component = LLMComponent(
            model_id="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"],
        )
        
        component()
        
        # 验证参数是否通过 runtime_config 正确传递给了 ModelManager 实例
        mock_get_model.assert_called_once()
        _, kwargs = mock_get_model.call_args
        config = kwargs["runtime_config"]
        
        assert config["temperature"] == scenario["temperature"]
        assert config["top_p"] == scenario["top_p"]
        assert config["max_tokens"] == scenario["max_tokens"]


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