"""
Testy propagacji zmian narzędzi z backendu przez Gateway do klienta.

Scenariusz: backend MCP dynamicznie dodaje nowy tool i wysyła
notifications/tools/list_changed do Gateway. Gateway:
1. Wykrywa notyfikację przez _ConnectionWorker (patch _received_notification)
2. Wywołuje Orchestrator.refresh_backend() — aktualizuje Registry
3. Wysyła notifications/tools/list_changed do podłączonych klientów (np. Claude)

Serwer testowy ma endpoint /control/add_tool który dodaje narzędzie
i broadcastuje notyfikację — całość działa w pętli eventowej serwera,
bez problemów z cross-loop.
"""

import asyncio
import contextlib
import threading
import time
import weakref
from collections.abc import Callable

import httpx
import mcp.types as types
import pytest
import uvicorn
from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client
from mcp.server.lowlevel import Server
from mcp.server.session import ServerSession
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager
from mcp.types import ServerNotification, ToolListChangedNotification
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Mount, Route

from ..connection_pool import ConnectionPool
from ..executor import Executor
from ..gateway import Gateway
from ..orchestrator import Orchestrator
from ..registry import Registry

DYNAMIC_BACKEND_HOST = "127.0.0.1"
DYNAMIC_BACKEND_PORT = 18132
DYNAMIC_BACKEND_URL = f"http://{DYNAMIC_BACKEND_HOST}:{DYNAMIC_BACKEND_PORT}/mcp/"
DYNAMIC_BACKEND_CONTROL = f"http://{DYNAMIC_BACKEND_HOST}:{DYNAMIC_BACKEND_PORT}/control"

GATEWAY_HOST = "127.0.0.1"
GATEWAY_PORT = 18133
GATEWAY_URL = f"http://{GATEWAY_HOST}:{GATEWAY_PORT}/mcp/"


class DynamicMCPServer:
    """Testowy serwer MCP który potrafi dynamicznie dodawać narzędzia.

    Endpoint /control/add_tool dodaje narzędzie i wysyła notifications/tools/list_changed
    do wszystkich podłączonych klientów. Całość działa w jednej pętli eventowej,
    bez problemów z cross-loop.
    """

    def __init__(self):
        self._tools: dict[str, types.Tool] = {}
        self._server = Server("dynamic-backend")
        self._sessions: weakref.WeakSet[ServerSession] = weakref.WeakSet()
        self._session_manager = StreamableHTTPSessionManager(
            app=self._server, json_response=False, stateless=False
        )
        self._setup_handlers()

    def _setup_handlers(self) -> None:
        @self._server.list_tools()
        async def list_tools() -> list[types.Tool]:
            try:
                self._sessions.add(self._server.request_context.session)
            except LookupError:
                pass
            return list(self._tools.values())

        @self._server.call_tool()
        async def call_tool(name: str, arguments: dict | None) -> list[types.TextContent]:
            if name not in self._tools:
                raise ValueError(f"Unknown tool: {name}")
            return [types.TextContent(type="text", text=f"{name}: ok")]

    def register_tool(self, name: str, description: str = "") -> None:
        self._tools[name] = types.Tool(
            name=name,
            description=description,
            inputSchema={"type": "object", "properties": {}},
        )

    async def _control_add_tool(self, request: Request) -> JSONResponse:
        data = await request.json()
        self.register_tool(data["name"], data.get("description", ""))
        for session in list(self._sessions):
            try:
                await session.send_tool_list_changed()
            except Exception:
                pass
        return JSONResponse({"ok": True})

    def create_app(self) -> Starlette:
        @contextlib.asynccontextmanager
        async def lifespan(app: Starlette):
            async with self._session_manager.run():
                yield

        return Starlette(
            lifespan=lifespan,
            routes=[
                Mount("/mcp", app=self._session_manager.handle_request),
                Route("/control/add_tool", self._control_add_tool, methods=["POST"]),
            ],
        )


@pytest.fixture(scope="module")
def dynamic_backend():
    server = DynamicMCPServer()
    server.register_tool("echo", "Zwraca tekst")

    config = uvicorn.Config(
        server.create_app(), host=DYNAMIC_BACKEND_HOST, port=DYNAMIC_BACKEND_PORT, log_level="warning"
    )
    uvi = uvicorn.Server(config)
    thread = threading.Thread(target=uvi.run, daemon=True)
    thread.start()

    deadline = time.time() + 5
    while time.time() < deadline:
        try:
            httpx.get(f"http://{DYNAMIC_BACKEND_HOST}:{DYNAMIC_BACKEND_PORT}/")
            break
        except Exception:
            time.sleep(0.1)

    yield {"url": DYNAMIC_BACKEND_URL, "control": DYNAMIC_BACKEND_CONTROL}

    uvi.should_exit = True


@pytest.fixture(scope="module")
def infrastructure(dynamic_backend):
    registry = Registry()
    pool = ConnectionPool()
    executor = Executor(pool)
    gateway = Gateway(registry, executor)
    orchestrator = Orchestrator(registry, pool, executor, gateway=gateway)

    app = gateway.create_app()
    config = uvicorn.Config(app, host=GATEWAY_HOST, port=GATEWAY_PORT, log_level="warning")
    server = uvicorn.Server(config)
    thread = threading.Thread(target=server.run, daemon=True)
    thread.start()

    deadline = time.time() + 5
    while time.time() < deadline:
        try:
            httpx.get(f"http://{GATEWAY_HOST}:{GATEWAY_PORT}/")
            break
        except Exception:
            time.sleep(0.1)

    yield {
        "orchestrator": orchestrator,
        "registry": registry,
        "backend_url": dynamic_backend["url"],
        "backend_control": dynamic_backend["control"],
    }

    server.should_exit = True


@pytest.fixture
async def system(infrastructure):
    yield infrastructure
    orch = infrastructure["orchestrator"]
    registry = infrastructure["registry"]
    for backend_id in list(registry._backends.keys()):
        await orch.remove_backend(backend_id)


def make_notification_collector(session: ClientSession) -> asyncio.Event:
    received = asyncio.Event()
    original = session._received_notification

    async def patched(notification: ServerNotification) -> None:
        if isinstance(notification.root, ToolListChangedNotification):
            received.set()
        await original(notification)

    session._received_notification = patched
    return received


@pytest.mark.anyio
async def test_initial_tools_visible_via_gateway(system):
    """Narzędzia zarejestrowane przed połączeniem widoczne przez Gateway."""
    await system["orchestrator"].add_backend("dynamic", system["backend_url"])

    async with streamable_http_client(GATEWAY_URL) as (read, write, _):
        async with ClientSession(read, write) as session:
            await session.initialize()
            result = await session.list_tools()

    names = [t.name for t in result.tools]
    assert "dynamic__echo" in names


@pytest.mark.anyio
async def test_gateway_client_notified_when_backend_adds_tool(system):
    """Klient Gateway dostaje notifications/tools/list_changed gdy backend dodaje narzędzie."""
    await system["orchestrator"].add_backend("dynamic", system["backend_url"])

    async with streamable_http_client(GATEWAY_URL) as (read, write, _):
        async with ClientSession(read, write) as session:
            await session.initialize()
            await session.list_tools()  # rejestruje sesję w Gateway

            notified = make_notification_collector(session)

            async with httpx.AsyncClient() as http:
                await http.post(
                    f"{system['backend_control']}/add_tool",
                    json={"name": "new_tool"},
                )

            await asyncio.wait_for(notified.wait(), timeout=5.0)
            assert notified.is_set()


@pytest.mark.anyio
async def test_after_backend_adds_tool_gateway_list_tools_updated(system):
    """Po notyfikacji z backendu, Gateway zwraca zaktualizowaną listę narzędzi."""
    await system["orchestrator"].add_backend("dynamic", system["backend_url"])

    async with streamable_http_client(GATEWAY_URL) as (read, write, _):
        async with ClientSession(read, write) as session:
            await session.initialize()
            await session.list_tools()

            notified = make_notification_collector(session)

            async with httpx.AsyncClient() as http:
                await http.post(
                    f"{system['backend_control']}/add_tool",
                    json={"name": "surprise_tool"},
                )

            await asyncio.wait_for(notified.wait(), timeout=5.0)

            updated = await session.list_tools()
            names = [t.name for t in updated.tools]
            assert "dynamic__surprise_tool" in names
            assert "dynamic__echo" in names


@pytest.mark.anyio
async def test_full_propagation_chain(system):
    """
    Weryfikuje pełny łańcuch:
    backend dodaje tool
      → _ConnectionWorker wykrywa ToolListChangedNotification
        → Orchestrator.refresh_backend()
          → Registry zaktualizowane
            → Gateway.notify_tool_list_changed()
              → klient Claude dostaje powiadomienie i widzi nowe narzędzie
    """
    await system["orchestrator"].add_backend("dynamic", system["backend_url"])

    async with streamable_http_client(GATEWAY_URL) as (read, write, _):
        async with ClientSession(read, write) as session:
            await session.initialize()

            before = await session.list_tools()
            names_before = {t.name for t in before.tools}

            notified = make_notification_collector(session)

            async with httpx.AsyncClient() as http:
                await http.post(
                    f"{system['backend_control']}/add_tool",
                    json={"name": "chain_tool", "description": "test propagacji"},
                )

            await asyncio.wait_for(notified.wait(), timeout=5.0)

            after = await session.list_tools()
            names_after = {t.name for t in after.tools}

            assert "dynamic__chain_tool" not in names_before
            assert "dynamic__chain_tool" in names_after
