"""
测试工具 Schema 转换修复

验证修复后的工具 Schema 转换是否正确
"""

import pytest
import json
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))

from tools.base import BaseTool
from tools.filesystem import ReadTool, WriteTool
from pydantic import BaseModel, Field


class TestToolSchemaFix:
    """测试工具 Schema 修复"""

    def test_read_tool_schema(self):
        """测试 ReadTool 的 Schema 转换"""
        tool = ReadTool()

        # 生成原始 Pydantic schema
        raw_schema = tool.args_schema.model_json_schema()
        print(f"\nRaw Pydantic Schema for ReadTool:")
        print(json.dumps(raw_schema, indent=2))

        # 清理 schema（模拟修复后的逻辑）
        clean_schema = {
            "type": raw_schema.get("type", "object"),
            "properties": raw_schema.get("properties", {}),
            "required": raw_schema.get("required", []),
        }

        print(f"\nClean Anthropic Schema for ReadTool:")
        print(json.dumps(clean_schema, indent=2))

        # 验证
        assert clean_schema["type"] == "object"
        assert "properties" in clean_schema
        assert "required" in clean_schema
        assert "title" not in clean_schema

        # 验证必需字段
        required = clean_schema["required"]
        assert "file_path" in required

    def test_write_tool_schema(self):
        """测试 WriteTool 的 Schema 转换"""
        tool = WriteTool()

        # 生成原始 Pydantic schema
        raw_schema = tool.args_schema.model_json_schema()

        # 清理 schema
        clean_schema = {
            "type": raw_schema.get("type", "object"),
            "properties": raw_schema.get("properties", {}),
            "required": raw_schema.get("required", []),
        }

        print(f"\nClean Anthropic Schema for WriteTool:")
        print(json.dumps(clean_schema, indent=2))

        # 验证
        assert len(clean_schema["required"]) == 2
        assert "file_path" in clean_schema["required"]
        assert "content" in clean_schema["required"]

    def test_anthropic_tool_definition_format(self):
        """测试 Anthropic 工具定义格式"""
        from core.llm.base import ToolDefinition

        tool = ReadTool()

        # 生成 clean schema
        raw_schema = tool.args_schema.model_json_schema()
        clean_schema = {
            "type": raw_schema.get("type", "object"),
            "properties": raw_schema.get("properties", {}),
            "required": raw_schema.get("required", []),
        }

        # 创建 ToolDefinition
        tool_def = ToolDefinition(
            name=tool.name,
            description=tool.description,
            parameters=clean_schema,
        )

        # 转换为 Anthropic 格式
        anthropic_format = {
            "name": tool_def.name,
            "description": tool_def.description,
            "input_schema": tool_def.parameters,
        }

        print(f"\nAnthropic Tool Definition:")
        print(json.dumps(anthropic_format, indent=2))

        # 验证
        assert anthropic_format["name"] == "read"
        assert "input_schema" in anthropic_format
        assert anthropic_format["input_schema"]["type"] == "object"
        assert "title" not in anthropic_format["input_schema"]

    def test_multiple_tools_anthropic_format(self):
        """测试多个工具的 Anthropic 格式"""
        tools = [ReadTool(), WriteTool()]

        anthropic_tools = []
        for tool in tools:
            raw_schema = tool.args_schema.model_json_schema()
            clean_schema = {
                "type": raw_schema.get("type", "object"),
                "properties": raw_schema.get("properties", {}),
                "required": raw_schema.get("required", []),
            }

            anthropic_tools.append(
                {
                    "name": tool.name,
                    "description": tool.description,
                    "input_schema": clean_schema,
                }
            )

        print(f"\nAnthropic Tools Array:")
        print(json.dumps(anthropic_tools, indent=2))

        # 验证
        assert len(anthropic_tools) == 2
        assert all("input_schema" in tool for tool in anthropic_tools)
        assert all("title" not in tool["input_schema"] for tool in anthropic_tools)


class TestAgentToolSetup:
    """测试 Agent 工具设置"""

    def test_agent_tool_definitions_with_fixed_schema(self):
        """测试 Agent 工具定义（使用修复后的 schema）"""
        # 直接测试工具 schema 转换，不实例化 Agent

        tools = [ReadTool(), WriteTool()]

        tool_defs = []
        for tool in tools:
            raw_schema = tool.args_schema.model_json_schema()
            clean_schema = {
                "type": raw_schema.get("type", "object"),
                "properties": raw_schema.get("properties", {}),
                "required": raw_schema.get("required", []),
            }

            from core.llm.base import ToolDefinition

            tool_defs.append(
                ToolDefinition(
                    name=tool.name,
                    description=tool.description,
                    parameters=clean_schema,
                )
            )

        print(f"\nAgent Tool Definitions Count: {len(tool_defs)}")
        for tool_def in tool_defs:
            print(f"\nTool: {tool_def.name}")
            print(json.dumps(tool_def.parameters, indent=2))

        # 验证所有工具定义
        assert len(tool_defs) == 2
        for tool_def in tool_defs:
            assert tool_def.parameters["type"] == "object"
            assert "title" not in tool_def.parameters
            assert "properties" in tool_def.parameters
            assert "required" in tool_def.parameters


class TestPydanticSchemaFeatures:
    """测试 Pydantic Schema 特性"""

    def test_schema_with_default_values(self):
        """测试带默认值的 Schema"""
        from pydantic import BaseModel, Field

        class ArgsWithDefaults(BaseModel):
            required_field: str = Field(..., description="Required field")
            optional_field: str = Field("default_value", description="Optional field")

        raw_schema = ArgsWithDefaults.model_json_schema()
        clean_schema = {
            "type": raw_schema.get("type", "object"),
            "properties": raw_schema.get("properties", {}),
            "required": raw_schema.get("required", []),
        }

        print(f"\nSchema with Default Values:")
        print(json.dumps(clean_schema, indent=2))

        # 验证
        assert len(clean_schema["required"]) == 1
        assert "required_field" in clean_schema["required"]
        assert "optional_field" not in clean_schema["required"]
        assert (
            clean_schema["properties"]["optional_field"]["default"] == "default_value"
        )

    def test_schema_with_nested_objects(self):
        """测试嵌套对象的 Schema"""
        from pydantic import BaseModel, Field

        class NestedConfig(BaseModel):
            key1: str = Field(..., description="Key 1")
            key2: int = Field(1, description="Key 2")

        class ArgsWithNested(BaseModel):
            name: str = Field(..., description="Name")
            config: NestedConfig = Field(..., description="Config")

        raw_schema = ArgsWithNested.model_json_schema()
        clean_schema = {
            "type": raw_schema.get("type", "object"),
            "properties": raw_schema.get("properties", {}),
            "required": raw_schema.get("required", []),
        }

        print(f"\nSchema with Nested Objects:")
        print(json.dumps(clean_schema, indent=2))

        # 验证
        assert "config" in clean_schema["properties"]
        # 注意：嵌套对象可能使用 $ref，所以我们不验证 type
        # 只要 properties 存在即可
        assert clean_schema["properties"]["config"] is not None

    def test_schema_with_enum(self):
        """测试带枚举的 Schema"""
        from pydantic import BaseModel, Field
        from enum import Enum

        class Color(str, Enum):
            RED = "red"
            GREEN = "green"
            BLUE = "blue"

        class ArgsWithEnum(BaseModel):
            color: Color = Field(Color.RED, description="Color")

        raw_schema = ArgsWithEnum.model_json_schema()
        clean_schema = {
            "type": raw_schema.get("type", "object"),
            "properties": raw_schema.get("properties", {}),
            "required": raw_schema.get("required", []),
        }

        print(f"\nSchema with Enum:")
        print(json.dumps(clean_schema, indent=2))

        # 验证
        assert "color" in clean_schema["properties"]
        # 注意：枚举可能使用 $ref，所以我们只验证属性存在
        assert clean_schema["properties"]["color"] is not None


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