"""
Pula połączeń do backendowych serwerów MCP (transport HTTP/SSE).

Odpowiedzialność:
- Tworzy i cache'uje ClientSession (MCP SDK) dla każdego backendu
- Sprawdza żywotność połączenia przed zwróceniem (ping)
- Odtwarza połączenie jeśli zerwane (reconnect)
- Zamyka połączenia przy usuwaniu backendu
- Udostępnia metody: get(backend_id, url) -> ClientSession, remove(backend_id), close_all()
- Nie zna narzędzi ani logiki wywołań — tylko zarządza połączeniami
- Używany wyłącznie przez Executor
"""

import asyncio
import logging
from collections.abc import Callable

from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client
from mcp.types import (
    ResourceListChangedNotification,
    ResourceUpdatedNotification,
    ServerNotification,
    ToolListChangedNotification,
)

logger = logging.getLogger(__name__)


class _ConnectionWorker:
    """Zarządza cyklem życia pojedynczego połączenia MCP w dedykowanym tasku.

    Otwiera i zamyka streamable_http_client + ClientSession zawsze w tym samym
    tasku, co pozwala uniknąć problemów z propagacją cancel scope anyio.
    """

    def __init__(
        self,
        url: str,
        on_tool_list_changed: Callable[[], None] | None = None,
        on_disconnect: Callable[[], None] | None = None,
        on_resource_list_changed: Callable[[], None] | None = None,
        on_resource_updated: Callable[[str], None] | None = None,
    ):
        self._url = url
        self._on_tool_list_changed = on_tool_list_changed
        self._on_disconnect = on_disconnect
        self._on_resource_list_changed = on_resource_list_changed
        self._on_resource_updated = on_resource_updated
        self._session: ClientSession | None = None
        self._ready: asyncio.Event = asyncio.Event()
        self._close_event: asyncio.Event = asyncio.Event()
        self._task: asyncio.Task | None = None
        self._error: BaseException | None = None
        self._connected = False  # True po pomyślnym initialize()

    async def _run(self) -> None:
        try:
            async with streamable_http_client(self._url) as (read, write, _):
                async with ClientSession(read, write) as session:
                    if (
                        self._on_tool_list_changed
                        or self._on_resource_list_changed
                        or self._on_resource_updated
                    ):
                        original = session._received_notification

                        async def patched(notification: ServerNotification) -> None:
                            root = notification.root
                            if (
                                isinstance(root, ToolListChangedNotification)
                                and self._on_tool_list_changed
                            ):
                                self._on_tool_list_changed()
                            elif (
                                isinstance(root, ResourceListChangedNotification)
                                and self._on_resource_list_changed
                            ):
                                self._on_resource_list_changed()
                            elif (
                                isinstance(root, ResourceUpdatedNotification)
                                and self._on_resource_updated
                            ):
                                self._on_resource_updated(str(root.params.uri))
                            await original(notification)

                        session._received_notification = patched
                    await session.initialize()
                    self._session = session
                    self._connected = True
                    self._ready.set()
                    await self._close_event.wait()
        except Exception as e:
            self._error = e
            self._ready.set()
        finally:
            if self._connected and not self._close_event.is_set() and self._on_disconnect:
                self._on_disconnect()

    async def start(self) -> ClientSession:
        self._task = asyncio.get_event_loop().create_task(self._run())
        await self._ready.wait()
        if self._error:
            raise self._error
        return self._session

    async def close(self) -> None:
        self._close_event.set()
        if self._task:
            try:
                await asyncio.wait_for(self._task, timeout=5.0)
            except (asyncio.TimeoutError, BaseException):
                pass

    @property
    def session(self) -> ClientSession | None:
        return self._session


class ConnectionPool:
    def __init__(self):
        self._workers: dict[str, _ConnectionWorker] = {}
        self._on_tool_list_changed: Callable[[str], None] | None = None
        self._on_disconnect: Callable[[str], None] | None = None
        self._on_resource_list_changed: Callable[[str], None] | None = None
        self._on_resource_updated: Callable[[str, str], None] | None = None

    def set_on_tool_list_changed(self, callback: Callable[[str], None]) -> None:
        self._on_tool_list_changed = callback

    def set_on_disconnect(self, callback: Callable[[str], None]) -> None:
        self._on_disconnect = callback

    def set_on_resource_list_changed(self, callback: Callable[[str], None]) -> None:
        self._on_resource_list_changed = callback

    def set_on_resource_updated(self, callback: Callable[[str, str], None]) -> None:
        self._on_resource_updated = callback

    async def get(self, backend_id: str, url: str) -> ClientSession:
        if backend_id in self._workers:
            if await self._is_alive(backend_id):
                return self._workers[backend_id].session
            await self.remove(backend_id)

        logger.debug(f"Connecting to backend {backend_id} at {url}")
        tl_cb = (lambda: self._on_tool_list_changed(backend_id)) if self._on_tool_list_changed else None
        dc_cb = (lambda: self._handle_disconnect(backend_id)) if self._on_disconnect else None
        rl_cb = (lambda: self._on_resource_list_changed(backend_id)) if self._on_resource_list_changed else None
        ru_cb = (
            (lambda uri: self._on_resource_updated(backend_id, uri))
            if self._on_resource_updated
            else None
        )
        worker = _ConnectionWorker(
            url,
            on_tool_list_changed=tl_cb,
            on_disconnect=dc_cb,
            on_resource_list_changed=rl_cb,
            on_resource_updated=ru_cb,
        )
        session = await worker.start()
        self._workers[backend_id] = worker
        return session

    def _handle_disconnect(self, backend_id: str) -> None:
        self._workers.pop(backend_id, None)
        if self._on_disconnect:
            self._on_disconnect(backend_id)

    async def remove(self, backend_id: str) -> None:
        worker = self._workers.pop(backend_id, None)
        if worker:
            await worker.close()

    async def close_all(self) -> None:
        for backend_id in list(self._workers):
            await self.remove(backend_id)

    async def _is_alive(self, backend_id: str) -> bool:
        try:
            await self._workers[backend_id].session.send_ping()
            return True
        except Exception:
            return False
