import os
from unittest.mock import Mock, patch
from langchain_core.documents import Document
import pytest

from alayaflow.component.rerank_node import JinaRerankerComponent


class TestJinaRerankerComponent:
    """测试JinaRerankerComponent类的测试类"""

    @pytest.fixture
    def mock_reranker(self):
        """创建一个模拟的JinaRerank模型"""
        mock_reranker = Mock()
        mock_reranker.compress_documents = Mock()
        return mock_reranker

    @pytest.fixture
    def mock_model_manager(self, mock_reranker):
        """创建一个模拟的ModelManager"""
        mock_manager = Mock()
        mock_manager.get_model = Mock(return_value=mock_reranker)
        return mock_manager

    @pytest.fixture
    def test_data(self):
        """测试数据集"""
        return {
            "scenario_1": {
                "query": "如何冲泡一杯好喝的绿茶？",
                "docs": [
                    "冲泡美味绿茶的关键是控制水温。最好使用80-85摄氏度的热水，避免使用沸腾的100度开水，以免烫伤茶叶导致苦涩。先将茶叶放入杯中，注入热水，浸泡约2-3分钟后即可饮用。",
                    "绿茶的口感与水质、茶具都有关系。建议选用软水（如矿泉水或过滤水），并使用陶瓷或玻璃杯具冲泡，这样能更好地激发绿茶的清香。",
                    "绿茶的种类繁多，如龙井、碧螺春、毛峰等。不同产地的绿茶风味各异，建议多尝试找到自己喜欢的品种。",
                ],
                "expected_reranked": [
                    Document(page_content="冲泡美味绿茶的关键是控制水温。最好使用80-85摄氏度的热水，避免使用沸腾的100度开水，以免烫伤茶叶导致苦涩。先将茶叶放入杯中，注入热水，浸泡约2-3分钟后即可饮用。"),
                    Document(page_content="绿茶的口感与水质、茶具都有关系。建议选用软水（如矿泉水或过滤水），并使用陶瓷或玻璃杯具冲泡，这样能更好地激发绿茶的清香。"),
                    Document(page_content="绿茶的种类繁多，如龙井、碧螺春、毛峰等。不同产地的绿茶风味各异，建议多尝试找到自己喜欢的品种。"),
                ],
            },
            "scenario_2": {
                "query": "测试查询",
                "docs": [
                    Document(page_content="文档1", metadata={"id": 1}),
                    Document(page_content="文档2", metadata={"id": 2}),
                ],
                "expected_reranked": [
                    Document(page_content="文档2", metadata={"id": 2}),
                    Document(page_content="文档1", metadata={"id": 1}),
                ],
            },
            "scenario_3": {
                "query": "测试查询",
                "docs": [
                    {"page_content": "字典文档1", "metadata": {"source": "test1"}},
                    {"content": "字典文档2", "metadata": {"source": "test2"}},
                ],
                "expected_reranked": [
                    Document(page_content="字典文档2", metadata={"source": "test2"}),
                    Document(page_content="字典文档1", metadata={"source": "test1"}),
                ],
            },
        }

    def test_rerank_with_empty_query(self, mock_reranker, mock_model_manager, test_data):
        """测试空查询的情况"""
        scenario = test_data["scenario_1"]
        query = ""
        docs = scenario["docs"]

        with patch("alayaflow.component.rerank_node.ModelManager", return_value=mock_model_manager):
            component = JinaRerankerComponent(model_id="test-reranker", top_n=None)
            result = component(query=query, docs=docs)

        # 验证结果：空查询时应该直接返回原始文档
        assert len(result) == len(docs)
        # 验证reranker没有被调用
        mock_reranker.compress_documents.assert_not_called()

    def test_rerank_with_empty_docs(self, mock_reranker, mock_model_manager, test_data):
        """测试空文档列表的情况"""
        scenario = test_data["scenario_1"]
        query = scenario["query"]
        docs = []

        with patch("alayaflow.component.rerank_node.ModelManager", return_value=mock_model_manager):
            component = JinaRerankerComponent(model_id="test-reranker", top_n=None)
            result = component(query=query, docs=docs)

        # 验证结果：空文档时应该直接返回空列表
        assert result == []
        # 验证reranker没有被调用
        mock_reranker.compress_documents.assert_not_called()

    def test_rerank_successful(self, mock_reranker, mock_model_manager, test_data):
        """测试成功rerank的情况"""
        scenario = test_data["scenario_1"]
        query = scenario["query"]
        docs = scenario["docs"]
        expected_reranked = scenario["expected_reranked"]

        # 设置mock返回值
        mock_reranker.compress_documents.return_value = expected_reranked

        with patch("alayaflow.component.rerank_node.ModelManager", return_value=mock_model_manager):
            component = JinaRerankerComponent(model_id="test-reranker", top_n=None)
            result = component(query=query, docs=docs)

        # 验证结果
        assert result == expected_reranked
        assert len(result) == len(expected_reranked)
        # 验证reranker被调用
        mock_reranker.compress_documents.assert_called_once()
        
        # 验证调用参数
        call_args = mock_reranker.compress_documents.call_args
        call_kwargs = call_args[1]  # 关键字参数
        assert "documents" in call_kwargs
        assert "query" in call_kwargs
        assert call_kwargs["query"] == query
        assert len(call_kwargs["documents"]) == len(docs)

    def test_rerank_with_top_n(self, mock_reranker, mock_model_manager, test_data):
        """测试使用top_n参数的情况"""
        scenario = test_data["scenario_1"]
        query = scenario["query"]
        docs = scenario["docs"]
        top_n = 2
        expected_reranked = scenario["expected_reranked"][:top_n]

        # 设置mock返回值
        mock_reranker.compress_documents.return_value = expected_reranked

        with patch("alayaflow.component.rerank_node.ModelManager", return_value=mock_model_manager):
            component = JinaRerankerComponent(model_id="test-reranker", top_n=top_n)
            result = component(query=query, docs=docs)

        # 验证结果
        assert len(result) == top_n
        # 验证ModelManager.get_model被调用时传入了top_n
        mock_model_manager.get_model.assert_called_once()
        call_args = mock_model_manager.get_model.call_args
        assert call_args[0][0] == "test-reranker"  # model_id
        assert call_args[1]["runtime_config"]["top_n"] == top_n

    def test_to_documents_with_strings(self, mock_reranker, mock_model_manager):
        """测试_to_documents方法处理字符串输入"""
        docs = ["文档1", "文档2", "文档3"]
        
        with patch("alayaflow.component.rerank_node.ModelManager", return_value=mock_model_manager):
            component = JinaRerankerComponent(model_id="test-reranker", top_n=None)
            result = component._to_documents(docs)

        # 验证结果
        assert len(result) == 3
        assert all(isinstance(d, Document) for d in result)
        assert result[0].page_content == "文档1"
        assert result[1].page_content == "文档2"
        assert result[2].page_content == "文档3"
        assert all(d.metadata == {} for d in result)

    def test_to_documents_with_documents(self, mock_reranker, mock_model_manager):
        """测试_to_documents方法处理Document对象"""
        docs = [
            Document(page_content="文档1", metadata={"id": 1}),
            Document(page_content="文档2", metadata={"id": 2}),
        ]
        
        with patch("alayaflow.component.rerank_node.ModelManager", return_value=mock_model_manager):
            component = JinaRerankerComponent(model_id="test-reranker", top_n=None)
            result = component._to_documents(docs)

        # 验证结果：应该直接返回相同的Document对象
        assert len(result) == 2
        assert result[0] == docs[0]
        assert result[1] == docs[1]

    def test_to_documents_with_dicts(self, mock_reranker, mock_model_manager):
        """测试_to_documents方法处理字典输入"""
        docs = [
            {"page_content": "字典文档1", "metadata": {"source": "test1"}},
            {"content": "字典文档2", "metadata": {"source": "test2"}},
            {"page_content": "", "metadata": {}},  # 空内容
        ]
        
        with patch("alayaflow.component.rerank_node.ModelManager", return_value=mock_model_manager):
            component = JinaRerankerComponent(model_id="test-reranker", top_n=None)
            result = component._to_documents(docs)

        # 验证结果
        assert len(result) == 3
        assert all(isinstance(d, Document) for d in result)
        assert result[0].page_content == "字典文档1"
        assert result[0].metadata == {"source": "test1"}
        assert result[1].page_content == "字典文档2"  # 支持content字段
        assert result[1].metadata == {"source": "test2"}
        assert result[2].page_content == ""
        assert result[2].metadata == {}

    def test_rerank_multiple_scenarios(self, mock_reranker, mock_model_manager, test_data):
        """测试多个场景"""
        scenarios_to_test = ["scenario_1", "scenario_2", "scenario_3"]

        for scenario_name in scenarios_to_test:
            scenario = test_data[scenario_name]
            query = scenario["query"]
            docs = scenario["docs"]
            expected_reranked = scenario["expected_reranked"]

            # 重置mock
            mock_reranker.reset_mock()
            mock_model_manager.reset_mock()

            # 设置模拟响应
            mock_reranker.compress_documents.return_value = expected_reranked

            with patch("alayaflow.component.rerank_node.ModelManager", return_value=mock_model_manager):
                component = JinaRerankerComponent(model_id="test-reranker", top_n=None)
                result = component(query=query, docs=docs)

            # 验证结果
            assert result == expected_reranked
            assert len(result) == len(expected_reranked)

            # 验证reranker被调用
            mock_reranker.compress_documents.assert_called_once()

            # 验证调用参数
            call_args = mock_reranker.compress_documents.call_args
            call_kwargs = call_args[1]
            assert call_kwargs["query"] == query
            assert len(call_kwargs["documents"]) == len(docs)

    def test_get_reranker_with_top_n(self, mock_reranker, mock_model_manager):
        """测试_get_reranker方法正确传递top_n参数"""
        top_n = 5
        
        with patch("alayaflow.component.rerank_node.ModelManager", return_value=mock_model_manager):
            component = JinaRerankerComponent(model_id="test-reranker", top_n=top_n)
            result = component._get_reranker()

        # 验证结果
        assert result == mock_reranker
        # 验证ModelManager.get_model被调用时传入了top_n
        mock_model_manager.get_model.assert_called_once()
        call_args = mock_model_manager.get_model.call_args
        assert call_args[0][0] == "test-reranker"
        assert call_args[1]["runtime_config"]["top_n"] == top_n

    def test_get_reranker_without_top_n(self, mock_reranker, mock_model_manager):
        """测试_get_reranker方法不传递top_n参数"""
        with patch("alayaflow.component.rerank_node.ModelManager", return_value=mock_model_manager):
            component = JinaRerankerComponent(model_id="test-reranker", top_n=None)
            result = component._get_reranker()

        # 验证结果
        assert result == mock_reranker
        # 验证ModelManager.get_model被调用时没有top_n
        mock_model_manager.get_model.assert_called_once()
        call_args = mock_model_manager.get_model.call_args
        assert call_args[0][0] == "test-reranker"
        assert call_args[1]["runtime_config"] == {}

    @pytest.mark.skipif(
        not os.getenv("ENABLE_REAL_LLM_TEST"),
        reason="未开启真实LLM测试开关（设置ENABLE_REAL_LLM_TEST=true可运行）"
    )
    def test_rerank_real_model(self, test_data):
        """测试使用真实模型的情况"""
        from alayaflow.api import Flow
        from alayaflow.component.model.schemas import ModelProfile
        from pydantic import SecretStr
        
        scenario = test_data["scenario_1"]
        query = scenario["query"]
        docs = scenario["docs"]

        # 检查环境变量
        jina_api_key = os.getenv("JINA_API_KEY")
        if not jina_api_key:
            pytest.skip("需要设置环境变量 JINA_API_KEY 才能运行真实模型测试")

        # 初始化Flow并注册模型
        flow = Flow()
        flow.init({
            "alayahub_url": "http://your-alayahub-url",
        })
        
        flow.register_models([
            {
                "name": "Jina Rerank",
                "model_id": "jina-reranker",
                "provider_name": "Jina",
                "model_type": "JinaRerank",
                "model_name": "jina-reranker-v3",
                "base_url": "null",
                "api_key": jina_api_key,
            }
        ])

        # 创建组件并测试
        component = JinaRerankerComponent(model_id="jina-reranker", top_n=3)
        result = component(query=query, docs=docs)

        # 验证结果
        assert isinstance(result, list)
        assert len(result) > 0
        assert all(isinstance(d, Document) for d in result)
        assert len(result) <= 3  # top_n=3


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

