"""LSP server lifecycle pool.

Manages lazy-spawning, deduplication, and shutdown of LSP server processes.
Keyed by (server_id, root_path) — one client per server per project root.
"""

import asyncio
import shutil
import time
from pathlib import Path

from otto.log import get_logger
from otto.lsp.client import Diagnostic, DocumentSymbol, LSPClient
from otto.lsp.languages import ServerDef, find_servers, language_id

log = get_logger(__name__)

_INSTALL_TIMEOUT = 30
_IDLE_TTL = 600  # 10 minutes
_REAP_INTERVAL = 120  # 2 minutes


class LSPPool:
    """Pool of LSP server clients, lazily spawned per (server, root) pair."""

    def __init__(self) -> None:
        self._clients: dict[tuple[str, str], LSPClient] = {}
        self._spawning: dict[tuple[str, str], asyncio.Task] = {}
        self._broken: set[tuple[str, str]] = set()
        self._installed: set[str] = set()  # server IDs we've attempted install for
        self._reap_task: asyncio.Task | None = None

    async def get_clients(
        self, path: Path, *, block: bool = False, timeout: float = 10.0
    ) -> list[LSPClient]:
        """Get LSP clients for a file path.

        Ready clients are returned immediately.  If a server needs spawning
        (install + initialize), the spawn is kicked off in the background and
        the caller gets an empty list on that first call.  Subsequent calls
        return the now-warm client.  Pass ``block=True`` to wait for spawning
        (capped at *timeout* seconds to avoid stalling tool results).
        """
        servers = find_servers(path)
        clients: list[LSPClient] = []

        for server_def, root in servers:
            key = (server_def.id, str(root))
            if key in self._broken:
                continue

            client = self._clients.get(key)
            if client and not client.broken:
                clients.append(client)
                continue

            # Already spawning — optionally wait
            if key in self._spawning:
                if block:
                    try:
                        await asyncio.wait_for(asyncio.shield(self._spawning[key]), timeout=timeout)
                    except (TimeoutError, Exception):
                        pass
                    client = self._clients.get(key)
                    if client and not client.broken:
                        clients.append(client)
                continue

            # Kick off background spawn — don't block the tool result
            self._spawn_background(server_def, root)
            if block:
                task = self._spawning.get(key)
                if task:
                    try:
                        await asyncio.wait_for(asyncio.shield(task), timeout=timeout)
                    except (TimeoutError, Exception):
                        pass
                    client = self._clients.get(key)
                    if client and not client.broken:
                        clients.append(client)

        return clients

    def _spawn_background(self, server_def: ServerDef, root: Path) -> None:
        """Kick off a background spawn task (fire-and-forget)."""
        key = (server_def.id, str(root))
        if key in self._spawning or key in self._broken:
            return

        async def _do_spawn() -> None:
            try:
                client = await self._spawn_client(server_def, root)
                if client and not client.broken:
                    log.info("LSP server ready: %s", server_def.id)
            except Exception as exc:
                log.info("background spawn failed for %s: %s", server_def.id, exc)

        task = asyncio.create_task(_do_spawn())
        self._spawning[key] = task

        def _cleanup(t: asyncio.Task) -> None:
            self._spawning.pop(key, None)

        task.add_done_callback(_cleanup)

    def _start_reaper(self) -> None:
        """Start the idle-client reaper loop if not already running."""
        if self._reap_task is not None and not self._reap_task.done():
            return

        async def _reap_loop() -> None:
            while True:
                await asyncio.sleep(_REAP_INTERVAL)
                await self._reap_idle()

        self._reap_task = asyncio.create_task(_reap_loop())

    async def _reap_idle(self) -> None:
        """Shutdown clients that have been idle longer than _IDLE_TTL."""
        now = time.monotonic()
        to_reap = [
            key for key, client in self._clients.items() if now - client.last_used > _IDLE_TTL
        ]
        for key in to_reap:
            client = self._clients.pop(key, None)
            if client is not None:
                log.info("reaping idle LSP client: %s", key[0])
                try:
                    await client.shutdown()
                except Exception:
                    pass

    async def touch_file(self, path: Path, content: str) -> None:
        """Notify matching LSP clients about a file's content.

        Blocks for in-progress spawns so even the first write gets diagnostics.
        """
        lang = language_id(path)
        if lang is None:
            return
        resolved = path.resolve()
        clients = await self.get_clients(resolved, block=True)
        log.debug("touch_file: %s lang=%s clients=%d", resolved, lang, len(clients))
        for client in clients:
            try:
                await client.did_open(resolved, content, lang)
            except Exception:
                pass

    async def get_diagnostics(self, path: Path) -> list[Diagnostic]:
        """Collect diagnostics from all clients tracking this file."""
        resolved = path.resolve()
        clients = await self.get_clients(resolved, block=True)
        all_diags: list[Diagnostic] = []
        for client in clients:
            try:
                diags = await client.get_diagnostics(resolved)
                all_diags.extend(diags)
            except Exception:
                pass
        return all_diags

    async def get_document_symbols(self, path: Path) -> list[DocumentSymbol]:
        """Get document symbols from the first capable client."""
        resolved = path.resolve()
        lang = language_id(path)
        if lang is None:
            return []

        clients = await self.get_clients(resolved, block=True)
        for client in clients:
            try:
                # Open the file first so the server knows about it
                try:
                    text = resolved.read_text(encoding="utf-8", errors="replace")
                except OSError:
                    continue
                await client.did_open(resolved, text, lang)
                symbols = await client.document_symbols(resolved)
                if symbols:
                    return symbols
            except Exception:
                continue
        return []

    async def close(self) -> None:
        """Shutdown all server processes."""
        # Cancel the idle reaper
        if self._reap_task is not None and not self._reap_task.done():
            self._reap_task.cancel()
            try:
                await self._reap_task
            except asyncio.CancelledError:
                pass
        self._reap_task = None

        # Cancel any in-progress spawns
        for task in self._spawning.values():
            task.cancel()
        self._spawning.clear()

        # Shutdown all clients
        for client in self._clients.values():
            try:
                await client.shutdown()
            except Exception:
                pass
        self._clients.clear()

    async def _spawn_client(self, server_def: ServerDef, root: Path) -> LSPClient | None:
        """Spawn a new LSP client for a server definition."""
        key = (server_def.id, str(root))

        # Check if the command is available
        cmd_name = server_def.command[0]
        if not shutil.which(cmd_name):
            installed = await self._try_install(server_def)
            if not installed:
                log.info("server %s not available, skipping", server_def.id)
                self._broken.add(key)
                return None

        client = LSPClient(server_def.id, server_def.command, root)
        await client.initialize()

        if client.broken:
            self._broken.add(key)
            return None

        self._clients[key] = client
        self._start_reaper()
        return client

    async def _try_install(self, server_def: ServerDef) -> bool:
        """Attempt to install a missing LSP server. Returns True on success."""
        if server_def.install_cmd is None:
            return False
        if server_def.id in self._installed:
            return False  # Already tried once this session

        self._installed.add(server_def.id)
        log.info("auto-installing LSP server: %s", server_def.id)

        # Build list of install commands to try (with fallbacks for pip)
        attempts = [server_def.install_cmd]
        if server_def.install_cmd[0] == "pip":
            rest = server_def.install_cmd[1:]
            # Try uv pip first (fast), then pip3
            uv = shutil.which("uv")
            if uv:
                attempts.insert(0, (uv, "pip", "install", "--quiet") + rest[1:])
            pip3 = shutil.which("pip3")
            if pip3:
                attempts.append((pip3,) + rest)

        for cmd in attempts:
            if not shutil.which(cmd[0]):
                continue
            try:
                proc = await asyncio.create_subprocess_exec(
                    *cmd,
                    stdout=asyncio.subprocess.DEVNULL,
                    stderr=asyncio.subprocess.DEVNULL,
                )
                await asyncio.wait_for(proc.wait(), timeout=_INSTALL_TIMEOUT)
                if proc.returncode == 0:
                    if shutil.which(server_def.command[0]) is not None:
                        log.info("installed LSP server: %s", server_def.id)
                        return True
            except (TimeoutError, FileNotFoundError, OSError) as exc:
                log.info("install attempt failed for %s with %s: %s", server_def.id, cmd[0], exc)

        return False
