"""Internal MCP server that lets Navi push UI components to the webclient.

The server exposes a single tool, ``render_component``.  When Navi calls it,
the payload is forwarded over the active WebSocket for the session as a
``ui_component`` event.  The webclient receives the event and renders the
named component with the supplied data.

This server is started by ``navi.main:lifespan`` on a dedicated port
(``NAVI_UI_MCP_PORT``, default 8001) before the main container is built so
that :class:`navi.mcp.McpManager` can connect to it during startup like any
other MCP server.
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any

import structlog
from mcp.server import FastMCP

logger = logging.getLogger(__name__)

# Shared holder: the lifespan task sets the orchestrator once the main
# FastAPI container is ready.  The tool waits (with a timeout) so that early
# MCP health-check pings or tool calls do not crash before startup completes.
_orchestrator: Any | None = None
_orchestrator_ready = asyncio.Event()

_ORCHESTRATOR_TIMEOUT = 10.0

mcp = FastMCP(
    "navi_ui",
    host="127.0.0.1",
    port=8001,
    streamable_http_path="/mcp",
    instructions="""\
Internal Navi UI server. Use this to render structured data in the webclient.

Tool: render_component(component_name, payload, session_id)
- component_name: identifier of the UI component to render.
- payload: JSON-serializable object with the data for that component.
- session_id: target Navi session (injected automatically by the agent).

Only call this when the user explicitly asks for a graphical / structured
view, or when a tool response is naturally represented as a component.
""",
)


@mcp.tool()
async def render_component(
    component_name: str,
    payload: dict[str, Any],
    session_id: str | None = None,
) -> str:
    """Render a UI component in the webclient for the given session.

    Args:
        component_name: Identifier of the registered UI component (e.g. "table").
        payload: JSON-serializable data object consumed by that component.
        session_id: Navi session id. Injected automatically by the agent context.

    Returns:
        A short confirmation string, or an error message if forwarding failed.
    """
    if not component_name or not isinstance(component_name, str):
        return "Error: component_name must be a non-empty string"
    if not isinstance(payload, dict):
        return "Error: payload must be a JSON object (dict)"
    if not session_id:
        return "Error: session_id is required (it is normally injected by the agent)"

    try:
        await asyncio.wait_for(
            _orchestrator_ready.wait(),
            timeout=_ORCHESTRATOR_TIMEOUT,
        )
    except asyncio.TimeoutError:
        return "Error: UI server is not ready (orchestrator unavailable)"

    orchestrator = _orchestrator
    if orchestrator is None:
        return "Error: UI server orchestrator reference is missing"

    try:
        await orchestrator._notify_session(
            session_id,
            {
                "type": "ui_component",
                "component": component_name,
                "payload": payload,
            },
        )
    except Exception as exc:
        logger.warning("render_component failed for session %s: %s", session_id, exc)
        return f"Error: failed to send component to session: {exc}"

    return f"Component {component_name!r} rendered for session {session_id}"


def set_orchestrator(orchestrator: Any) -> None:
    """Wire the active orchestrator so the tool can reach WebSocket clients."""
    global _orchestrator
    _orchestrator = orchestrator
    _orchestrator_ready.set()
    logger.info("navi_ui orchestrator wired")


def clear_orchestrator() -> None:
    """Clear the orchestrator reference, e.g. during graceful shutdown."""
    global _orchestrator
    _orchestrator = None
    _orchestrator_ready.clear()


async def start_ui_server(host: str, port: int) -> None:
    """Start the StreamableHTTP MCP server and run forever.

    This coroutine is intended to be scheduled as a background task from the
    main application lifespan.  It only returns when the server is stopped.
    """
    log = structlog.get_logger()
    log.info("navi_ui_mcp_server_starting", host=host, port=port)
    # FastMCP copies constructor args into its settings object, so update them
    # explicitly to honour runtime configuration.
    mcp.settings.host = host
    mcp.settings.port = port
    await mcp.run_streamable_http_async()