"""LSP server lifecycle manager.

Manages running LSP server instances with:

- Reference-counted shared instances (multiple actors can share one server).
- Health checks via process liveness.
- Automatic restart on crash with re-initialization.
- Graceful shutdown when the last reference is released.

Based on ``docs/specification.md`` LSP Server Lifecycle (lines 20744-20758)
and issue #826.
"""

from __future__ import annotations

import threading
from typing import Any

import structlog

from cleveragents.lsp.client import LspClient
from cleveragents.lsp.errors import LspServerNotFoundError
from cleveragents.lsp.models import LspServerConfig
from cleveragents.lsp.transport import StdioTransport

logger = structlog.get_logger(__name__)


class _ManagedServer:
    """Internal state for a running LSP server instance.

    Attributes:
        config: The server configuration.
        transport: The stdio transport.
        client: The LSP protocol client.
        workspace_path: The workspace root path.
        ref_count: Number of active references.
    """

    __slots__ = ("client", "config", "ref_count", "transport", "workspace_path")

    def __init__(
        self,
        config: LspServerConfig,
        transport: StdioTransport,
        client: LspClient,
        workspace_path: str,
    ) -> None:
        self.config = config
        self.transport = transport
        self.client = client
        self.workspace_path = workspace_path
        self.ref_count = 1


class LspLifecycleManager:
    """Manages running LSP server instances with reference counting.

    Thread-safe.  Each server is identified by its namespaced name.
    Multiple callers can acquire/release the same server; the process
    is only stopped when the last reference is released.
    """

    def __init__(self) -> None:
        self._servers: dict[str, _ManagedServer] = {}
        self._lock = threading.Lock()

    def start_server(
        self,
        config: LspServerConfig,
        workspace_path: str,
    ) -> LspClient:
        """Start or acquire a reference to an LSP server.

        If a server with the same name is already running for the same
        workspace, its reference count is incremented and the existing
        client is returned.

        Args:
            config: Server configuration.
            workspace_path: Root directory for the workspace.

        Returns:
            An :class:`LspClient` connected to the running server.

        Raises:
            LspError: If the server process cannot be started.
        """
        # Phase 1: Check for existing server (short lock)
        with self._lock:
            existing = self._servers.get(config.name)
            if existing is not None and existing.transport.is_alive:
                existing.ref_count += 1
                logger.info(
                    "lsp.lifecycle.reused",
                    server=config.name,
                    ref_count=existing.ref_count,
                )
                return existing.client

        # Phase 2: Start and initialize OUTSIDE the lock (may take seconds)
        transport = StdioTransport(
            command=config.command,
            args=config.args,
            env=config.env,
            cwd=workspace_path,
        )
        transport.start()

        client = LspClient(transport, server_name=config.name)
        try:
            client.initialize(workspace_path)
        except Exception:
            transport.stop()
            raise

        # Phase 3: Register the running server (short lock)
        managed = _ManagedServer(
            config=config,
            transport=transport,
            client=client,
            workspace_path=workspace_path,
        )
        with self._lock:
            # Another thread may have started the same server concurrently
            race_winner = self._servers.get(config.name)
            if race_winner is not None and race_winner.transport.is_alive:
                # We lost the race — stop our instance, reuse theirs
                transport.stop()
                race_winner.ref_count += 1
                return race_winner.client

            self._servers[config.name] = managed

        logger.info(
            "lsp.lifecycle.started",
            server=config.name,
            workspace=workspace_path,
        )
        return client

    def stop_server(self, name: str) -> None:
        """Release a reference to an LSP server.

        Decrements the reference count.  When it reaches zero, the
        server is shut down gracefully.

        Args:
            name: Namespaced server name.

        Raises:
            LspServerNotFoundError: If no server with *name* is running.
        """
        with self._lock:
            managed = self._servers.get(name)
            if managed is None:
                raise LspServerNotFoundError(name)

            managed.ref_count -= 1
            if managed.ref_count > 0:
                logger.info(
                    "lsp.lifecycle.released",
                    server=name,
                    ref_count=managed.ref_count,
                )
                return

            # Last reference — shut down
            self._shutdown_server(managed)
            del self._servers[name]

    def _shutdown_server(self, managed: _ManagedServer) -> None:
        """Gracefully shut down a managed server."""
        name = managed.config.name
        logger.info("lsp.lifecycle.shutting_down", server=name)
        try:
            managed.client.shutdown()
        except Exception:
            logger.warning(
                "lsp.lifecycle.shutdown_error",
                server=name,
                exc_info=True,
            )
        managed.transport.stop()

    def get_client(self, name: str) -> LspClient:
        """Get the client for a running server.

        Args:
            name: Namespaced server name.

        Returns:
            The :class:`LspClient` for the server.

        Raises:
            LspServerNotFoundError: If no server with *name* is running.
        """
        with self._lock:
            managed = self._servers.get(name)
            if managed is None:
                raise LspServerNotFoundError(name)
            return managed.client

    def health_check(self, name: str) -> bool:
        """Check whether the named server is alive.

        Args:
            name: Namespaced server name.

        Returns:
            ``True`` if the server process is running.
        """
        with self._lock:
            managed = self._servers.get(name)
            if managed is None:
                return False
            return managed.transport.is_alive

    def restart_server(self, name: str) -> LspClient:
        """Restart a crashed or unresponsive server.

        Stops the current process (if alive), starts a fresh one, and
        re-initializes the LSP handshake.

        Uses the same 3-phase lock pattern as :meth:`start_server` to avoid
        holding ``self._lock`` during blocking I/O:

        - **Phase 1** (short lock): Read current state and remove the old
          entry so other threads see the server as gone during restart.
        - **Phase 2** (no lock): Stop the old transport, spawn the new
          process, and perform the LSP handshake.  These operations can
          block for up to 60 seconds and must not hold the lock.
        - **Phase 3** (short lock): Commit the new client into shared state.

        Args:
            name: Namespaced server name.

        Returns:
            A new :class:`LspClient` for the restarted server.

        Raises:
            LspServerNotFoundError: If the server was never started.
            LspError: If the restart fails.
        """
        # Phase 1: Read current state under lock (short critical section)
        with self._lock:
            managed = self._servers.get(name)
            if managed is None:
                raise LspServerNotFoundError(name)

            logger.warning("lsp.lifecycle.restarting", server=name)

            # Snapshot the fields we need for Phase 2 so we can release the lock
            old_transport = managed.transport
            config = managed.config
            workspace_path = managed.workspace_path
            ref_count = managed.ref_count

            # Remove the entry so concurrent callers see the server as absent
            # during the blocking restart.  Phase 3 will re-insert it.
            del self._servers[name]

        # Phase 2: Blocking I/O — no lock held
        old_transport.stop()

        transport = StdioTransport(
            command=config.command,
            args=config.args,
            env=config.env,
            cwd=workspace_path,
        )
        transport.start()

        client = LspClient(transport, server_name=name)
        try:
            client.initialize(workspace_path)
        except Exception:
            transport.stop()
            raise

        # Phase 3: Commit new state under lock (short critical section)
        new_managed = _ManagedServer(
            config=config,
            transport=transport,
            client=client,
            workspace_path=workspace_path,
        )
        new_managed.ref_count = ref_count

        with self._lock:
            self._servers[name] = new_managed

        logger.info("lsp.lifecycle.restarted", server=name)
        return client

    def stop_all(self) -> None:
        """Shut down all running servers.  Used during cleanup."""
        with self._lock:
            for name, managed in list(self._servers.items()):
                try:
                    self._shutdown_server(managed)
                except Exception:
                    logger.warning(
                        "lsp.lifecycle.stop_all_error",
                        server=name,
                        exc_info=True,
                    )
            self._servers.clear()

    def list_running(self) -> list[dict[str, Any]]:
        """Return status info for all running servers."""
        with self._lock:
            result: list[dict[str, Any]] = []
            for name, managed in self._servers.items():
                result.append(
                    {
                        "name": name,
                        "workspace": managed.workspace_path,
                        "alive": managed.transport.is_alive,
                        "ref_count": managed.ref_count,
                        "initialized": managed.client.is_initialized,
                    }
                )
            return result


__all__ = [
    "LspLifecycleManager",
]
