from __future__ import annotations

from contextvars import ContextVar
from datetime import UTC, datetime
import json
import logging
from logging.config import dictConfig
from time import perf_counter
from uuid import uuid4


request_id_context: ContextVar[str] = ContextVar("request_id", default="-")


class JsonFormatter(logging.Formatter):
    def format(self, record: logging.LogRecord) -> str:
        payload = {
            "timestamp": datetime.now(UTC).isoformat(),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
            "request_id": getattr(record, "request_id", None) or request_id_context.get(),
        }
        for key in ("method", "path", "status_code", "duration_ms", "client_ip"):
            value = getattr(record, key, None)
            if value is not None:
                payload[key] = value
        if record.exc_info:
            payload["exception"] = self.formatException(record.exc_info)
        return json.dumps(payload, ensure_ascii=False, separators=(",", ":"))


def configure_logging(level: str = "INFO") -> None:
    dictConfig(
        {
            "version": 1,
            "disable_existing_loggers": False,
            "formatters": {"json": {"()": JsonFormatter}},
            "handlers": {"console": {"class": "logging.StreamHandler", "formatter": "json", "stream": "ext://sys.stdout"}},
            "loggers": {
                "app": {"handlers": ["console"], "level": level, "propagate": False},
                "app.access": {"handlers": ["console"], "level": level, "propagate": False},
            },
        }
    )


class RequestObservabilityMiddleware:
    def __init__(self, app) -> None:
        self.app = app
        self.logger = logging.getLogger("app.access")

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

        headers = {key.lower(): value for key, value in scope.get("headers", [])}
        incoming_id = headers.get(b"x-request-id", b"").decode("ascii", errors="ignore").strip()
        request_id = incoming_id[:64] if incoming_id else uuid4().hex
        token = request_id_context.set(request_id)
        started_at = perf_counter()
        status_code = 500
        response_finished = False

        async def send_with_request_id(message) -> None:
            nonlocal status_code, response_finished
            if message["type"] == "http.response.start":
                status_code = int(message["status"])
                message.setdefault("headers", []).append((b"x-request-id", request_id.encode("ascii")))
            elif message["type"] == "http.response.body" and not message.get("more_body", False):
                response_finished = True
            await send(message)

        try:
            await self.app(scope, receive, send_with_request_id)
        except Exception:
            self._log(scope, request_id, 500, started_at, logging.ERROR, "request_failed", exc_info=True)
            raise
        finally:
            if response_finished:
                self._log(scope, request_id, status_code, started_at, logging.INFO, "request_completed")
                from app.services.request_traffic_service import RequestTrafficService

                await RequestTrafficService.record(
                    scope.get("path", ""),
                    status_code,
                    (perf_counter() - started_at) * 1000,
                )
            request_id_context.reset(token)

    def _log(self, scope, request_id: str, status_code: int, started_at: float, level: int, message: str, **kwargs) -> None:
        client = scope.get("client")
        self.logger.log(
            level,
            message,
            extra={
                "request_id": request_id,
                "method": scope.get("method"),
                "path": scope.get("path"),
                "status_code": status_code,
                "duration_ms": round((perf_counter() - started_at) * 1000, 2),
                "client_ip": client[0] if client else "unknown",
            },
            **kwargs,
        )
