import logging
import sys
from contextvars import ContextVar


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


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


def set_trace_id(trace_id: str) -> None:
    TRACE_ID.set(trace_id)


def setup_logging(level: str = "INFO", log_file: str | None = None) -> None:
    handlers: list[logging.Handler] = [logging.StreamHandler(sys.stdout)]
    if log_file:
        handlers.append(logging.FileHandler(log_file, encoding="utf-8"))
    trace_filter = TraceIdFilter()
    for handler in handlers:
        handler.addFilter(trace_filter)
    logging.basicConfig(
        level=getattr(logging, level.upper(), logging.INFO),
        format="%(asctime)s | %(levelname)s | %(name)s | trace=%(trace_id)s | %(message)s",
        handlers=handlers,
        force=True,
    )