"""Tests for the MCP server wrapper."""

from unittest.mock import AsyncMock, MagicMock

import pytest

from aal_sdk.mcp import wrap_mcp_server
from aal_sdk.session import AALSession


def _make_session() -> AALSession:
    """Create a test session with a mock push function."""
    return AALSession("sess-1", "agent-1", MagicMock())


def _make_server(tools: dict) -> MagicMock:
    """Create a mock MCP server with registered tools."""
    server = MagicMock()
    server._registeredTools = tools
    return server


async def test_wrap_mcp_server_wraps_all_tools():
    """Test that wrap_mcp_server replaces all registered tool handlers."""
    original_handler_a = AsyncMock(return_value={"result": "a"})
    original_handler_b = AsyncMock(return_value={"result": "b"})
    server = _make_server({
        "tool_a": original_handler_a,
        "tool_b": original_handler_b,
    })
    session = _make_session()

    wrap_mcp_server(server, session)

    # Handlers should be replaced (not the same object)
    assert server._registeredTools["tool_a"] is not original_handler_a
    assert server._registeredTools["tool_b"] is not original_handler_b


async def test_wrapped_tool_logs_success():
    """Test that a wrapped tool logs SUCCESS on successful execution."""
    original = AsyncMock(return_value={"weather": "sunny"})
    server = _make_server({"get_weather": original})
    session = _make_session()

    wrap_mcp_server(server, session)

    result = await server._registeredTools["get_weather"]({"city": "London"})

    assert result == {"weather": "sunny"}
    original.assert_called_once_with({"city": "London"})
    session._push_fn.assert_called_once()
    event = session._push_fn.call_args[0][0]
    assert event["action_type"] == "TOOL_CALL"
    assert event["tool_name"] == "get_weather"
    assert event["parameters"] == {"city": "London"}
    assert event["response"] == {"weather": "sunny"}
    assert event["status"] == "SUCCESS"
    assert "latency_ms" in event


async def test_wrapped_tool_logs_error_on_exception():
    """Test that a wrapped tool logs ERROR when an exception occurs."""
    original = AsyncMock(side_effect=ValueError("bad input"))
    server = _make_server({"broken_tool": original})
    session = _make_session()

    wrap_mcp_server(server, session)

    with pytest.raises(ValueError, match="bad input"):
        await server._registeredTools["broken_tool"]({"x": 1})

    event = session._push_fn.call_args[0][0]
    assert event["status"] == "ERROR"
    assert event["tool_name"] == "broken_tool"
    assert event["response"] == {"error": "bad input"}


async def test_original_exception_is_reraised():
    """Test that the original exception is re-raised after logging."""
    original_error = RuntimeError("something broke")
    original = AsyncMock(side_effect=original_error)
    server = _make_server({"failing_tool": original})
    session = _make_session()

    wrap_mcp_server(server, session)

    with pytest.raises(RuntimeError) as exc_info:
        await server._registeredTools["failing_tool"]()

    assert exc_info.value is original_error
