"""FastMCP v3 middleware that broadcasts tool call results to the dashboard.

This module provides:
- ``EventType`` -- StrEnum of broadcast event categories
- ``EventPayload`` -- Pydantic model matching the dashboard ``IngestEventPayload``
- ``TOOL_EVENT_MAP`` -- mapping of MCP tool names to ``EventType`` values
- ``EventEmitterMiddleware`` -- FastMCP v3 ``Middleware`` subclass that fires
  fire-and-forget POSTs to ``/api/events/ingest`` for mapped tools
- ``get_dashboard_url`` -- reads ``DASHBOARD_URL`` env var

Architecture: ADR-001 (middleware interceptor), ADR-002 (dashboard/MCP decoupling),
ADR-003 (fire-and-forget), ADR-005 (conditional registration).
"""

from __future__ import annotations

import asyncio
import datetime
import json
import logging
import os
from enum import StrEnum
from typing import TYPE_CHECKING, Any

import httpx
from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext
from mcp.types import ImageContent, TextContent
from pydantic import BaseModel, Field

if TYPE_CHECKING:
    from fastmcp.tools.tool import ToolResult
    from mcp import types as mt

_log = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# EventType and EventPayload — must match dashboard IngestEventType /
# IngestEventPayload structurally (ADR-002: defined independently).
# ---------------------------------------------------------------------------


class EventType(StrEnum):
    """Event categories broadcast to WebSocket clients."""

    JOB_CREATED = "job_created"
    JOB_UPDATED = "job_updated"
    STEP_ADDED = "step_added"
    SCREENSHOT_CAPTURED = "screenshot_captured"
    VM_STATE_CHANGED = "vm_state_changed"
    FLEET_STATUS = "fleet_status"
    INSPECTOR_REPORT = "inspector_report"
    TODO_UPDATED = "todo_updated"
    STEERING_MESSAGE = "steering_message"


class EventPayload(BaseModel):
    """Event payload POSTed to the dashboard ingest endpoint.

    Field structure is identical to the dashboard ``IngestEventPayload`` so
    that both ends accept the same JSON body (ADR-002 parity requirement).
    """

    event_type: EventType
    vm_name: str = ""
    job_id: str = ""
    data: dict[str, Any] = Field(default_factory=dict)
    timestamp: str = ""

    model_config = {
        "json_schema_extra": {
            "examples": [
                {
                    "event_type": "job_created",
                    "vm_name": "my-vm",
                    "job_id": "job-123",
                    "data": {"status": "active"},
                    "timestamp": "2026-03-20T12:00:00+00:00",
                }
            ]
        }
    }


# ---------------------------------------------------------------------------

TOOL_EVENT_MAP: dict[str, EventType] = {
    "vm_job_start": EventType.JOB_CREATED,
    "vm_job_pause": EventType.JOB_UPDATED,
    "vm_job_resume": EventType.JOB_UPDATED,
    "vm_job_complete": EventType.JOB_UPDATED,
    "vm_job_step": EventType.STEP_ADDED,
    "vm_screenshot": EventType.SCREENSHOT_CAPTURED,
    "vm_screenshot_api": EventType.SCREENSHOT_CAPTURED,
    "vagrant_up": EventType.VM_STATE_CHANGED,
    "vagrant_destroy": EventType.VM_STATE_CHANGED,
    "vm_todo_add": EventType.TODO_UPDATED,
    "vm_todo_update": EventType.TODO_UPDATED,
    "vm_todo_reorder": EventType.TODO_UPDATED,
    "vm_steering_send": EventType.STEERING_MESSAGE,
    "vm_viewport_action": EventType.SCREENSHOT_CAPTURED,
}


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def get_dashboard_url() -> str | None:
    """Return the dashboard base URL from ``DASHBOARD_URL`` env var.

    Returns:
        The URL string, or ``None`` if the variable is unset or empty.
    """
    url = os.environ.get("DASHBOARD_URL")
    return url or None


def _extract_payload_from_tool_result(
    tool_name: str, event_type: EventType, result: ToolResult, arguments: dict[str, Any] | None = None
) -> EventPayload:
    """Build an ``EventPayload`` from a completed tool result.

    For ``TextContent`` results (dict-returning tools): parses the JSON text
    and extracts ``vm_name`` and ``job_id`` from the parsed dict.

    For ``ImageContent`` results (screenshot tools): sets
    ``data={"return_type": "Image"}`` and extracts ``vm_name`` from the tool
    call arguments.

    Args:
        tool_name: Name of the MCP tool that produced the result.
        event_type: The ``EventType`` mapped to this tool.
        result: The ``ToolResult`` returned by the tool.
        arguments: The tool call arguments dict (used for ImageContent tools).

    Returns:
        A populated ``EventPayload`` with ISO-8601 timestamp.
    """
    timestamp = datetime.datetime.now(datetime.UTC).isoformat()
    vm_name = ""
    job_id = ""
    data: dict[str, Any] = {}

    if result.content:
        first = result.content[0]
        if isinstance(first, TextContent):
            try:
                parsed: dict[str, Any] = json.loads(first.text)
                vm_name = str(parsed.get("vm_name", ""))
                job_id = str(parsed.get("job_id", ""))
                data = parsed
            except (json.JSONDecodeError, TypeError):
                data = {"raw": first.text}
        elif isinstance(first, ImageContent):
            data = {"return_type": "Image"}
            if arguments:
                vm_name = str(arguments.get("vm_name", ""))

    return EventPayload(event_type=event_type, vm_name=vm_name, job_id=job_id, data=data, timestamp=timestamp)


async def _post_event(client: httpx.AsyncClient, url: str, payload: EventPayload) -> None:
    """POST an event payload to the dashboard ingest endpoint.

    Fire-and-forget: all exceptions are caught and logged as warnings.
    No retry logic per ADR-003.

    Args:
        client: Shared ``httpx.AsyncClient`` instance.
        url: Full URL of the dashboard ingest endpoint.
        payload: The ``EventPayload`` to POST.
    """
    try:
        timeout = httpx.Timeout(connect=2.0, read=2.0, write=2.0, pool=2.0)
        await client.post(url, json=payload.model_dump(), timeout=timeout)
    except Exception:  # noqa: BLE001
        _log.warning(
            "Event emission failed for event_type=%s vm_name=%s", payload.event_type, payload.vm_name, exc_info=True
        )


# ---------------------------------------------------------------------------
# EventEmitterMiddleware
# ---------------------------------------------------------------------------


class EventEmitterMiddleware(Middleware):
    """FastMCP v3 middleware that broadcasts tool results to the dashboard.

    Intercepts ``on_call_tool`` for tools listed in ``TOOL_EVENT_MAP``.
    When a mapped tool completes, fires a background POST to the dashboard
    ``/api/events/ingest`` endpoint (fire-and-forget via ``asyncio.create_task``).

    The original tool result is always returned unchanged.

    When ``dashboard_url`` is ``None`` or empty, ``on_call_tool`` is a pure
    pass-through with zero side effects (ADR-005).

    Args:
        dashboard_url: Base URL of the dashboard server (e.g.
            ``"http://localhost:5173"``). Pass ``None`` to disable emission.
    """

    def __init__(self, dashboard_url: str | None = None) -> None:
        """Initialise the middleware with an optional dashboard base URL."""
        self._dashboard_url = dashboard_url
        self._ingest_url: str | None = f"{dashboard_url.rstrip('/')}/api/events/ingest" if dashboard_url else None
        self._client: httpx.AsyncClient | None = None
        # Keeps strong references to in-flight tasks so GC cannot cancel them
        # before they complete (required by RUF006).
        self._tasks: set[asyncio.Task[None]] = set()

    def _get_client(self) -> httpx.AsyncClient:
        """Lazily create a shared ``httpx.AsyncClient`` on first use."""
        if self._client is None:
            self._client = httpx.AsyncClient()
        return self._client

    async def on_call_tool(
        self,
        context: MiddlewareContext[mt.CallToolRequestParams],
        call_next: CallNext[mt.CallToolRequestParams, ToolResult],
    ) -> ToolResult:
        """Intercept tool calls and emit events for mapped tools.

        Args:
            context: Middleware context carrying the ``CallToolRequestParams``
                message with ``.name`` and ``.arguments``.
            call_next: Callable to invoke the next handler in the pipeline.

        Returns:
            The original ``ToolResult`` from the tool, always unchanged.
        """
        result = await call_next(context)

        if not self._ingest_url:
            return result

        tool_name: str = context.message.name
        event_type = TOOL_EVENT_MAP.get(tool_name)
        if event_type is None:
            return result

        arguments: dict[str, Any] | None = context.message.arguments
        payload = _extract_payload_from_tool_result(
            tool_name=tool_name, event_type=event_type, result=result, arguments=arguments
        )
        task = asyncio.create_task(_post_event(self._get_client(), self._ingest_url, payload), name=f"emit-{tool_name}")
        self._tasks.add(task)
        task.add_done_callback(self._tasks.discard)

        return result
