"""Tests for preserving typed schemas on complex tool parameters.

AXM-704: Ensure _register_one preserves discriminated union types
(e.g. list[ReplaceOp | CreateOp]) in the JSON-Schema generated by
FastMCP, and that simple params remain unchanged.

Covers:
- Unit: complex param schema preserved, simple params unchanged
- Functional: all tools register, batch_edit schema via MCP
"""

from __future__ import annotations

import inspect
from typing import Any, Literal

import pydantic

from axm_mcp.discovery import _register_one, discover_tools, register_tools

# ────────────────────────────── Helpers ──────────────────────────────


class FakeMCP:
    """Minimal FastMCP stand-in that captures registered tools."""

    def __init__(self) -> None:
        self.tools: dict[str, Any] = {}

    def tool(self, *, name: str) -> Any:
        def decorator(fn: Any) -> Any:
            self.tools[name] = fn
            return fn

        return decorator


class ReplaceOp(pydantic.BaseModel):
    """Replace operation."""

    op: Literal["replace"] = "replace"
    file: str
    edits: list[dict[str, str]]


class CreateOp(pydantic.BaseModel):
    """Create operation."""

    op: Literal["create"] = "create"
    file: str
    content: str


class DeleteOp(pydantic.BaseModel):
    """Delete operation."""

    op: Literal["delete"] = "delete"
    file: str


# ────────────────────────────── Unit tests ───────────────────────────


class TestComplexParamSchemaPreserved:
    """Complex union params must retain discriminated type info."""

    def test_complex_param_schema_preserved(self) -> None:
        """Register tool with list[ReplaceOp | CreateOp] param.

        The wrapper signature must preserve the annotation so FastMCP
        can generate a JSON-Schema with discriminated union types.
        """

        class ComplexTool:
            def execute(
                self,
                *,
                path: str,
                operations: list[ReplaceOp | CreateOp | DeleteOp],
            ) -> Any:
                """Apply batch operations.

                Args:
                    path: Target project path.
                    operations: List of edit operations.

                Returns:
                    ToolResult with results.
                """

                class _R:
                    success = True
                    data: dict[str, Any] = {"applied": len(operations)}
                    error = None

                return _R()

        fake_mcp = FakeMCP()
        _register_one(fake_mcp, "complex_tool", ComplexTool())

        wrapper = fake_mcp.tools["complex_tool"]
        sig = inspect.signature(wrapper)

        # The 'operations' parameter must exist
        assert "operations" in sig.parameters
        ann = sig.parameters["operations"].annotation

        # The annotation must NOT be reduced to a plain type —
        # it should preserve the union/list structure
        assert ann is not inspect.Parameter.empty, (
            "Complex param annotation was stripped"
        )

        # Verify the annotation string contains the union types
        ann_str = str(ann)
        assert "ReplaceOp" in ann_str, f"ReplaceOp lost from annotation: {ann_str}"
        assert "CreateOp" in ann_str, f"CreateOp lost from annotation: {ann_str}"
        assert "DeleteOp" in ann_str, f"DeleteOp lost from annotation: {ann_str}"


class TestSimpleParamUnchanged:
    """Simple str/int params must keep their existing behavior."""

    def test_simple_param_unchanged(self) -> None:
        """Register tool with str, int params — schema identical."""

        class SimpleTool:
            def execute(
                self,
                *,
                path: str,
                limit: int = 10,
                verbose: bool = False,
            ) -> Any:
                """Do something simple.

                Args:
                    path: Target path.
                    limit: Max results.
                    verbose: Enable verbose output.

                Returns:
                    ToolResult with data.
                """

                class _R:
                    success = True
                    data = {"path": path, "limit": limit}
                    error = None

                return _R()

        fake_mcp = FakeMCP()
        _register_one(fake_mcp, "simple_tool", SimpleTool())

        wrapper = fake_mcp.tools["simple_tool"]
        sig = inspect.signature(wrapper)
        params = sig.parameters

        assert params["path"].annotation is str
        assert params["limit"].annotation is int
        assert params["limit"].default == 10
        assert params["verbose"].annotation is bool
        assert params["verbose"].default is False


# ──────────────── Functional / regression tests ──────────────────────


class TestAllExistingToolsRegister:
    """All discovered tools must register without error."""

    def test_all_existing_tools_register(self) -> None:
        """discover_tools() + register_tools() succeeds for every tool."""
        tools = discover_tools()
        assert len(tools) > 0, "No tools discovered"

        fake_mcp = FakeMCP()
        # Must not raise
        register_tools(fake_mcp, tools)

        # Every discovered tool should be registered
        for name in tools:
            assert name in fake_mcp.tools, (
                f"Tool '{name}' discovered but not registered"
            )


class TestBatchEditSchemaViaMCP:
    """batch_edit schema must expose typed operation fields via _register_one."""

    def test_batch_edit_schema_via_mcp(self) -> None:
        """Register a batch_edit-like tool and verify typed schema fields.

        Uses a self-contained fake tool with the same signature shape as
        the real batch_edit (path + operations with union types), so the
        test doesn't depend on axm-edit being installed.
        """

        class FakeBatchEditTool:
            def execute(
                self,
                *,
                path: str,
                operations: list[ReplaceOp | CreateOp | DeleteOp],
            ) -> Any:
                """Apply batch file operations.

                Args:
                    path: Project root directory.
                    operations: List of edit operations.

                Returns:
                    ToolResult with results.
                """

                class _R:
                    success = True
                    data: dict[str, Any] = {"applied": len(operations)}
                    error = None

                return _R()

        fake_mcp = FakeMCP()
        _register_one(fake_mcp, "batch_edit", FakeBatchEditTool())

        wrapper = fake_mcp.tools["batch_edit"]
        sig = inspect.signature(wrapper)

        # batch_edit must expose 'path' and 'operations' params
        param_names = list(sig.parameters.keys())
        assert "path" in param_names, (
            f"'path' missing from batch_edit params: {param_names}"
        )
        assert "operations" in param_names, (
            f"'operations' missing from batch_edit params: {param_names}"
        )

        # The 'operations' annotation must preserve the typed schema
        ops_ann = sig.parameters["operations"].annotation
        assert ops_ann is not inspect.Parameter.empty, (
            "batch_edit 'operations' annotation was stripped"
        )
