"""结构化 JSON 日志 + trace_id。"""
import logging
import sys
import uuid
from contextvars import ContextVar

from pythonjsonlogger import jsonlogger

trace_id_var: ContextVar[str] = ContextVar("trace_id", default="")


def get_trace_id() -> str:
    t = trace_id_var.get()
    if not t:
        t = str(uuid.uuid4())
        trace_id_var.set(t)
    return t


def set_trace_id(tid: str) -> None:
    trace_id_var.set(tid)


class TraceIdFilter(logging.Filter):
    def filter(self, record: logging.LogRecord) -> bool:
        record.trace_id = get_trace_id()
        return True


class JsonFormatter(jsonlogger.JsonFormatter):
    def add_fields(self, log_record: dict, record: logging.LogRecord, message_dict: dict) -> None:
        super().add_fields(log_record, record, message_dict)
        log_record["trace_id"] = getattr(record, "trace_id", "")
        log_record["level"] = record.levelname
        if record.exc_info:
            log_record["exception"] = self.formatException(record.exc_info)


def setup_logging(log_level: str = "INFO", log_json: bool = True) -> None:
    root = logging.getLogger()
    root.handlers.clear()
    handler = logging.StreamHandler(sys.stdout)
    if log_json:
        handler.setFormatter(JsonFormatter("%(timestamp)s %(level)s %(message)s %(trace_id)s"))
    else:
        handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s [%(trace_id)s] %(message)s"))
    handler.addFilter(TraceIdFilter())
    root.addHandler(handler)
    root.setLevel(getattr(logging, log_level.upper(), logging.INFO))