"""Server lifecycle management — start, connect, and shutdown MCP servers.

Wraps the existing StdioSimulator to manage multiple MCP server processes
for the agent runtime. Handles:
- Starting servers from path (reads hdx.toml), command, or URL
- Tool discovery via MCP tools/list
- Graceful shutdown with process cleanup
"""

from __future__ import annotations

import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from hatchdx import HdxError
from hatchdx.agent.config import ServerConfig
from hatchdx.agent.runtime.tools import ToolDefinition
from hatchdx.harness.simulator import SimulatorError, StdioSimulator, ToolInfo


# ---------------------------------------------------------------------------
# Errors
# ---------------------------------------------------------------------------


class ServerConnectionError(HdxError):
    """Failed to start or connect to an MCP server."""


# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------


@dataclass
class ManagedServer:
    """A running MCP server managed by the agent runtime."""

    config: ServerConfig
    simulator: StdioSimulator | None = None
    tools: list[ToolInfo] = field(default_factory=list)
    transport: str = "stdio"
    server_info: dict[str, Any] = field(default_factory=dict)

    @property
    def is_running(self) -> bool:
        if self.simulator is None:
            return False
        return self.simulator._initialized


# ---------------------------------------------------------------------------
# Server manager
# ---------------------------------------------------------------------------


class ServerManager:
    """Manage the lifecycle of multiple MCP servers for an agent.

    Usage::

        async with ServerManager() as manager:
            servers = await manager.start_all(config.servers)
            tools = await manager.discover_all_tools()
            # ... use tools ...
        # all servers shut down on exit
    """

    def __init__(self) -> None:
        self._servers: list[ManagedServer] = []

    # -- Context manager ------------------------------------------------------

    async def __aenter__(self) -> ServerManager:
        return self

    async def __aexit__(self, *exc: object) -> None:
        await self.shutdown_all()

    # -- Properties -----------------------------------------------------------

    @property
    def servers(self) -> list[ManagedServer]:
        return list(self._servers)

    # -- Start servers --------------------------------------------------------

    async def start_server(
        self,
        config: ServerConfig,
        base_dir: Path | None = None,
    ) -> ManagedServer:
        """Start a single MCP server from its configuration.

        Args:
            config: Server configuration from agent.yaml.
            base_dir: Base directory for resolving relative paths.
                      Defaults to CWD.

        Returns:
            A ManagedServer with an active connection.

        Raises:
            ServerConnectionError: If the server fails to start.
        """
        base = base_dir or Path.cwd()

        if config.url:
            return self._register_remote_server(config)

        try:
            command, cwd = self._resolve_command(config, base)
            env = self._build_env(config)

            simulator = StdioSimulator(
                command,
                timeout=10.0,
                env=env,
                cwd=str(cwd) if cwd else None,
            )
            server_info = await simulator.start()

            server = ManagedServer(
                config=config,
                simulator=simulator,
                transport="stdio",
                server_info=server_info,
            )
            self._servers.append(server)
            return server

        except SimulatorError as e:
            raise ServerConnectionError(
                f"Failed to start server '{config.name}': {e}"
            ) from e
        except Exception as e:
            raise ServerConnectionError(
                f"Failed to start server '{config.name}': {e}"
            ) from e

    async def start_all(
        self,
        configs: list[ServerConfig],
        base_dir: Path | None = None,
    ) -> list[ManagedServer]:
        """Start all configured servers.

        If any server fails to start, all previously started servers
        are shut down before the error is raised.
        """
        started: list[ManagedServer] = []
        try:
            for config in configs:
                server = await self.start_server(config, base_dir=base_dir)
                started.append(server)
        except (ServerConnectionError, Exception):
            # Clean up any servers we already started and remove from manager
            for server in started:
                await self._shutdown_server(server)
                if server in self._servers:
                    self._servers.remove(server)
            raise
        return started

    # -- Tool discovery -------------------------------------------------------

    async def discover_tools(
        self, server: ManagedServer
    ) -> list[ToolDefinition]:
        """Discover tools from a single running server.

        Calls tools/list via the MCP protocol and converts the results
        to ToolDefinition objects for the agent runtime.
        """
        if server.simulator is None:
            return []

        try:
            tool_infos = await server.simulator.list_tools()
            server.tools = tool_infos

            return [
                ToolDefinition(
                    name=ti.name,
                    description=ti.description,
                    server_name=server.config.name,
                    input_schema=ti.input_schema,
                )
                for ti in tool_infos
            ]
        except SimulatorError as e:
            raise ServerConnectionError(
                f"Failed to discover tools from server '{server.config.name}': {e}"
            ) from e

    async def discover_all_tools(self) -> dict[str, list[ToolDefinition]]:
        """Discover tools from all managed servers.

        Returns:
            Mapping of server_name -> list of ToolDefinition, ready for
            collect_and_filter_tools().
        """
        result: dict[str, list[ToolDefinition]] = {}
        for server in self._servers:
            result[server.config.name] = await self.discover_tools(server)
        return result

    # -- Shutdown -------------------------------------------------------------

    async def shutdown_server(self, server: ManagedServer) -> None:
        """Shut down a single server and remove it from management."""
        await self._shutdown_server(server)
        if server in self._servers:
            self._servers.remove(server)

    async def shutdown_all(self) -> None:
        """Shut down all managed servers."""
        for server in self._servers:
            await self._shutdown_server(server)
        self._servers.clear()

    # -- Internal helpers -----------------------------------------------------

    def _register_remote_server(self, config: ServerConfig) -> ManagedServer:
        """Register a remote (URL-based) server without starting a process."""
        server = ManagedServer(
            config=config,
            simulator=None,
            transport="http",
        )
        self._servers.append(server)
        return server

    def _resolve_command(
        self, config: ServerConfig, base_dir: Path
    ) -> tuple[list[str], Path | None]:
        """Resolve the server command and working directory.

        Returns:
            Tuple of (command_list, working_directory).
        """
        if config.path:
            return self._resolve_path_command(config, base_dir)
        elif config.command:
            return self._resolve_direct_command(config)
        else:
            raise ServerConnectionError(
                f"Server '{config.name}': no path, command, or url specified"
            )

    def _resolve_path_command(
        self, config: ServerConfig, base_dir: Path
    ) -> tuple[list[str], Path]:
        """Resolve command from a server path (reads hdx.toml/pyproject.toml)."""
        from hatchdx.utils.config import ConfigError, load_config, resolve_server_command

        server_path = Path(config.path)  # type: ignore[arg-type]
        if not server_path.is_absolute():
            server_path = base_dir / server_path

        if not server_path.is_dir():
            raise ServerConnectionError(
                f"Server '{config.name}': path does not exist: {server_path}\n"
                "Check the 'path' value in agent.yaml"
            )

        try:
            project_config = load_config(server_path)
        except ConfigError as e:
            raise ServerConnectionError(
                f"Server '{config.name}': failed to read config from {server_path}: {e}"
            ) from e

        command = project_config.server_command.split()
        command = resolve_server_command(command, server_path)
        return command, server_path

    def _resolve_direct_command(
        self, config: ServerConfig
    ) -> tuple[list[str], None]:
        """Resolve a direct command string into a command list."""
        command = config.command.split()  # type: ignore[union-attr]
        if config.args:
            command.extend(config.args)
        return command, None

    @staticmethod
    def _build_env(config: ServerConfig) -> dict[str, str] | None:
        """Build environment dict for the subprocess.

        If the server config has env vars, merges them with the current
        environment (config values override). Returns None if no custom
        env is needed (inherits parent environment).
        """
        if not config.env:
            return None
        env = dict(os.environ)
        env.update(config.env)
        return env

    @staticmethod
    async def _shutdown_server(server: ManagedServer) -> None:
        """Shut down a single server's simulator (if running)."""
        if server.simulator is not None:
            try:
                await server.simulator.shutdown()
            except Exception:
                pass  # Best-effort cleanup
