"""Server API endpoints — discover and invoke MCP server tools."""

from __future__ import annotations

import time
from pathlib import Path
from typing import Any

from aiohttp import web

from hatchdx.harness.simulator import StdioSimulator, SimulatorError
from hatchdx.utils.config import resolve_server_command
from hatchdx.utils.workspace import discover_servers

server_routes = web.RouteTableDef()

# Cache of active simulators keyed by server command string
_simulators: dict[str, StdioSimulator] = {}


def _get_server_config(
    workspace_root: Path | None = None,
) -> dict[str, dict[str, Any]]:
    """Load server configs from workspace servers or single project."""
    from hatchdx.utils.config import load_config

    if workspace_root is not None:
        configs: dict[str, dict[str, Any]] = {}
        for server_dir in discover_servers(workspace_root):
            try:
                config = load_config(project_dir=server_dir)
                configs[config.server_name] = {
                    "command": config.server_command.split(),
                    "name": config.server_name,
                    "project_dir": str(server_dir),
                }
            except Exception:
                continue
        return configs

    # Non-workspace: single project at CWD
    try:
        config = load_config()
        return {
            config.server_name: {
                "command": config.server_command.split(),
                "name": config.server_name,
                "project_dir": str(Path.cwd()),
            }
        }
    except Exception:
        return {}


async def _get_simulator(
    name: str, workspace_root: Path | None = None
) -> StdioSimulator:
    """Get or create a simulator for the named server."""
    if name in _simulators:
        return _simulators[name]

    servers = _get_server_config(workspace_root=workspace_root)
    if name not in servers:
        raise web.HTTPNotFound(text=f"Server '{name}' not found in config")

    project_dir = Path(servers[name]["project_dir"])
    cmd = resolve_server_command(servers[name]["command"], project_dir)
    sim = StdioSimulator(cmd, timeout=10.0, cwd=str(project_dir))
    await sim.__aenter__()
    _simulators[name] = sim
    return sim


@server_routes.get("/api/servers")
async def list_servers(request: web.Request) -> web.Response:
    workspace_root: Path | None = request.app.get("workspace_root")
    servers = _get_server_config(workspace_root=workspace_root)
    result = []
    for name, cfg in servers.items():
        connected = name in _simulators
        tool_count = 0
        if connected:
            try:
                tools = await _simulators[name].list_tools()
                tool_count = len(tools)
            except Exception:
                connected = False

        result.append({
            "name": name,
            "status": "connected" if connected else "disconnected",
            "transport": "stdio",
            "tool_count": tool_count,
        })

    return web.json_response({"servers": result})


@server_routes.get("/api/servers/{name}/tools")
async def list_tools(request: web.Request) -> web.Response:
    name = request.match_info["name"]
    workspace_root: Path | None = request.app.get("workspace_root")
    try:
        sim = await _get_simulator(name, workspace_root=workspace_root)
        tools = await sim.list_tools()
        return web.json_response({
            "server": name,
            "tools": [
                {
                    "name": t.name,
                    "description": t.description,
                    "input_schema": t.input_schema,
                }
                for t in tools
            ],
        })
    except SimulatorError as e:
        return web.json_response({"error": str(e)}, status=502)
    except web.HTTPNotFound:
        return web.json_response({"error": f"Server '{name}' not found"}, status=404)


@server_routes.post("/api/servers/{name}/tools/{tool}/call")
async def call_tool(request: web.Request) -> web.Response:
    name = request.match_info["name"]
    tool = request.match_info["tool"]

    try:
        body = await request.json()
    except Exception:
        body = {}

    arguments = body.get("arguments", {})

    workspace_root: Path | None = request.app.get("workspace_root")
    try:
        sim = await _get_simulator(name, workspace_root=workspace_root)
        start = time.monotonic()
        result = await sim.call_tool(tool, arguments)
        latency_ms = (time.monotonic() - start) * 1000

        return web.json_response({
            "content": result.content,
            "is_error": result.is_error,
            "latency_ms": round(latency_ms, 1),
        })
    except SimulatorError as e:
        return web.json_response({"error": str(e), "is_error": True}, status=502)
    except web.HTTPNotFound:
        return web.json_response({"error": f"Server '{name}' not found"}, status=404)


@server_routes.get("/api/servers/{name}/health")
async def server_health(request: web.Request) -> web.Response:
    name = request.match_info["name"]
    workspace_root: Path | None = request.app.get("workspace_root")
    servers = _get_server_config(workspace_root=workspace_root)

    if name not in servers:
        return web.json_response({"server": name, "status": "not_found"}, status=404)

    if name in _simulators:
        try:
            await _simulators[name].list_tools()
            return web.json_response({"server": name, "status": "healthy"})
        except Exception:
            # Clean up dead simulator
            try:
                await _simulators[name].__aexit__(None, None, None)
            except Exception:
                pass
            del _simulators[name]

    return web.json_response({"server": name, "status": "disconnected"})
