"""
Test cases for Step-KB Binding feature.
[Step-KB-Binding] 步骤关联知识库功能的测试用例

测试覆盖：
1. 步骤配置的增删改查与参数校验
2. 配置步骤KB范围后，检索仅在范围内发生
3. 未配置时回退原逻辑
4. 多知识库同名内容场景下，步骤约束生效
5. trace 字段完整性校验
"""

import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from dataclasses import dataclass
from typing import Any


class TestStepKbBindingModel:
    """测试步骤KB绑定数据模型"""
    
    def test_flow_step_with_kb_binding_fields(self):
        """测试 FlowStep 包含 KB 绑定字段"""
        from app.models.entities import FlowStep
        
        step = FlowStep(
            step_no=1,
            content="测试步骤",
            allowed_kb_ids=["kb-1", "kb-2"],
            preferred_kb_ids=["kb-1"],
            kb_query_hint="查找产品相关信息",
            max_kb_calls_per_step=2,
        )
        
        assert step.allowed_kb_ids == ["kb-1", "kb-2"]
        assert step.preferred_kb_ids == ["kb-1"]
        assert step.kb_query_hint == "查找产品相关信息"
        assert step.max_kb_calls_per_step == 2
    
    def test_flow_step_without_kb_binding(self):
        """测试 FlowStep 不配置 KB 绑定时的默认值"""
        from app.models.entities import FlowStep
        
        step = FlowStep(
            step_no=1,
            content="测试步骤",
        )
        
        assert step.allowed_kb_ids is None
        assert step.preferred_kb_ids is None
        assert step.kb_query_hint is None
        assert step.max_kb_calls_per_step is None
    
    def test_max_kb_calls_validation(self):
        """测试 max_kb_calls_per_step 的范围校验"""
        from app.models.entities import FlowStep
        from pydantic import ValidationError
        
        # 有效范围 1-5
        step = FlowStep(step_no=1, content="test", max_kb_calls_per_step=3)
        assert step.max_kb_calls_per_step == 3
        
        # 超出上限
        with pytest.raises(Exception):  # ValidationError
            FlowStep(step_no=1, content="test", max_kb_calls_per_step=10)


class TestStepKbConfig:
    """测试 StepKbConfig 数据类"""
    
    def test_step_kb_config_creation(self):
        """测试 StepKbConfig 创建"""
        from app.services.mid.kb_search_dynamic_tool import StepKbConfig
        
        config = StepKbConfig(
            allowed_kb_ids=["kb-1", "kb-2"],
            preferred_kb_ids=["kb-1"],
            kb_query_hint="查找产品信息",
            max_kb_calls=2,
            step_id="flow-1_step_1",
        )
        
        assert config.allowed_kb_ids == ["kb-1", "kb-2"]
        assert config.preferred_kb_ids == ["kb-1"]
        assert config.kb_query_hint == "查找产品信息"
        assert config.max_kb_calls == 2
        assert config.step_id == "flow-1_step_1"
    
    def test_step_kb_config_defaults(self):
        """测试 StepKbConfig 默认值"""
        from app.services.mid.kb_search_dynamic_tool import StepKbConfig
        
        config = StepKbConfig()
        
        assert config.allowed_kb_ids is None
        assert config.preferred_kb_ids is None
        assert config.kb_query_hint is None
        assert config.max_kb_calls == 1
        assert config.step_id is None


class TestKbSearchDynamicToolWithStepConfig:
    """测试 KbSearchDynamicTool 与步骤配置的集成"""
    
    @pytest.mark.asyncio
    async def test_kb_search_with_allowed_kb_ids(self):
        """测试配置 allowed_kb_ids 后检索范围受限"""
        from app.services.mid.kb_search_dynamic_tool import (
            KbSearchDynamicTool,
            KbSearchDynamicConfig,
            StepKbConfig,
        )
        
        mock_session = MagicMock()
        mock_timeout_governor = MagicMock()
        
        tool = KbSearchDynamicTool(
            session=mock_session,
            timeout_governor=mock_timeout_governor,
            config=KbSearchDynamicConfig(enabled=True),
        )
        
        step_config = StepKbConfig(
            allowed_kb_ids=["kb-allowed-1", "kb-allowed-2"],
            step_id="test_step",
        )
        
        with patch.object(tool, '_do_retrieve', new_callable=AsyncMock) as mock_retrieve:
            mock_retrieve.return_value = [
                {"id": "1", "content": "test", "score": 0.8, "metadata": {"kb_id": "kb-allowed-1"}}
            ]
            
            result = await tool.execute(
                query="测试查询",
                tenant_id="tenant-1",
                step_kb_config=step_config,
            )
            
            # 验证检索调用时传入了正确的 kb_ids
            call_args = mock_retrieve.call_args
            assert call_args[1]['step_kb_config'] == step_config
            
            # 验证返回结果包含 step_kb_binding 信息
            assert result.step_kb_binding is not None
            assert result.step_kb_binding['allowed_kb_ids'] == ["kb-allowed-1", "kb-allowed-2"]
    
    @pytest.mark.asyncio
    async def test_kb_search_without_step_config(self):
        """测试未配置步骤KB时的回退行为"""
        from app.services.mid.kb_search_dynamic_tool import (
            KbSearchDynamicTool,
            KbSearchDynamicConfig,
        )
        
        mock_session = MagicMock()
        mock_timeout_governor = MagicMock()
        
        tool = KbSearchDynamicTool(
            session=mock_session,
            timeout_governor=mock_timeout_governor,
            config=KbSearchDynamicConfig(enabled=True),
        )
        
        with patch.object(tool, '_do_retrieve', new_callable=AsyncMock) as mock_retrieve:
            mock_retrieve.return_value = [
                {"id": "1", "content": "test", "score": 0.8, "metadata": {}}
            ]
            
            result = await tool.execute(
                query="测试查询",
                tenant_id="tenant-1",
            )
            
            # 验证检索调用时未传入 step_kb_config
            call_args = mock_retrieve.call_args
            assert call_args[1]['step_kb_config'] is None
            
            # 验证返回结果不包含 step_kb_binding
            assert result.step_kb_binding is None
    
    @pytest.mark.asyncio
    async def test_kb_search_result_includes_used_kb_ids(self):
        """测试检索结果包含实际使用的知识库ID"""
        from app.services.mid.kb_search_dynamic_tool import (
            KbSearchDynamicTool,
            KbSearchDynamicConfig,
            StepKbConfig,
        )
        
        mock_session = MagicMock()
        mock_timeout_governor = MagicMock()
        
        tool = KbSearchDynamicTool(
            session=mock_session,
            timeout_governor=mock_timeout_governor,
            config=KbSearchDynamicConfig(enabled=True),
        )
        
        step_config = StepKbConfig(
            allowed_kb_ids=["kb-1", "kb-2"],
            step_id="test_step",
        )
        
        with patch.object(tool, '_do_retrieve', new_callable=AsyncMock) as mock_retrieve:
            mock_retrieve.return_value = [
                {"id": "1", "content": "test1", "score": 0.9, "metadata": {"kb_id": "kb-1"}},
                {"id": "2", "content": "test2", "score": 0.8, "metadata": {"kb_id": "kb-1"}},
                {"id": "3", "content": "test3", "score": 0.7, "metadata": {"kb_id": "kb-2"}},
            ]
            
            result = await tool.execute(
                query="测试查询",
                tenant_id="tenant-1",
                step_kb_config=step_config,
            )
            
            # 验证 used_kb_ids 包含所有命中的知识库
            assert result.step_kb_binding is not None
            assert set(result.step_kb_binding['used_kb_ids']) == {"kb-1", "kb-2"}
            assert result.step_kb_binding['kb_hit'] is True


class TestTraceInfoStepKbBinding:
    """测试 TraceInfo 中的 step_kb_binding 字段"""
    
    def test_trace_info_with_step_kb_binding(self):
        """测试 TraceInfo 包含 step_kb_binding 字段"""
        from app.models.mid.schemas import TraceInfo, ExecutionMode
        
        trace = TraceInfo(
            mode=ExecutionMode.AGENT,
            step_kb_binding={
                "step_id": "flow-1_step_2",
                "allowed_kb_ids": ["kb-1", "kb-2"],
                "used_kb_ids": ["kb-1"],
                "kb_hit": True,
            },
        )
        
        assert trace.step_kb_binding is not None
        assert trace.step_kb_binding['step_id'] == "flow-1_step_2"
        assert trace.step_kb_binding['allowed_kb_ids'] == ["kb-1", "kb-2"]
        assert trace.step_kb_binding['used_kb_ids'] == ["kb-1"]
    
    def test_trace_info_without_step_kb_binding(self):
        """测试 TraceInfo 默认不包含 step_kb_binding"""
        from app.models.mid.schemas import TraceInfo, ExecutionMode
        
        trace = TraceInfo(mode=ExecutionMode.AGENT)
        
        assert trace.step_kb_binding is None


class TestFlowStepKbBindingIntegration:
    """测试流程步骤与KB绑定的集成"""
    
    def test_script_flow_steps_with_kb_binding(self):
        """测试 ScriptFlow 的 steps 包含 KB 绑定配置"""
        from app.models.entities import ScriptFlowCreate
        
        flow_create = ScriptFlowCreate(
            name="测试流程",
            steps=[
                {
                    "step_no": 1,
                    "content": "步骤1",
                    "allowed_kb_ids": ["kb-1"],
                    "preferred_kb_ids": None,
                    "kb_query_hint": "查找产品信息",
                    "max_kb_calls_per_step": 2,
                },
                {
                    "step_no": 2,
                    "content": "步骤2",
                    # 不配置 KB 绑定
                },
            ],
        )
        
        assert flow_create.steps[0]['allowed_kb_ids'] == ["kb-1"]
        assert flow_create.steps[0]['kb_query_hint'] == "查找产品信息"
        assert flow_create.steps[1].get('allowed_kb_ids') is None


class TestKbBindingLogging:
    """测试 KB 绑定的日志记录"""
    
    @pytest.mark.asyncio
    async def test_step_kb_config_logging(self, caplog):
        """测试步骤KB配置的日志记录"""
        import logging
        from app.services.mid.kb_search_dynamic_tool import (
            KbSearchDynamicTool,
            KbSearchDynamicConfig,
            StepKbConfig,
        )
        
        mock_session = MagicMock()
        mock_timeout_governor = MagicMock()
        
        tool = KbSearchDynamicTool(
            session=mock_session,
            timeout_governor=mock_timeout_governor,
            config=KbSearchDynamicConfig(enabled=True),
        )
        
        step_config = StepKbConfig(
            allowed_kb_ids=["kb-1"],
            step_id="flow-1_step_1",
        )
        
        with patch.object(tool, '_do_retrieve', new_callable=AsyncMock) as mock_retrieve:
            mock_retrieve.return_value = []
            
            with caplog.at_level(logging.INFO):
                await tool.execute(
                    query="测试",
                    tenant_id="tenant-1",
                    step_kb_config=step_config,
                )
        
        # 验证日志包含 Step-KB-Binding 标记
        assert any("Step-KB-Binding" in record.message for record in caplog.records)


# 运行测试的入口
if __name__ == "__main__":
    pytest.main([__file__, "-v"])