"""
ASGI middleware that tracks requests-in-flight.

This is a raw ASGI middleware (not Starlette's BaseHTTPMiddleware) to avoid
the known double-buffering and streaming issues with the base class.
"""

import re
import time
from typing import List, Optional, Tuple

from starlette.types import ASGIApp, Receive, Scope, Send

from asgi_runway._metrics import (
    APP_NAME,
    REQUEST_DURATION_SECONDS,
    REQUESTS_IN_FLIGHT,
    REQUESTS_IN_FLIGHT_BY_ROUTE,
    REQUESTS_TOTAL,
)

# A route group is a (compiled-regex, label) pair.
# Paths matching the pattern are tracked under the given label.
RouteGroup = Tuple[re.Pattern, str]


class RunwayMiddleware:
    """
    ASGI middleware that instruments every HTTP request with:

    - ``runway_requests_in_flight`` (Gauge) — primary autoscaling signal.
      Represents how many requests this pod is actively handling right now.
    - ``runway_requests_total`` (Counter) — requests by method + status code.
    - ``runway_request_duration_seconds`` (Histogram) — latency by method.

    Usage with FastAPI::

        from fastapi import FastAPI
        from asgi_runway import RunwayMiddleware, metrics_router

        app = FastAPI()
        app.add_middleware(RunwayMiddleware)
        app.include_router(metrics_router)

    For per-route granularity, pass ``route_groups``::

        app.add_middleware(
            RunwayMiddleware,
            route_groups=[
                (r"^/api/infer", "inference"),
                (r"^/api/embed", "embedding"),
            ],
        )

    Parameters
    ----------
    app:
        The wrapped ASGI application.
    exclude_paths:
        Paths that should NOT be counted in the metrics (e.g. health checks,
        the /metrics endpoint itself). Defaults to ``["/metrics", "/healthz", "/health"]``.
    route_groups:
        Optional list of ``(regex_pattern, label)`` tuples. When a request path
        matches a pattern, ``runway_requests_in_flight_by_route`` is also updated
        with the corresponding label. First match wins.
    """

    def __init__(
        self,
        app: ASGIApp,
        exclude_paths: Optional[List[str]] = None,
        route_groups: Optional[List[Tuple[str, str]]] = None,
    ) -> None:
        self.app = app
        self._exclude = set(exclude_paths or ["/metrics", "/healthz", "/health"])
        self._route_groups: List[RouteGroup] = [
            (re.compile(pattern), label) for pattern, label in (route_groups or [])
        ]

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        path: str = scope.get("path", "")
        if path in self._exclude:
            await self.app(scope, receive, send)
            return

        method: str = scope.get("method", "UNKNOWN").upper()
        route_label: Optional[str] = self._match_route(path)

        # --- Increment in-flight counters ---
        REQUESTS_IN_FLIGHT.inc()
        if route_label:
            REQUESTS_IN_FLIGHT_BY_ROUTE.labels(app_name=APP_NAME, route=route_label).inc()

        start = time.perf_counter()
        status_code = 500

        async def _send_wrapper(message: dict) -> None:
            nonlocal status_code
            if message["type"] == "http.response.start":
                status_code = message.get("status", 500)
            await send(message)

        try:
            await self.app(scope, receive, _send_wrapper)
        finally:
            duration = time.perf_counter() - start

            # --- Decrement in-flight counters ---
            REQUESTS_IN_FLIGHT.dec()
            if route_label:
                REQUESTS_IN_FLIGHT_BY_ROUTE.labels(app_name=APP_NAME, route=route_label).dec()

            # --- Update observability counters ---
            REQUESTS_TOTAL.labels(app_name=APP_NAME, method=method, status=str(status_code)).inc()
            REQUEST_DURATION_SECONDS.labels(app_name=APP_NAME, method=method).observe(duration)

    def _match_route(self, path: str) -> Optional[str]:
        for pattern, label in self._route_groups:
            if pattern.search(path):
                return label
        return None
