"""Middleware for cost tracking, auth, and request logging."""

from __future__ import annotations

import time

import structlog
from proxym.config import get_settings
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import Response

logger = structlog.get_logger()


class AuthMiddleware(BaseHTTPMiddleware):
    """Simple bearer token auth (matches LiteLLM master_key pattern)."""

    def __init__(self, app, master_key: str):
        super().__init__(app)
        self.master_key = master_key

    def _accepted_tokens(self) -> set[str]:
        settings = get_settings()
        return {
            self.master_key,
            settings.master_key,
            settings.default_api_key,
            "sk-proxy-local-dev",
        }

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        # Skip auth for health/docs/dashboard/favicon/config
        if request.url.path in ("/", "/health", "/docs", "/openapi.json", "/metrics", "/proxym-dashboard.jsx", "/dashboard-api.js", "/dashboard-hook.js", "/favicon.ico", "/config") or request.url.path.startswith("/dashboard") or request.url.path.startswith("/mcp/"):
            return await call_next(request)

        auth = request.headers.get("authorization", "")
        token = auth.removeprefix("Bearer ").strip()

        if not token or token not in self._accepted_tokens():
            return Response(
                content='{"error": "invalid api key"}',
                status_code=401,
                media_type="application/json",
            )

        return await call_next(request)


class CostTrackingMiddleware(BaseHTTPMiddleware):
    """Log request timing and attach cost headers to responses."""

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        t0 = time.monotonic()
        response = await call_next(request)
        elapsed = time.monotonic() - t0

        # Add timing header
        response.headers["X-Request-Time-Ms"] = f"{elapsed * 1000:.1f}"

        if request.url.path.startswith("/v1/"):
            logger.info(
                "request_completed",
                path=request.url.path,
                method=request.method,
                status=response.status_code,
                elapsed_ms=f"{elapsed * 1000:.1f}",
            )

        return response
