"""Centralized logging configuration for the SAE pipeline.

Usage:
    from core.logging_setup import setup_logging, get_logger
    setup_logging()  # call once at CLI entry point
    logger = get_logger(__name__)  # call in each module
"""

import logging
import os
from datetime import datetime

# _DEFAULT_LOG_DIR = "/workspace/SAE/logs"
_DEFAULT_LOG_DIR = "/shared/ssd_14T_data/kai/zhiyan/SAE/sae_pipeline/logs"
def setup_logging(
    level: str = "INFO",
    log_file: str | None = None,
    log_dir: str = _DEFAULT_LOG_DIR,
    suffix: str | None = None,
):
    """Configure logging to both console and file.

    Args:
        level: Logging level (DEBUG, INFO, WARNING, ERROR).
        log_file: Explicit log file path. If None, auto-generates
                  based on timestamp in log_dir.
        log_dir: Directory for auto-generated log files.
        suffix: Optional suffix for the log filename (e.g., "rank0", "main").
                Used to create per-process log files in multi-GPU runs.
    """
    os.makedirs(log_dir, exist_ok=True)

    if log_file is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        tag = f"_{suffix}" if suffix else ""
        log_file = os.path.join(log_dir, f"sae_pipeline_{timestamp}{tag}.log")

    fmt = "%(asctime)s | %(levelname)-7s | %(name)s | %(message)s"
    datefmt = "%Y-%m-%d %H:%M:%S"

    # In spawned worker processes, basicConfig may have already been called
    # by the parent. Clear existing handlers to ensure clean per-process setup.
    root = logging.getLogger()
    for h in root.handlers[:]:
        root.removeHandler(h)

    logging.basicConfig(
        level=getattr(logging, level.upper()),
        format=fmt,
        datefmt=datefmt,
        handlers=[
            logging.StreamHandler(),
            logging.FileHandler(log_file),
        ],
    )

    logging.getLogger().info(f"Logging to {log_file}")


def get_logger(name: str) -> logging.Logger:
    """Get a logger with the given name."""
    return logging.getLogger(name)