"""
cortex/api/middleware/usage_tracker.py — Usage tracking middleware for Cortex HTTP API.

Records every API request to stats.db for usage metering (Stripe billing,
dashboard analytics, system health). Fire-and-forget: DB write never
blocks the HTTP response.

Dependencies: cortex.api.stats_db (StatsDB).
"""
from __future__ import annotations

import asyncio
import logging
import time
from typing import Any

logger = logging.getLogger(__name__)


class UsageTrackingMiddleware:
    """Records every API request to stats.db for usage metering.

    Implements the raw ASGI interface (not BaseHTTPMiddleware) to avoid
    the response-body buffering overhead that BaseHTTPMiddleware incurs.
    The DB write is fire-and-forget: latency is measured before `send` so
    the response is never delayed.
    """

    def __init__(self, app: Any, stats_db: Any) -> None:
        self.app = app
        self._stats_db = stats_db

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

        start = time.perf_counter()
        status_code: int = 200

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

        try:
            await self.app(scope, receive, send_wrapper)
        finally:
            latency_ms = (time.perf_counter() - start) * 1000
            asyncio.ensure_future(
                self._record(scope, status_code, latency_ms)
            )

    async def _record(self, scope: Any, status_code: int, latency_ms: float) -> None:
        """Fire-and-forget DB write. Errors are logged, never raised."""
        try:
            # Extract owner_id from request state (set by auth middleware) or fall back
            owner_id: str = "anonymous"
            state = scope.get("state", {})
            if isinstance(state, dict):
                owner_id = state.get("owner_id", "anonymous")
            else:
                owner_id = getattr(state, "owner_id", "anonymous")

            path: str = scope.get("path", "/")
            method: str = scope.get("method", "GET")
            tokens_used: int = 0

            await self._stats_db.record_request(
                owner_id=owner_id,
                endpoint=path,
                method=method,
                status_code=status_code,
                latency_ms=round(latency_ms, 3),
                tokens_used=tokens_used,
            )
        except Exception:
            logger.exception("UsageTrackingMiddleware._record failed (non-fatal)")
