"""Tests for ToolRegistry."""

import pytest

from koan.tools.decorator import tool
from koan.tools.registry import ToolRegistry
from koan.tools.types import ToolExecutionError, ToolValidationError


@tool
def add_numbers(a: int, b: int) -> int:
    """Add two numbers."""
    return a + b


@tool
async def async_concat(x: str, y: str) -> str:
    """Concatenate two strings."""
    return x + y


@tool
def failing_tool(msg: str) -> str:
    """Always fails."""
    raise RuntimeError(msg)


class TestRegistryBasics:
    def test_add_and_len(self) -> None:
        reg = ToolRegistry()
        reg.add(add_numbers)
        assert len(reg) == 1
        assert "add_numbers" in reg

    def test_register_decorator(self) -> None:
        reg = ToolRegistry()

        @reg.register
        @tool
        def local_tool(x: int) -> int:
            return x * 2

        assert "local_tool" in reg
        assert len(reg) == 1

    def test_add_non_tool_raises(self) -> None:
        reg = ToolRegistry()

        def plain_func() -> None:
            pass

        with pytest.raises(ValueError, match="not decorated with @tool"):
            reg.add(plain_func)

    def test_get(self) -> None:
        reg = ToolRegistry()
        reg.add(add_numbers)
        assert reg.get("add_numbers") is add_numbers
        assert reg.get("nonexistent") is None

    def test_definitions(self) -> None:
        reg = ToolRegistry()
        reg.add(add_numbers)
        reg.add(async_concat)

        defns = reg.definitions()
        assert len(defns) == 2
        names = {d.name for d in defns}
        assert names == {"add_numbers", "async_concat"}

    def test_to_openai_tools(self) -> None:
        reg = ToolRegistry()
        reg.add(add_numbers)

        tools = reg.to_openai_tools()
        assert len(tools) == 1
        assert tools[0]["type"] == "function"
        assert tools[0]["function"]["name"] == "add_numbers"
        assert "parameters" in tools[0]["function"]


class TestRegistryExecution:
    async def test_execute_sync_tool(self) -> None:
        reg = ToolRegistry()
        reg.add(add_numbers)

        result = await reg.execute("add_numbers", {"a": 3, "b": 5})
        assert result.output == 8
        assert result.tool_name == "add_numbers"
        assert not result.is_error

    async def test_execute_async_tool(self) -> None:
        reg = ToolRegistry()
        reg.add(async_concat)

        result = await reg.execute("async_concat", {"x": "hello", "y": " world"})
        assert result.output == "hello world"

    async def test_execute_with_json_string(self) -> None:
        reg = ToolRegistry()
        reg.add(add_numbers)

        result = await reg.execute("add_numbers", '{"a": 10, "b": 20}')
        assert result.output == 30

    async def test_execute_unknown_tool(self) -> None:
        reg = ToolRegistry()
        with pytest.raises(KeyError, match="Tool not found"):
            await reg.execute("nonexistent", {})

    async def test_execute_invalid_args(self) -> None:
        reg = ToolRegistry()
        reg.add(add_numbers)

        with pytest.raises(ToolValidationError, match="Validation failed"):
            await reg.execute("add_numbers", {"a": "not_an_int", "b": "also_not"})

    async def test_execute_missing_required_arg(self) -> None:
        reg = ToolRegistry()
        reg.add(add_numbers)

        with pytest.raises(ToolValidationError, match="Validation failed"):
            await reg.execute("add_numbers", {"a": 1})

    async def test_execute_tool_that_raises(self) -> None:
        reg = ToolRegistry()
        reg.add(failing_tool)

        with pytest.raises(ToolExecutionError, match="boom"):
            await reg.execute("failing_tool", {"msg": "boom"})

    async def test_execute_with_context(self) -> None:
        class RunContext:
            def __init__(self, user_id: str) -> None:
                self.user_id = user_id

        @tool
        def ctx_tool(query: str, ctx: RunContext) -> str:
            """Uses context."""
            return f"{ctx.user_id}: {query}"

        reg = ToolRegistry()
        reg.add(ctx_tool)

        context = RunContext(user_id="u123")
        result = await reg.execute("ctx_tool", {"query": "hello"}, context=context)
        assert result.output == "u123: hello"
