"""Tests for tool preparation and schema handling."""

from typing import Any
from unittest.mock import Mock

from swarm.tool import prepare_tools_for_completion


class TestPrepareToolsForCompletion:
    """Test the prepare_tools_for_completion function."""

    def test_empty_tools_list(self):
        """Test handling of empty tools list."""
        result = prepare_tools_for_completion([])
        assert result == []

    def test_basic_function_without_special_params(self):
        """Test function with no special parameters."""

        def simple_function(name: str, age: int) -> str:
            """A simple test function."""
            return f"Hello {name}, you are {age} years old"

        tools = prepare_tools_for_completion([simple_function])

        assert len(tools) == 1
        tool = tools[0]

        # Check basic structure
        assert tool["type"] == "function"
        assert "function" in tool

        func_def = tool["function"]
        assert func_def["name"] == "simple_function"
        assert func_def["description"] == "A simple test function."

        # Check parameters
        params = func_def["parameters"]
        assert "name" in params["properties"]
        assert "age" in params["properties"]
        assert "name" in params["required"]
        assert "age" in params["required"]

    def test_function_with_ctx(self):
        """Test that ctx is stripped from schema."""

        def function_with_ctx(name: str, ctx: Any) -> str:
            """Function that uses ctx parameter."""
            return f"Hello {name}"

        tools = prepare_tools_for_completion([function_with_ctx])

        assert len(tools) == 1
        tool = tools[0]

        params = tool["function"]["parameters"]
        properties = params["properties"]
        required = params["required"]

        # ctx should be removed
        assert "ctx" not in properties
        assert "ctx" not in required

        # Other params should remain
        assert "name" in properties
        assert "name" in required

    def test_function_with_only_ctx(self):
        """Test function with only ctx parameter."""

        def func_only_ctx(ctx: Any) -> str:
            """Function with only ctx."""
            return "ctx only"

        tools = prepare_tools_for_completion([func_only_ctx])

        assert len(tools) == 1
        tool = tools[0]

        params = tool["function"]["parameters"]
        properties = params["properties"]
        required = params.get("required", [])

        # ctx should be removed, leaving empty properties
        assert "ctx" not in properties
        assert len(properties) == 0
        assert len(required) == 0

    def test_function_with_ctx_and_multiple_params(self):
        """Test function with ctx and multiple regular parameters."""

        def function_with_ctx_and_params(
            name: str,
            age: int,
            ctx: Any,
            city: str = "NYC",
        ) -> str:
            """Function with ctx and params."""
            return f"Hello {name}, age {age}, from {city}"

        tools = prepare_tools_for_completion([function_with_ctx_and_params])

        assert len(tools) == 1
        tool = tools[0]

        params = tool["function"]["parameters"]
        properties = params["properties"]
        required = params["required"]

        # ctx should be removed
        assert "ctx" not in properties
        assert "ctx" not in required

        # Regular params should remain
        assert "name" in properties
        assert "age" in properties
        assert "city" in properties
        assert "name" in required
        assert "age" in required

        # Should have 3 properties (name, age, city)
        assert len(properties) == 3

    def test_multiple_functions(self):
        """Test multiple functions in one call."""

        def func1(name: str, ctx: Any) -> str:
            return f"Func1: {name}"

        def func2(age: int) -> str:
            return f"Func2: {age}"

        def func3(city: str, ctx: Any) -> str:
            return f"Func3: {city}"

        tools = prepare_tools_for_completion([func1, func2, func3])

        assert len(tools) == 3

        # Check func1 - ctx should be removed
        func1_params = tools[0]["function"]["parameters"]
        assert "name" in func1_params["properties"]
        assert "ctx" not in func1_params["properties"]

        # Check func2 - no ctx
        func2_params = tools[1]["function"]["parameters"]
        assert "age" in func2_params["properties"]

        # Check func3 - ctx should be removed
        func3_params = tools[2]["function"]["parameters"]
        assert "city" in func3_params["properties"]
        assert "ctx" not in func3_params["properties"]

    def test_malformed_schema_handling(self):
        """Test handling of malformed schemas."""
        # Create a mock function that returns a malformed schema
        mock_func = Mock()
        mock_func.__name__ = "mock_function"
        mock_func.__doc__ = "Mock function"

        # Create a signature that will cause issues
        def malformed_schema_func():
            pass

        # This should not crash even if function introspection fails
        tools = prepare_tools_for_completion([malformed_schema_func])
        assert len(tools) == 1
        assert tools[0]["function"]["name"] == "malformed_schema_func"


class TestToolSchemaConsistency:
    """Test that tool schemas are consistent with OpenAI format."""

    def test_openai_format_structure(self):
        """Test that generated schemas match OpenAI format."""

        def sample_function(name: str, age: int) -> str:
            """Sample function for testing."""
            return f"{name} is {age}"

        tools = prepare_tools_for_completion([sample_function])
        tool = tools[0]

        # Check top-level structure
        assert "type" in tool
        assert tool["type"] == "function"
        assert "function" in tool

        # Check function definition structure
        func_def = tool["function"]
        assert "name" in func_def
        assert "description" in func_def
        assert "parameters" in func_def

        # Check parameters structure
        params = func_def["parameters"]
        assert "type" in params
        assert params["type"] == "object"
        assert "properties" in params
        assert "required" in params

        # Properties should be a dict
        assert isinstance(params["properties"], dict)
        # Required should be a list
        assert isinstance(params["required"], list)

    def test_ctx_stripped_from_mixed_params(self):
        """Test ctx is stripped when mixed with regular parameters."""

        def mixed_params(
            regular1: str,
            regular2: int,
            ctx: Any,
            regular3: str = "default",
        ) -> str:
            return "mixed"

        tools = prepare_tools_for_completion([mixed_params])
        tool = tools[0]

        params = tool["function"]["parameters"]
        properties = params["properties"]
        required = params["required"]

        # Should have only regular parameters
        expected_properties = {"regular1", "regular2", "regular3"}
        actual_properties = set(properties.keys())
        assert actual_properties == expected_properties

        # Required should only include required regular params
        expected_required = {"regular1", "regular2"}
        actual_required = set(required)
        assert actual_required == expected_required
