"""Tests for MCP resilience — verifying the sidecar survives hostile SDK behavior.

Every test here simulates a failure mode from the MCP SDK (BaseExceptionGroup,
hanging cleanup, dead connections, injected CancelledError) and verifies that
the MCPManager and tool proxies contain the error without crashing.
"""

from __future__ import annotations

import asyncio
from contextlib import AsyncExitStack
from types import SimpleNamespace
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from icarus.mcp.client import (
    MCPManager,
    _extract_leaf_error,
    _safe_aclose,
    _wrap_tool,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _make_mock_tool(name: str = "test_tool", response_format: str | None = None):
    """Create a mock LangChain tool with an async _arun."""
    tool = MagicMock()
    tool.name = name
    tool._arun = AsyncMock(return_value="ok")
    if response_format:
        tool.response_format = response_format
    return tool


def _make_base_exception_group(msg: str = "connection failed") -> BaseExceptionGroup:
    """Create a BaseExceptionGroup like anyio produces."""
    inner = ConnectionError(msg)
    return BaseExceptionGroup("unhandled errors in a TaskGroup", [inner])


# ---------------------------------------------------------------------------
# _safe_aclose
# ---------------------------------------------------------------------------

class TestSafeAclose:
    @pytest.mark.asyncio
    async def test_catches_exception(self):
        stack = AsyncExitStack()
        stack.aclose = AsyncMock(side_effect=RuntimeError("cleanup failed"))
        # Should not raise
        await _safe_aclose(stack)

    @pytest.mark.asyncio
    async def test_catches_base_exception_group(self):
        stack = AsyncExitStack()
        stack.aclose = AsyncMock(side_effect=_make_base_exception_group())
        # Should not raise
        await _safe_aclose(stack)

    @pytest.mark.asyncio
    async def test_empty_stack_succeeds(self):
        stack = AsyncExitStack()
        await _safe_aclose(stack)


# ---------------------------------------------------------------------------
# _extract_leaf_error
# ---------------------------------------------------------------------------

class TestExtractLeafError:
    def test_plain_exception(self):
        exc = ValueError("test")
        assert _extract_leaf_error(exc) is exc

    def test_single_nested(self):
        inner = ConnectionError("refused")
        group = BaseExceptionGroup("outer", [inner])
        assert _extract_leaf_error(group) is inner

    def test_double_nested(self):
        inner = ConnectionError("refused")
        mid = BaseExceptionGroup("mid", [inner])
        outer = BaseExceptionGroup("outer", [mid])
        assert _extract_leaf_error(outer) is inner


# ---------------------------------------------------------------------------
# _wrap_tool — proxy behavior
# ---------------------------------------------------------------------------

class TestWrapTool:
    @pytest.mark.asyncio
    async def test_success_passes_through(self):
        tool = _make_mock_tool()
        state: dict[str, bool] = {}
        wrapped = _wrap_tool(tool, "server1", {}, state, None, None, None)
        result = await wrapped._arun()
        assert result == "ok"

    @pytest.mark.asyncio
    async def test_exception_returns_error_string(self):
        tool = _make_mock_tool()
        tool._arun = AsyncMock(side_effect=ConnectionError("refused"))
        state: dict[str, bool] = {}
        wrapped = _wrap_tool(tool, "server1", {}, state, None, None, None)
        result = await wrapped._arun()
        assert "[Tool unavailable]" in result
        assert "server1" in result

    @pytest.mark.asyncio
    async def test_base_exception_group_returns_error_string(self):
        tool = _make_mock_tool()
        tool._arun = AsyncMock(side_effect=_make_base_exception_group("dead"))
        state: dict[str, bool] = {}
        wrapped = _wrap_tool(tool, "server1", {}, state, None, None, None)
        result = await wrapped._arun()
        assert "[Tool unavailable]" in result
        assert "dead" in result

    @pytest.mark.asyncio
    async def test_content_and_artifact_returns_tuple(self):
        tool = _make_mock_tool(response_format="content_and_artifact")
        tool._arun = AsyncMock(side_effect=ConnectionError("down"))
        state: dict[str, bool] = {}
        wrapped = _wrap_tool(tool, "server1", {}, state, None, None, None)
        result = await wrapped._arun()
        assert isinstance(result, tuple)
        assert len(result) == 2
        assert "[Tool unavailable]" in result[0]
        assert result[1] is None

    @pytest.mark.asyncio
    async def test_on_failure_called_once_per_server(self):
        """Multiple tool failures on same server should fire on_failure once."""
        failures: list[str] = []
        state: dict[str, bool] = {}

        tool1 = _make_mock_tool("tool1")
        tool1._arun = AsyncMock(side_effect=ConnectionError("down"))
        tool2 = _make_mock_tool("tool2")
        tool2._arun = AsyncMock(side_effect=ConnectionError("down"))

        on_fail = lambda name, msg: failures.append(name)
        wrapped1 = _wrap_tool(tool1, "server1", {}, state, on_fail, None, None)
        wrapped2 = _wrap_tool(tool2, "server1", {}, state, on_fail, None, None)

        await wrapped1._arun()
        await wrapped2._arun()

        assert len(failures) == 1  # not 2

    @pytest.mark.asyncio
    async def test_auto_recovery(self):
        """Success after failure should call on_recovery."""
        recoveries: list[str] = []
        state: dict[str, bool] = {}

        tool = _make_mock_tool()
        original_arun = tool._arun

        # First call fails
        call_count = 0
        async def flaky(*args, **kwargs):
            nonlocal call_count
            call_count += 1
            if call_count == 1:
                raise ConnectionError("down")
            return "ok"

        tool._arun = flaky
        on_fail = lambda name, msg: None
        on_recover = lambda name: recoveries.append(name)
        wrapped = _wrap_tool(tool, "server1", {}, state, on_fail, on_recover, None)

        # Fail first
        await wrapped._arun()
        assert state.get("degraded") is True

        # Succeed second
        await wrapped._arun()
        assert state.get("degraded") is False
        assert len(recoveries) == 1

    @pytest.mark.asyncio
    async def test_sets_server_name(self):
        tool = _make_mock_tool()
        state: dict[str, bool] = {}
        wrapped = _wrap_tool(tool, "myserver", {}, state, None, None, None)
        assert wrapped._server_name == "myserver"

    @pytest.mark.asyncio
    async def test_degraded_short_circuits_when_reconnecting(self):
        """When reconnecting flag is set, tool returns immediately."""
        tool = _make_mock_tool()
        state: dict[str, bool] = {"degraded": True, "reconnecting": True}
        wrapped = _wrap_tool(tool, "server1", {}, state, None, None, None)
        result = await wrapped._arun()
        assert "Reconnecting" in str(result)

    @pytest.mark.asyncio
    async def test_concurrent_tools_fire_reconnect_once(self):
        """Multiple degraded tools should fire on_reconnect only once."""
        reconnects: list[str] = []
        state: dict[str, bool] = {"degraded": True}
        # Use a URL that's reachable (localhost with a listening port)
        import socket
        with socket.socket() as s:
            s.bind(("127.0.0.1", 0))
            port = s.getsockname()[1]
            s.listen(1)

            cfg = {"url": f"http://127.0.0.1:{port}/mcp"}

            t1 = _make_mock_tool("tool1")
            t2 = _make_mock_tool("tool2")
            on_reconnect = lambda name: reconnects.append(name)
            w1 = _wrap_tool(t1, "s1", cfg, state, None, None, on_reconnect)
            w2 = _wrap_tool(t2, "s1", cfg, state, None, None, on_reconnect)

            await asyncio.gather(w1._arun(), w2._arun())

        assert len(reconnects) == 1


# ---------------------------------------------------------------------------
# MCPManager.start — per-server isolation
# ---------------------------------------------------------------------------

class TestMCPManagerStart:
    @pytest.mark.asyncio
    async def test_unreachable_server_returns_failed(self):
        """Server that fails _check_reachable should be in the failed list."""
        with patch("icarus.mcp.client._check_reachable", side_effect=ConnectionError("refused")):
            mgr = MCPManager({"bad-server": {"url": "http://127.0.0.1:9999/mcp"}})
            tools, failed = await mgr.start()
            assert tools == []
            assert failed == ["bad-server"]
            assert mgr.server_status["bad-server"] == "error"
            assert "refused" in mgr.server_errors["bad-server"]

    @pytest.mark.asyncio
    async def test_base_exception_group_caught(self):
        """BaseExceptionGroup from anyio should be caught per-server."""
        with patch("icarus.mcp.client._open_session", side_effect=_make_base_exception_group()):
            mgr = MCPManager({"bad": {"url": "http://127.0.0.1:9999/mcp"}})
            tools, failed = await mgr.start()
            assert failed == ["bad"]
            assert mgr.server_status["bad"] == "error"

    @pytest.mark.asyncio
    async def test_healthy_server_not_affected_by_failed(self):
        """A healthy server should connect even if another fails."""
        mock_session = AsyncMock()
        mock_session.initialize = AsyncMock()
        mock_tools = [_make_mock_tool("tool1")]

        call_count = 0
        async def open_session_side_effect(cfg, stack):
            nonlocal call_count
            call_count += 1
            if cfg.get("url") == "http://bad:1/mcp":
                raise ConnectionError("refused")
            return mock_session

        with patch("icarus.mcp.client._open_session", side_effect=open_session_side_effect), \
             patch("langchain_mcp_adapters.tools.load_mcp_tools", return_value=mock_tools):
            mgr = MCPManager({
                "good": {"url": "http://good:1/mcp"},
                "bad": {"url": "http://bad:1/mcp"},
            })
            tools, failed = await mgr.start()
            assert failed == ["bad"]
            assert len(tools) == 1
            assert mgr.server_status["good"] == "connected"
            assert mgr.server_status["bad"] == "error"

    @pytest.mark.asyncio
    async def test_disabled_servers_not_connected(self):
        mgr = MCPManager({}, disabled={"off": {"url": "http://x/mcp"}})
        tools, failed = await mgr.start()
        assert tools == []
        assert failed == []
        assert mgr.server_status["off"] == "disabled"


# ---------------------------------------------------------------------------
# MCPManager.stop — per-server teardown
# ---------------------------------------------------------------------------

class TestMCPManagerStop:
    @pytest.mark.asyncio
    async def test_stop_clears_state(self):
        mgr = MCPManager({})
        mgr._tools = [1, 2, 3]
        mgr._stacks["a"] = AsyncExitStack()
        await mgr.stop()
        assert mgr._tools == []
        assert mgr._stacks == {}

    @pytest.mark.asyncio
    async def test_stop_survives_cleanup_error(self):
        """stop() should not raise even if a stack cleanup fails."""
        bad_stack = AsyncExitStack()
        bad_stack.aclose = AsyncMock(side_effect=_make_base_exception_group())
        mgr = MCPManager({})
        mgr._stacks["bad"] = bad_stack
        # Should not raise
        await mgr.stop()


# ---------------------------------------------------------------------------
# MCPManager.mark_degraded / mark_recovered
# ---------------------------------------------------------------------------

class TestMCPManagerDegradedRecovery:
    def test_mark_degraded(self):
        mgr = MCPManager({})
        mgr._server_status["s1"] = "connected"
        mgr.mark_degraded("s1", "timeout")
        assert mgr.server_status["s1"] == "degraded"
        assert mgr.server_errors["s1"] == "timeout"

    def test_mark_degraded_only_from_connected(self):
        mgr = MCPManager({})
        mgr._server_status["s1"] = "error"
        mgr.mark_degraded("s1", "timeout")
        assert mgr.server_status["s1"] == "error"  # unchanged

    def test_mark_recovered(self):
        mgr = MCPManager({})
        mgr._server_status["s1"] = "degraded"
        mgr._server_errors["s1"] = "old error"
        mgr.mark_recovered("s1")
        assert mgr.server_status["s1"] == "connected"
        assert "s1" not in mgr.server_errors

    def test_mark_recovered_only_from_degraded(self):
        mgr = MCPManager({})
        mgr._server_status["s1"] = "error"
        mgr.mark_recovered("s1")
        assert mgr.server_status["s1"] == "error"  # unchanged


# ---------------------------------------------------------------------------
# MCPManager.disconnect_server
# ---------------------------------------------------------------------------

class TestDisconnectServer:
    @pytest.mark.asyncio
    async def test_removes_tools_by_server_name(self):
        mgr = MCPManager({})
        tool1 = _make_mock_tool("t1")
        tool1._server_name = "s1"
        tool2 = _make_mock_tool("t2")
        tool2._server_name = "s2"
        mgr._tools = [tool1, tool2]
        mgr._stacks["s1"] = AsyncExitStack()

        await mgr.disconnect_server("s1")
        assert len(mgr._tools) == 1
        assert mgr._tools[0]._server_name == "s2"
        assert mgr.server_status["s1"] == "disabled"


# ---------------------------------------------------------------------------
# Sidecar ↔ Engine callback wiring
# ---------------------------------------------------------------------------

def _mock_checkpointer():
    cp = AsyncMock()
    cm = AsyncMock()
    cm.__aenter__ = AsyncMock(return_value=cp)
    cm.__aexit__ = AsyncMock(return_value=False)
    return cm, cp


def _boot_patches(engine, *, failed=None):
    """Context manager stack for a mock boot() with optional failed servers."""
    from contextlib import ExitStack
    mock_mcp = AsyncMock()
    mock_mcp.start = AsyncMock(return_value=(["tool1"], failed or []))
    mock_mcp.mark_degraded = MagicMock()
    mock_mcp.mark_recovered = MagicMock()
    mock_mcp.server_errors = {}
    cm, _ = _mock_checkpointer()
    stack = ExitStack()
    stack.enter_context(patch(
        "icarus.runtime.context.UserContext.with_workspace",
        return_value=MagicMock(config_dir=MagicMock()),
    ))
    stack.enter_context(patch("icarus.session.threads.get_checkpointer", return_value=cm))
    stack.enter_context(patch("icarus.mcp.MCPManager.from_config_file", return_value=mock_mcp))
    stack.enter_context(patch("icarus.config.app_config.settings"))
    stack.enter_context(patch.object(engine, "_boot_agent_background", new_callable=AsyncMock))
    return stack, mock_mcp


class TestSidecarCallbackWiring:
    """Verify sidecar MCP overrides call engine originals without recursion."""

    @pytest.mark.asyncio
    async def test_on_failure_no_recursion(self):
        from desktop.sidecar.main import Sidecar

        responses = []
        with patch("desktop.sidecar.main.write_response", responses.append):
            sidecar = Sidecar()
            stack, mock_mcp = _boot_patches(sidecar._engine)
            with stack:
                sidecar._install_mcp_callbacks()
                await sidecar._engine.boot()

            # This is what MCPManager's tool wrapper calls on failure
            sidecar._engine._on_mcp_server_failure("srv", "refused")

        warnings = [r for r in responses if r.get("type") == "connector_warning"]
        assert len(warnings) == 1
        assert warnings[0]["name"] == "srv"
        mock_mcp.mark_degraded.assert_called_once_with("srv", "refused")

    @pytest.mark.asyncio
    async def test_on_recovery_no_recursion(self):
        from desktop.sidecar.main import Sidecar

        responses = []
        with patch("desktop.sidecar.main.write_response", responses.append):
            sidecar = Sidecar()
            stack, mock_mcp = _boot_patches(sidecar._engine)
            with stack:
                sidecar._install_mcp_callbacks()
                await sidecar._engine.boot()

            sidecar._engine._on_mcp_server_recovery("srv")

        recovered = [r for r in responses if r.get("type") == "connector_recovered"]
        assert len(recovered) == 1
        mock_mcp.mark_recovered.assert_called_once_with("srv")

    @pytest.mark.asyncio
    async def test_boot_failures_awaited_before_ready(self):
        """Auto-disable must complete before the 'ready' event."""
        from desktop.sidecar.main import Sidecar

        order = []

        def capture(msg):
            if msg.get("type") in ("ready", "connector_warning"):
                order.append(msg["type"])

        with patch("desktop.sidecar.main.write_response", capture):
            sidecar = Sidecar()
            stack, _ = _boot_patches(sidecar._engine, failed=["bad"])
            with (
                stack,
                patch("icarus.mcp.config.read_mcp_config", return_value={
                    "mcpServers": {"bad": {"command": "fake"}},
                }),
                patch("icarus.mcp.config.write_mcp_config"),
                patch("icarus.mcp.config.toggle_server"),
            ):
                sidecar._install_mcp_callbacks()
                await sidecar.boot()

        assert "connector_warning" in order, "auto-disable didn't fire"
        assert "ready" in order, "ready was never sent"
        assert order.index("connector_warning") < order.index("ready"), (
            f"ready sent before auto-disable completed: {order}"
        )
