import logging as _logging


class _RankFilter(_logging.Filter):
    def __init__(self, active_rank: int):
        super().__init__()
        self.active_rank = active_rank

    def filter(self, record):
        return getattr(record, "rank", 0) == self.active_rank


_LOGGER = _logging.getLogger("torchlitex")
_WANDB_RUN = None


def configure_logging(level=_logging.INFO, fmt="%(asctime)s | %(levelname)s | %(message)s", active_rank: int = 0):
    if not _LOGGER.handlers:
        handler = _logging.StreamHandler()
        handler.setFormatter(_logging.Formatter(fmt))
        handler.addFilter(_RankFilter(active_rank))
        _LOGGER.addHandler(handler)
    _LOGGER.setLevel(level)
    _LOGGER.propagate = False
    return _LOGGER


def get_logger():
    return _LOGGER


def rank_print(message, rank: int = 0, active_rank: int = 0, level: str = "info"):
    if rank != active_rank:
        return

    if _LOGGER.handlers:
        log_fn = getattr(_LOGGER, level, _LOGGER.info)
        log_fn(message, extra={"rank": rank})
    else:
        print(message)


def init_wandb(project=None, **kwargs):
    """
    Lazy wandb init; safe to call when wandb is unavailable.
    """
    global _WANDB_RUN
    if _WANDB_RUN is not None:
        return _WANDB_RUN

    try:
        import wandb  # type: ignore
    except ImportError:
        return None

    _WANDB_RUN = wandb.init(project=project, **kwargs)
    return _WANDB_RUN


def wandb_log(data: dict, step: int | None = None, rank: int = 0, active_rank: int = 0, commit: bool | None = None):
    """
    Rank-aware wandb logging; no-op if wandb is unavailable or rank mismatch.
    """
    if rank != active_rank:
        return
    if _WANDB_RUN is None:
        return
    _WANDB_RUN.log(data, step=step, commit=commit)
