"""
Zarządza dynamicznym dodawaniem i usuwaniem backendów MCP.

Odpowiedzialność:
- Rejestruje nowy backend: łączy się przez Executor, pobiera list_tools, aktualizuje Registry
- Usuwa backend: czyści Registry, usuwa połączenie z ConnectionPool
- Odświeża narzędzia backendu (np. gdy backend dodał nowe narzędzia)
- Monitoruje połączenia: przy rozłączeniu usuwa narzędzia i notyfikuje, po powrocie przywraca
- Udostępnia metody: add_backend(id, url), remove_backend(id), refresh_backend(id)
- Jedyny komponent który ma prawo modyfikować Registry po starcie systemu
"""

import asyncio
import logging
from collections.abc import Callable

from .connection_pool import ConnectionPool
from .executor import Executor
from .gateway import Gateway
from .models import Backend, ResourceDescriptor, ResourceTemplateDescriptor, ToolDefinition
from .registry import Registry

logger = logging.getLogger(__name__)

RECONNECT_INTERVAL = 5.0


class Orchestrator:
    def __init__(
        self,
        registry: Registry,
        pool: ConnectionPool,
        executor: Executor,
        gateway: Gateway | None = None,
        on_change: Callable[[], None] | None = None,
    ):
        self._registry = registry
        self._pool = pool
        self._executor = executor
        self._gateway = gateway
        self._on_change = on_change
        self._reconnect_tasks: dict[str, asyncio.Task] = {}
        self._shutting_down = False
        if gateway:
            self._pool.set_on_tool_list_changed(self._schedule_refresh)
            self._pool.set_on_resource_list_changed(self._schedule_resource_refresh)
            self._pool.set_on_resource_updated(self._schedule_resource_updated)
        self._pool.set_on_disconnect(self._on_backend_disconnect)

    async def shutdown(self) -> None:
        """Cancel every outstanding reconnect-loop task and short-circuit
        further pool callbacks. Set first; pool teardown after server.serve()
        returns triggers on_disconnect for every backend, which would
        otherwise schedule fresh reconnect tasks into a dying event loop
        ("Task was destroyed but it is pending")."""
        self._shutting_down = True
        tasks = list(self._reconnect_tasks.values())
        for task in tasks:
            task.cancel()
        for task in tasks:
            try:
                await task
            except (asyncio.CancelledError, Exception):
                pass
        self._reconnect_tasks.clear()

    def _schedule_refresh(self, backend_id: str) -> None:
        if self._shutting_down:
            return
        asyncio.get_event_loop().create_task(self._refresh_and_notify(backend_id))

    async def _refresh_and_notify(self, backend_id: str) -> None:
        try:
            await self.refresh_backend(backend_id)
        except Exception as e:
            logger.warning(f"Failed to refresh backend {backend_id} after tool list change: {e}")

    def _schedule_resource_refresh(self, backend_id: str) -> None:
        if self._shutting_down:
            return
        asyncio.get_event_loop().create_task(self._refresh_resources_and_notify(backend_id))

    async def _refresh_resources_and_notify(self, backend_id: str) -> None:
        try:
            await self.refresh_backend_resources(backend_id)
        except Exception as e:
            logger.warning(f"Failed to refresh resources for {backend_id}: {e}")

    def _schedule_resource_updated(self, backend_id: str, original_uri: str) -> None:
        if self._shutting_down:
            return
        asyncio.get_event_loop().create_task(self._forward_resource_updated(backend_id, original_uri))

    async def _forward_resource_updated(self, backend_id: str, original_uri: str) -> None:
        if not self._gateway:
            return
        # Static resource: look up the public mapping. Template-resolved URIs
        # don't appear in `list_resources()` (they're concrete URIs the
        # backend produced from a template), so we forward them verbatim
        # — public form == original form for templates in our routing.
        public_uri = original_uri
        for mapping in self._registry.list_resources():
            if mapping.backend.id == backend_id and mapping.original_uri == original_uri:
                public_uri = mapping.public_uri
                break
        try:
            await self._gateway.notify_resource_updated(public_uri)
        except Exception as e:
            logger.debug(f"Failed to forward resource_updated for {public_uri}: {e}")

    def _on_backend_disconnect(self, backend_id: str) -> None:
        if self._shutting_down:
            return
        asyncio.get_event_loop().create_task(self._handle_disconnect(backend_id))

    async def _handle_disconnect(self, backend_id: str) -> None:
        if not self._registry.has_backend(backend_id):
            return
        self._registry.remove_backend_tools(backend_id)
        self._registry.remove_backend_resources(backend_id)
        self._registry.remove_backend_resource_templates(backend_id)
        logger.info(f"Backend {backend_id} disconnected — tools and resources removed")
        await self._notify()
        # Disconnect to stan połączenia, nie zmiana konfiguracji — config zapisujemy
        # tylko gdy user świadomie dodaje/usuwa/przełącza enabled.
        if backend_id not in self._reconnect_tasks:
            task = asyncio.get_event_loop().create_task(self._reconnect_loop(backend_id))
            self._reconnect_tasks[backend_id] = task

    async def _reconnect_loop(self, backend_id: str) -> None:
        try:
            while True:
                await asyncio.sleep(RECONNECT_INTERVAL)
                if not self._registry.has_backend(backend_id):
                    break
                backend = self._registry.get_backend(backend_id)
                if not backend.enabled:
                    break
                try:
                    backend = self._registry.get_backend(backend_id)
                    tools = await self._executor.list_tools(backend)
                    for tool in tools:
                        self._registry.register_tool(
                            backend,
                            ToolDefinition(
                                name=tool.name,
                                description=tool.description or "",
                                input_schema=tool.inputSchema,
                            ),
                            public_name=f"{backend_id}__{tool.name}",
                        )
                    resource_count = await self._load_resources(backend)
                    template_count = await self._load_resource_templates(backend)
                    logger.info(
                        f"Backend {backend_id} reconnected — {len(tools)} tools, "
                        f"{resource_count} resources, {template_count} templates restored"
                    )
                    await self._notify()
                    # Reconnect to stan połączenia, nie zmiana konfiguracji.
                    break
                except Exception as e:
                    logger.debug(f"Reconnect attempt for {backend_id} failed: {e}")
        finally:
            self._reconnect_tasks.pop(backend_id, None)

    async def _notify(self) -> None:
        if self._gateway:
            await self._gateway.notify_tool_list_changed()
            await self._gateway.notify_resource_list_changed()

    async def _load_resources(self, backend: Backend) -> int:
        """Pobiera resources z backendu i rejestruje. Zwraca liczbę zarejestrowanych.

        Jeśli backend nie deklaruje capability `resources` w handshake — pomija.
        Każdy inny błąd jest propagowany dalej.
        """
        session = await self._pool.get(backend.id, backend.url)
        capabilities = session.get_server_capabilities()
        # Diagnostic — pokazuje czego dokładnie się dowiedzieliśmy z handshake.
        # Do usunięcia po wykryciu źródła problemu z resources.
        logger.info(
            "Backend %s capabilities: %r (resources=%r)",
            backend.id,
            capabilities,
            getattr(capabilities, "resources", None) if capabilities else None,
        )
        if capabilities is None or capabilities.resources is None:
            return 0
        resources = await self._executor.list_resources(backend)
        for resource in resources:
            self._registry.register_resource(
                backend,
                ResourceDescriptor(
                    uri=str(resource.uri),
                    name=resource.name,
                    description=resource.description,
                    mime_type=resource.mimeType,
                ),
            )
        return len(resources)

    async def _load_resource_templates(self, backend: Backend) -> int:
        """Pobiera resource templates z backendu i rejestruje. Liczba zwrotu.

        Templates są aggregowane razem z statycznymi resources — backend musi
        deklarować `resources` capability, żeby się tym zająć. Zwraca 0 jeśli
        backend nie eksponuje żadnego template'a (typowe — większość MCP
        serwerów wystawia tylko statyczne URI)."""
        session = await self._pool.get(backend.id, backend.url)
        capabilities = session.get_server_capabilities()
        if capabilities is None or capabilities.resources is None:
            return 0
        try:
            templates = await self._executor.list_resource_templates(backend)
        except Exception as e:
            # Niektóre serwery nie wspierają tego endpointu nawet gdy
            # deklarują resources cap — traktujemy jako brak template'ów.
            logger.debug(f"Backend {backend.id}: list_resource_templates failed: {e}")
            return 0
        for template in templates:
            self._registry.register_resource_template(
                backend,
                ResourceTemplateDescriptor(
                    uri_template=str(template.uriTemplate),
                    name=template.name,
                    description=template.description,
                    mime_type=template.mimeType,
                ),
            )
        return len(templates)

    async def add_backend(
        self,
        backend_id: str,
        url: str,
        transient: bool = False,
        enabled: bool = True,
        persist: bool = True,
    ) -> None:
        """Rejestruje backend i (jeśli enabled) podłącza się do niego.

        `persist=False` używaj przy wczytywaniu listy z config.json przy starcie —
        plik jest już źródłem prawdy, nie ma sensu go nadpisywać. Wszystkie inne
        wywołania (z management toola, REST API itp.) używają persist=True i
        konfiguracja na dysku jest aktualizowana.
        """
        backend = Backend(id=backend_id, url=url, transient=transient, enabled=enabled)
        # Backend rejestrujemy ZAWSZE pierwszy — to fakt z konfiguracji,
        # niezależny od tego czy serwer akurat odpowiada. Wcześniejsza wersja
        # robiła list_tools przed add_backend, więc fail przy boot (serwer
        # jeszcze nie wstał) wywalał wyjątek, registry zostawało puste, a
        # najbliższy save_config (z innego on_change) zapisywał config.json
        # bez tego backendu — i przy następnym bootcie gateway już go nie znał.
        self._registry.add_backend(backend)
        if not enabled:
            logger.info(f"Added disabled backend {backend_id}")
            if persist and self._on_change and not transient:
                self._on_change()
            return

        try:
            tools = await self._executor.list_tools(backend)
            for tool in tools:
                self._registry.register_tool(
                    backend,
                    ToolDefinition(
                        name=tool.name,
                        description=tool.description or "",
                        input_schema=tool.inputSchema,
                    ),
                    public_name=f"{backend_id}__{tool.name}",
                )
            resource_count = await self._load_resources(backend)
            template_count = await self._load_resource_templates(backend)
            logger.info(
                f"Added backend {backend_id} with {len(tools)} tools, "
                f"{resource_count} resources, {template_count} templates"
            )
        except Exception as e:
            # Backend nieosiągalny — zostawiamy go w registry (config jest
            # autorytatywny), a reconnect_loop będzie próbował co RECONNECT_INTERVAL.
            # Bez tego pojedynczy nieżyjący serwer przy boot zaorałby konfigurację.
            logger.warning(
                f"Backend {backend_id} unreachable at startup: {e} — "
                f"will retry every {RECONNECT_INTERVAL}s"
            )
            await self._pool.remove(backend_id)
            if backend_id not in self._reconnect_tasks:
                task = asyncio.get_event_loop().create_task(self._reconnect_loop(backend_id))
                self._reconnect_tasks[backend_id] = task
        await self._notify()
        if persist and self._on_change and not transient:
            self._on_change()

    async def disable_backend(self, backend_id: str) -> None:
        backend = self._registry.get_backend(backend_id)
        backend.enabled = False
        task = self._reconnect_tasks.pop(backend_id, None)
        if task:
            task.cancel()
        self._registry.remove_backend_tools(backend_id)
        self._registry.remove_backend_resources(backend_id)
        await self._pool.remove(backend_id)
        logger.info(f"Disabled backend {backend_id}")
        await self._notify()
        if self._on_change and not backend.transient:
            self._on_change()

    async def enable_backend(self, backend_id: str) -> None:
        backend = self._registry.get_backend(backend_id)
        backend.enabled = True
        tools = await self._executor.list_tools(backend)
        for tool in tools:
            self._registry.register_tool(
                backend,
                ToolDefinition(
                    name=tool.name,
                    description=tool.description or "",
                    input_schema=tool.inputSchema,
                ),
                public_name=f"{backend_id}__{tool.name}",
            )
        resource_count = await self._load_resources(backend)
        logger.info(f"Enabled backend {backend_id} with {len(tools)} tools, {resource_count} resources")
        await self._notify()
        if self._on_change and not backend.transient:
            self._on_change()

    async def disable_all_backends(self) -> list[str]:
        """Wyłącza wszystkie backendy oprócz 'management'. Zwraca listę wyłączonych."""
        disabled = []
        for backend_id in list(self._registry._backends.keys()):
            if backend_id == "management":
                continue
            backend = self._registry._backends[backend_id]
            if backend.enabled:
                await self.disable_backend(backend_id)
                disabled.append(backend_id)
        return disabled

    async def remove_backend(self, backend_id: str) -> None:
        task = self._reconnect_tasks.pop(backend_id, None)
        if task:
            task.cancel()
        backend = self._registry.get_backend(backend_id) if self._registry.has_backend(backend_id) else None
        transient = backend.transient if backend else False
        self._registry.remove_backend(backend_id)
        await self._pool.remove(backend_id)
        logger.info(f"Removed backend {backend_id}")
        await self._notify()
        if self._on_change and not transient:
            self._on_change()

    async def refresh_backend(self, backend_id: str) -> None:
        backend = self._registry.get_backend(backend_id)
        url = backend.url
        # Usuwa bez notyfikacji — add_backend na końcu wyśle jedną finalną notyfikację
        self._registry.remove_backend(backend_id)
        await self._pool.remove(backend_id)
        await self.add_backend(backend_id, url)

    async def refresh_backend_resources(self, backend_id: str) -> None:
        """Odświeża listę resources i resource templates danego backendu."""
        backend = self._registry.get_backend(backend_id)
        self._registry.remove_backend_resources(backend_id)
        self._registry.remove_backend_resource_templates(backend_id)
        await self._load_resources(backend)
        await self._load_resource_templates(backend)
        if self._gateway:
            await self._gateway.notify_resource_list_changed()
