"""ServiceBridge middleware for FastAPI / Starlette.

Usage::

    from fastapi import FastAPI
    from service_bridge.http.fastapi import ServiceBridgeMiddleware

    app = FastAPI()
    app.add_middleware(ServiceBridgeMiddleware, client=svc)
"""

from __future__ import annotations

import threading
from typing import TYPE_CHECKING, Optional

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import Response
from starlette.types import ASGIApp

if TYPE_CHECKING:
    from service_bridge.servicebridge import ServiceBridge

# Request-scoped key for the SDK client
_CLIENT_KEY = "servicebridge_client"
_TRACE_ID_KEY = "servicebridge_trace_id"
_SPAN_ID_KEY = "servicebridge_span_id"


def get_client(request: Request) -> Optional["ServiceBridge"]:
    """Retrieve the injected ServiceBridge client from a FastAPI request."""
    return getattr(request.state, _CLIENT_KEY, None)


class ServiceBridgeMiddleware(BaseHTTPMiddleware):
    """Starlette/FastAPI middleware that adds trace propagation and HTTP catalog registration.

    Features:
    - Reads ``traceparent`` or ``x-trace-id`` from incoming headers.
    - Sets ``x-trace-id`` on the response.
    - Injects the SDK client into ``request.state.servicebridge_client``.
    - Auto-registers routes in the ServiceBridge HTTP catalog on first request.

    Args:
        app: The ASGI application.
        client: ServiceBridge SDK client.
        exclude_paths: Path prefixes to skip tracing.
        propagate_trace_header: Write ``x-trace-id`` to response (default: True).
        auto_register: Register routes on first hit (default: True).
    """

    def __init__(
        self,
        app: ASGIApp,
        client: "ServiceBridge",
        exclude_paths: tuple[str, ...] = (),
        propagate_trace_header: bool = True,
        auto_register: bool = True,
    ) -> None:
        super().__init__(app)
        self._client = client
        self._exclude_paths = exclude_paths
        self._propagate = propagate_trace_header
        self._auto_register = auto_register
        self._registered: set[str] = set()
        self._reg_lock = threading.Lock()

    async def dispatch(self, request: Request, call_next) -> Response:
        from service_bridge.http_catalog import extract_trace_from_headers

        path = request.url.path or "/"

        for prefix in self._exclude_paths:
            if path.startswith(prefix):
                setattr(request.state, _CLIENT_KEY, self._client)
                return await call_next(request)

        headers = dict(request.headers)
        trace_id, parent_span_id = extract_trace_from_headers(headers)

        span = self._client.start_http_span(
            method=request.method,
            path=path,
            trace_id=trace_id,
            parent_span_id=parent_span_id,
        )

        setattr(request.state, _CLIENT_KEY, self._client)
        setattr(request.state, _TRACE_ID_KEY, span.trace_id)
        setattr(request.state, _SPAN_ID_KEY, span.span_id)

        status_code = 500
        try:
            response = await call_next(request)
            status_code = response.status_code
        except Exception:
            span.end(status_code=status_code)
            raise

        if self._propagate:
            response.headers["x-trace-id"] = span.trace_id

        span.end(status_code=status_code)

        if self._auto_register:
            pattern = _extract_pattern(request) or path
            key = f"{request.method}:{pattern}"
            with self._reg_lock:
                if key not in self._registered:
                    self._registered.add(key)
                    method = request.method

                    def _register(m=method, r=pattern, k=key) -> None:
                        import asyncio
                        try:
                            asyncio.run(self._client.register_http_endpoint(m, r))
                        except Exception:
                            with self._reg_lock:
                                self._registered.discard(k)

                    threading.Thread(target=_register, daemon=True).start()

        return response


def _extract_pattern(request: Request) -> str:
    """Try to get the matched route pattern from the request scope."""
    route = request.scope.get("route")
    if route is not None:
        path_fmt = getattr(route, "path", None)
        if path_fmt:
            return path_fmt
    return ""
