"""FastAPI router with all A2A protocol endpoints."""

from __future__ import annotations

import logging
from collections.abc import AsyncIterable
from typing import TYPE_CHECKING, Any

from a2a.types import (
    MessageSendParams,
    Task,
    TaskState,
)
from fastapi import APIRouter, Depends, Header, HTTPException, Path, Query, Request
from fastapi.responses import JSONResponse
from fastapi.sse import EventSourceResponse, ServerSentEvent

from a2akit.agent_card import AgentCardConfig, build_agent_card, external_base_url
from a2akit.middleware import A2AMiddleware, RequestEnvelope
from a2akit.schema import DirectReply
from a2akit.storage.base import (
    ListTasksQuery,
    UnsupportedOperationError,
)

if TYPE_CHECKING:
    from collections.abc import AsyncGenerator

    from a2akit.schema import StreamEvent
    from a2akit.task_manager import TaskManager

SUPPORTED_A2A_VERSION = "0.3.0"

logger = logging.getLogger(__name__)


def _sanitize_task_for_client(task: Task) -> Task:
    """Strip framework-internal metadata keys before client serialization.

    Keys prefixed with ``_`` are internal (e.g. ``_a2akit_direct_reply``,
    ``_idempotency_key``) and must not leak to external clients.
    """
    md = task.metadata
    if not md:
        return task
    cleaned = {k: v for k, v in md.items() if not k.startswith("_")}
    if len(cleaned) == len(md):
        return task  # nothing to strip
    return task.model_copy(update={"metadata": cleaned or None})


def _wrap_stream_event(event: StreamEvent) -> str:
    """Serialize a stream event for SSE (HTTP+JSON/REST transport).

    REST transport uses the object's ``kind`` field as discriminator.
    No wrapper envelope needed (unlike JSON-RPC transport).
    """
    if isinstance(event, DirectReply):
        return event.message.model_dump_json(by_alias=True, exclude_none=True)
    if isinstance(event, Task):
        event = _sanitize_task_for_client(event)
    return event.model_dump_json(by_alias=True, exclude_none=True)


def _check_a2a_version(
    a2a_version: str | None = Header(None, alias="A2A-Version"),
) -> None:
    """Validate the A2A-Version request header (spec §3.6.2).

    Pre-1.0: match Major.Minor exactly (breaking changes between minors).
    Post-1.0: match Major only (semver guarantees within major).
    Missing header is tolerated — spec says assume 0.3.
    """
    if a2a_version is None:
        return

    def _major_minor(v: str) -> tuple[str, str]:
        parts = v.strip().split(".")
        return (parts[0], parts[1] if len(parts) > 1 else "0")

    client = _major_minor(a2a_version)
    supported = _major_minor(SUPPORTED_A2A_VERSION)

    if supported[0] == "0":
        compatible = client == supported
    else:
        compatible = client[0] == supported[0]

    if not compatible:
        raise HTTPException(
            status_code=400,
            detail={
                "code": -32009,
                "message": (
                    f"Unsupported A2A version: {a2a_version}. "
                    f"This server supports {SUPPORTED_A2A_VERSION}."
                ),
            },
        )


def _get_tm(request: Request) -> TaskManager:
    """Extract the TaskManager from app state."""
    tm: TaskManager | None = getattr(request.app.state, "task_manager", None)
    if tm is None:
        raise HTTPException(
            status_code=503,
            detail={"code": -32603, "message": "TaskManager not initialized"},
        )
    return tm


def _check_streaming(request: Request) -> None:
    """Raise UnsupportedOperationError if streaming is not enabled."""
    caps = getattr(request.app.state, "capabilities", None)
    if caps is not None and not caps.streaming:
        raise UnsupportedOperationError("Streaming is not supported by this agent")


def _check_push_supported(request: Request) -> None:
    """Raise 501 if push notifications are not enabled."""
    caps = getattr(request.app.state, "capabilities", None)
    if not caps or not caps.push_notifications:
        raise HTTPException(
            status_code=501,
            detail={"code": -32003, "message": "Push notifications are not supported"},
        )


def _get_push_store(request: Request) -> Any:
    """Extract the PushConfigStore from app state."""
    store = getattr(request.app.state, "push_store", None)
    if store is None:
        raise HTTPException(
            status_code=501,
            detail={"code": -32003, "message": "Push notifications are not configured"},
        )
    return store


def _get_storage(request: Request) -> Any:
    """Extract the Storage from app state."""
    storage = getattr(request.app.state, "storage", None)
    if storage is None:
        raise HTTPException(
            status_code=503,
            detail={"code": -32603, "message": "Storage not initialized"},
        )
    return storage


def _validate_ids(params: MessageSendParams) -> MessageSendParams:
    """Validate that messageId is present."""
    msg = params.message
    if not msg.message_id or not msg.message_id.strip():
        raise HTTPException(
            status_code=400,
            detail={"code": -32600, "message": "messageId is required."},
        )
    return params


# Paths whose handlers run their own per-request middleware pipeline
# (they build a params-aware envelope themselves) and MUST NOT be
# double-processed by the router-level auth dependency.
_MIDDLEWARE_SELF_HANDLED_PATHS: frozenset[str] = frozenset(
    {
        "/v1/message:send",
        "/v1/message:stream",
    }
)

# Paths that are intentionally unauthenticated (liveness / readiness).
_MIDDLEWARE_PUBLIC_PATHS: frozenset[str] = frozenset(
    {
        "/v1/health",
        "/v1/health/ready",
    }
)


async def _enforce_middleware_pipeline(
    request: Request,
) -> AsyncGenerator[None, None]:
    """Router dependency: run the A2A middleware pipeline on every A2A endpoint.

    Spec §4.4 requires the server to authenticate EVERY incoming request.
    The message/send and message/stream handlers build a params-aware
    envelope themselves and run the pipeline with real ``MessageSendParams``;
    for every other endpoint (``tasks/*``, ``pushNotificationConfig/*``,
    ``/v1/card``) this dependency runs the pipeline with an empty envelope
    so that auth middlewares (Bearer, ApiKey, custom) fire uniformly.

    ``RequestEnvelope.params`` is ``None`` in this mode — middlewares that
    need params (e.g. TracingMiddleware) are expected to guard for it.
    """
    path = request.url.path
    if path in _MIDDLEWARE_SELF_HANDLED_PATHS or path in _MIDDLEWARE_PUBLIC_PATHS:
        yield
        return

    middlewares: list[A2AMiddleware] = getattr(request.app.state, "middlewares", [])
    if not middlewares:
        yield
        return

    envelope = RequestEnvelope()
    started: list[A2AMiddleware] = []
    try:
        for mw in middlewares:
            await mw.before_dispatch(envelope, request)
            started.append(mw)
        yield
    finally:
        # Only call after_dispatch on middlewares whose before_dispatch
        # succeeded — otherwise a middleware that raised during setup
        # (e.g. auth rejection) would get an after_dispatch without a
        # matching before_dispatch, and any middleware upstream of it
        # would leak resources (OTel span, context token, ...) because
        # its own after_dispatch was never reached.
        for mw in reversed(started):
            await mw.after_dispatch(envelope)


async def _stream_setup(
    request: Request,
    params: MessageSendParams,
) -> AsyncGenerator[
    tuple[
        tuple[str | None, StreamEvent],
        AsyncGenerator[tuple[str | None, StreamEvent], None],
        list[A2AMiddleware],
        RequestEnvelope,
    ],
    None,
]:
    """Yield-based dependency: validate, run middleware, fetch the first event.

    Uses ``yield`` instead of ``return`` so that FastAPI's dependency
    cleanup guarantees the ``finally`` block runs even if the route
    handler is never reached (e.g. client disconnect during response
    header preparation).  This prevents middleware resource leaks
    (OpenTelemetry spans, auth tokens) when the gap between dependency
    return and route handler execution is interrupted.
    """
    _check_streaming(request)
    params = _validate_ids(params)
    tm = _get_tm(request)
    middlewares: list[A2AMiddleware] = getattr(request.app.state, "middlewares", [])

    envelope = RequestEnvelope(params=params)
    started: list[A2AMiddleware] = []
    agen: AsyncGenerator[tuple[str | None, StreamEvent], None] | None = None
    try:
        for mw in middlewares:
            await mw.before_dispatch(envelope, request)
            started.append(mw)

        assert envelope.params is not None  # message endpoints always carry params
        agen = tm.stream_message(envelope.params, request_context=envelope.context)
        try:
            first_pair = await anext(agen)
        except BaseException:
            await agen.aclose()
            agen = None
            raise
    except BaseException:
        # Run after_dispatch only for middlewares that successfully
        # completed before_dispatch, in reverse order.
        for mw in reversed(started):
            await mw.after_dispatch(envelope)
        raise

    try:
        yield first_pair, agen, started, envelope
    finally:
        # Cleanup: close the event generator and run after_dispatch.
        # This runs even if the route handler was never entered.
        try:
            if agen is not None:
                await agen.aclose()
        finally:
            for mw in reversed(started):
                await mw.after_dispatch(envelope)


async def _subscribe_setup(
    request: Request,
    task_id: str = Path(),
    last_event_id: str | None = Header(None, alias="Last-Event-ID"),
) -> AsyncGenerator[
    tuple[tuple[str | None, StreamEvent], AsyncGenerator[tuple[str | None, StreamEvent], None]],
    None,
]:
    """Yield-based dependency: look up task, validate state, fetch first event.

    Uses ``yield`` (like ``_stream_setup``) so FastAPI's dependency
    cleanup guarantees ``agen.aclose()`` runs even if the client
    disconnects before the route handler iterates the SSE generator.
    With a plain ``return``, an early disconnect would leak the
    underlying Redis Pub/Sub connection held by ``agen``.
    """
    _check_streaming(request)
    tm = _get_tm(request)
    agen = tm.subscribe_task(task_id, after_event_id=last_event_id)
    try:
        first_pair = await anext(agen)
    except BaseException:
        await agen.aclose()
        raise

    try:
        yield first_pair, agen
    finally:
        await agen.aclose()


def build_a2a_router() -> APIRouter:
    """Build and return the complete A2A API router."""
    router = APIRouter(
        dependencies=[
            Depends(_check_a2a_version),
            Depends(_enforce_middleware_pipeline),
        ]
    )

    @router.post("/v1/message:send", tags=["Messages"])
    async def message_send(request: Request, params: MessageSendParams) -> JSONResponse:
        """Submit a message and return the task or message directly."""
        params = _validate_ids(params)
        tm = _get_tm(request)
        middlewares: list[A2AMiddleware] = getattr(request.app.state, "middlewares", [])

        envelope = RequestEnvelope(params=params)
        started: list[A2AMiddleware] = []

        try:
            for mw in middlewares:
                await mw.before_dispatch(envelope, request)
                started.append(mw)

            assert envelope.params is not None  # message endpoints always carry params
            result = await tm.send_message(envelope.params, request_context=envelope.context)
        except Exception:
            for mw in reversed(started):
                await mw.after_dispatch(envelope)
            raise

        for mw in reversed(started):
            await mw.after_dispatch(envelope, result)

        if isinstance(result, Task):
            result = _sanitize_task_for_client(result)
        return JSONResponse(
            content=result.model_dump(mode="json", by_alias=True, exclude_none=True)
        )

    @router.post("/v1/message:stream", response_class=EventSourceResponse, tags=["Messages"])
    async def message_stream(
        setup: tuple[
            tuple[str | None, StreamEvent],
            AsyncGenerator[tuple[str | None, StreamEvent], None],
            list[A2AMiddleware],
            RequestEnvelope,
        ] = Depends(_stream_setup),
    ) -> AsyncIterable[ServerSentEvent]:
        """Submit a message and stream events via SSE.

        All fallible setup (validation, middleware, first event) runs in
        the ``_stream_setup`` dependency — in the normal request context
        where exceptions produce proper HTTP error responses.

        This generator body runs inside FastAPI's SSE producer TaskGroup
        and only contains yield statements.  Middleware cleanup and
        generator close are handled by ``_stream_setup``'s finally block.

        SSE ``id:`` fields use the event-bus-assigned ID (not a local
        counter) so that ``Last-Event-ID`` reconnection maps correctly.
        """
        first_pair, agen, _middlewares, _envelope = setup
        try:
            eid, first_event = first_pair
            if not isinstance(first_event, DirectReply):
                yield ServerSentEvent(raw_data=_wrap_stream_event(first_event), id=eid)
            async for eid, ev in agen:
                if isinstance(ev, DirectReply):
                    continue
                yield ServerSentEvent(raw_data=_wrap_stream_event(ev), id=eid)
        except Exception:
            logger.exception("SSE stream aborted")

    @router.get("/v1/tasks/{task_id}", tags=["Tasks"])
    async def tasks_get(
        request: Request,
        task_id: str = Path(),
        history_length: int | None = Query(None, alias="historyLength"),
    ) -> JSONResponse:
        """Get a single task by ID."""
        tm = _get_tm(request)
        t = await tm.get_task(task_id, history_length)
        if not t:
            raise HTTPException(
                status_code=404, detail={"code": -32001, "message": "Task not found"}
            )
        t = _sanitize_task_for_client(t)
        return JSONResponse(content=t.model_dump(mode="json", by_alias=True, exclude_none=True))

    @router.get("/v1/tasks", tags=["Tasks"])
    async def tasks_list(
        request: Request,
        context_id: str | None = Query(None, alias="contextId"),
        status: TaskState | None = None,
        page_size: int = Query(50, alias="pageSize"),
        page_token: str | None = Query(None, alias="pageToken"),
        history_length: int | None = Query(None, alias="historyLength"),
        status_timestamp_after: str | None = Query(None, alias="statusTimestampAfter"),
        include_artifacts: bool = Query(False, alias="includeArtifacts"),
    ) -> JSONResponse:
        """List tasks with optional filters and pagination."""
        tm = _get_tm(request)
        query = ListTasksQuery(
            context_id=context_id,
            status=status,
            page_size=page_size,
            page_token=page_token,
            history_length=history_length,
            status_timestamp_after=status_timestamp_after,
            include_artifacts=include_artifacts,
        )
        result = await tm.list_tasks(query)
        result.tasks = [_sanitize_task_for_client(t) for t in result.tasks]
        return JSONResponse(
            content=result.model_dump(mode="json", by_alias=True, exclude_none=True)
        )

    @router.post("/v1/tasks/{task_id}:cancel", tags=["Tasks"])
    async def tasks_cancel(
        request: Request,
        task_id: str = Path(),
    ) -> JSONResponse:
        """Cancel a task by ID."""
        tm = _get_tm(request)
        result = await tm.cancel_task(task_id)
        result = _sanitize_task_for_client(result)
        return JSONResponse(
            content=result.model_dump(mode="json", by_alias=True, exclude_none=True)
        )

    @router.post(
        "/v1/tasks/{task_id}:subscribe", response_class=EventSourceResponse, tags=["Tasks"]
    )
    async def tasks_subscribe(
        setup: tuple[
            tuple[str | None, StreamEvent],
            AsyncGenerator[tuple[str | None, StreamEvent], None],
        ] = Depends(_subscribe_setup),
    ) -> AsyncIterable[ServerSentEvent]:
        """Subscribe to updates for an existing task via SSE.

        All fallible setup (task lookup, terminal-state check, first event)
        runs in the ``_subscribe_setup`` dependency — in the normal request
        context where exceptions produce proper HTTP error responses.
        """
        first_pair, agen = setup
        try:
            eid, first_event = first_pair
            yield ServerSentEvent(raw_data=_wrap_stream_event(first_event), id=eid)
            async for eid, ev in agen:
                if isinstance(ev, DirectReply):
                    continue
                yield ServerSentEvent(raw_data=_wrap_stream_event(ev), id=eid)
        except Exception:
            logger.exception("SSE subscribe stream aborted")

    @router.post("/v1/tasks/{task_id}/pushNotificationConfigs", tags=["Push Notifications"])
    async def push_config_set(request: Request, task_id: str = Path()) -> JSONResponse:
        """Set a push notification config for a task."""
        _check_push_supported(request)
        push_store = _get_push_store(request)
        storage = _get_storage(request)
        body = await request.json()
        from a2akit.push.endpoints import _serialize_tpnc, handle_set_config

        result = await handle_set_config(push_store, storage, task_id, body)
        return JSONResponse(content=_serialize_tpnc(result))

    @router.get(
        "/v1/tasks/{task_id}/pushNotificationConfigs/{config_id}",
        tags=["Push Notifications"],
    )
    async def push_config_get_by_id(
        request: Request,
        task_id: str = Path(),
        config_id: str = Path(),
    ) -> JSONResponse:
        """Get a specific push notification config."""
        _check_push_supported(request)
        push_store = _get_push_store(request)
        storage = _get_storage(request)
        from a2akit.push.endpoints import _serialize_tpnc, handle_get_config

        result = await handle_get_config(push_store, storage, task_id, config_id)
        return JSONResponse(content=_serialize_tpnc(result))

    @router.get("/v1/tasks/{task_id}/pushNotificationConfigs", tags=["Push Notifications"])
    async def push_config_list(request: Request, task_id: str = Path()) -> JSONResponse:
        """List all push notification configs for a task."""
        _check_push_supported(request)
        push_store = _get_push_store(request)
        storage = _get_storage(request)
        from a2akit.push.endpoints import _serialize_tpnc, handle_list_configs

        configs = await handle_list_configs(push_store, storage, task_id)
        return JSONResponse(content=[_serialize_tpnc(c) for c in configs])

    @router.delete(
        "/v1/tasks/{task_id}/pushNotificationConfigs/{config_id}",
        tags=["Push Notifications"],
    )
    async def push_config_delete(
        request: Request,
        task_id: str = Path(),
        config_id: str = Path(),
    ) -> JSONResponse:
        """Delete a push notification config."""
        _check_push_supported(request)
        push_store = _get_push_store(request)
        storage = _get_storage(request)
        from a2akit.push.endpoints import handle_delete_config

        await handle_delete_config(push_store, storage, task_id, config_id)
        return JSONResponse(content=None)

    @router.get("/v1/card", tags=["Discovery"])
    async def get_authenticated_extended_card(request: Request) -> JSONResponse:
        """Return the authenticated extended agent card."""
        provider = getattr(request.app.state, "extended_card_provider", None)
        if provider is None:
            raise HTTPException(
                status_code=404,
                detail={"code": -32007, "message": "Authenticated Extended Card not configured"},
            )
        extended_config: AgentCardConfig = await provider(request)
        base_url = external_base_url(
            dict(request.headers),
            request.url.scheme,
            request.url.netloc,
        )
        extra_protos = getattr(request.app.state, "additional_protocols", None)
        card = build_agent_card(extended_config, base_url, extra_protos)
        return JSONResponse(content=card.model_dump(mode="json", by_alias=True, exclude_none=True))

    @router.get("/v1/health", tags=["Health"])
    async def health_check() -> dict[str, str]:
        """Return a simple liveness status."""
        return {"status": "ok"}

    @router.get("/v1/health/ready", tags=["Health"])
    async def readiness_check(request: Request) -> JSONResponse:
        """Deep readiness check — pings all backend components."""
        components: dict[str, Any] = {}

        for name in ("storage", "broker", "event_bus"):
            backend = getattr(request.app.state, name, None)
            if backend is None:
                components[name] = {"status": "unavailable"}
                continue
            try:
                result = await backend.health_check()
            except Exception as exc:
                result = {"status": "error", "error": str(exc)}
            result["type"] = type(backend).__name__
            components[name] = result

        all_ok = all(c.get("status") == "ok" for c in components.values())
        status = "ok" if all_ok else "degraded"
        code = 200 if all_ok else 503

        return JSONResponse(
            content={"status": status, "components": components},
            status_code=code,
        )

    return router


def build_discovery_router(
    card_config: AgentCardConfig,
    additional_protocols: list[str] | None = None,
) -> APIRouter:
    """Build the agent card discovery router."""

    router = APIRouter()

    @router.get("/.well-known/agent-card.json", tags=["Discovery"])
    async def get_agent_card(request: Request) -> JSONResponse:
        """Serve the agent discovery card with the correct base URL."""
        base_url = external_base_url(
            dict(request.headers),
            request.url.scheme,
            request.url.netloc,
        )
        card = build_agent_card(card_config, base_url, additional_protocols)
        return JSONResponse(content=card.model_dump(mode="json", by_alias=True, exclude_none=True))

    return router
