import logging
import sys
from contextvars import ContextVar

# Храним trace_id текущего потока/корутины.
TRACE_ID: ContextVar[str] = ContextVar("trace_id", default="-")


class TraceIdFilter(logging.Filter):
    # Добавляет trace_id в каждую запись лога для сквозной трассировки.
    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:
    # stderr — Docker и Celery prefork корректно его подхватывают.
    handlers: list[logging.Handler] = [logging.StreamHandler(sys.stderr)]
    if log_file:
        handlers.append(logging.FileHandler(log_file, encoding="utf-8"))
    trace_filter = TraceIdFilter()
    for handler in handlers:
        handler.addFilter(trace_filter)
        handler.setLevel(getattr(logging, level.upper(), logging.INFO))
    root = logging.getLogger()
    root.setLevel(getattr(logging, level.upper(), logging.INFO))
    # Убираем старые хендлеры, чтобы не дублировать после fork.
    for old_handler in list(root.handlers):
        try:
            old_handler.close()
        except Exception:
            pass
    root.handlers.clear()
    for handler in handlers:
        root.addHandler(handler)
    root.propagate = False
    fmt = logging.Formatter(
        "%(asctime)s | %(levelname)s | %(name)s | %(message)s"
    )
    for handler in root.handlers:
        handler.setFormatter(fmt)

    for logger_name in (
        "celery.app.trace",
        "celery.worker.request",
        "celery.worker.strategy",
    ):
        noisy_logger = logging.getLogger(logger_name)
        noisy_logger.handlers.clear()
        noisy_logger.propagate = False
        noisy_logger.disabled = True