from __future__ import annotations

import contextvars
import json
import logging
import sys
import time


request_id_var: contextvars.ContextVar[str] = contextvars.ContextVar(
    "request_id", default="-"
)


class _JsonFormatter(logging.Formatter):
    def format(self, record: logging.LogRecord) -> str:
        ts = time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime(record.created))
        data: dict = {
            "ts": f"{ts}.{int(record.msecs):03d}Z",
            "level": record.levelname,
            "logger": record.name,
            "msg": record.getMessage(),
            "request_id": request_id_var.get(),
        }
        if record.exc_info:
            data["exc"] = self.formatException(record.exc_info)
        for key, val in record.__dict__.items():
            if key.startswith("ctx_"):
                data[key[4:]] = val
        return json.dumps(data, ensure_ascii=False)


def configure_logging(level: str = "INFO") -> None:
    level = (level or "INFO").upper()
    handler = logging.StreamHandler(sys.stdout)
    handler.setFormatter(_JsonFormatter())

    root = logging.getLogger()
    root.handlers.clear()
    root.addHandler(handler)
    root.setLevel(level)

    # Align uvicorn access/error logs with our JSON formatter.
    for name in ("uvicorn", "uvicorn.error", "uvicorn.access", "fastapi"):
        lg = logging.getLogger(name)
        lg.handlers.clear()
        lg.propagate = True
        lg.setLevel(level)

    # Trim noisy libraries.
    logging.getLogger("websockets").setLevel("WARNING")
    logging.getLogger("websockets.client").setLevel("WARNING")


def get_logger(name: str = "lingma_gateway") -> logging.Logger:
    return logging.getLogger(name)