import logging
import sys
import os
from warnings import warn
from mpi4py import MPI

logger = None  # Global logger instance

# Custom log levels mapping
CUSTOM_LOG_LEVELS = {
    1: logging.INFO,   # Basic workflow info
    2: logging.INFO,   # Extended workflow info (1 tab)
    3: logging.INFO,   # Deep workflow info (2 tabs)
    5: logging.DEBUG,  # Debug info
    7: logging.WARNING,
    9: logging.CRITICAL,
}

def configure_logger(log_path: str, log_detail: int = 1) -> None:
    """
    Configures a unique logger for each MPI process.
    Ensures the log folder exists before creating log files.

    Parameters
    ----------
    log_path : str
        Path to the log files. Each process will have its own log file.
        <rank> will be replaced by the corresponding rank with 3 digits
    log_detail : int, optional
        Level of log detail (1 = basic, 2 = extended, 3 = deep, 5 = debug). Default is 1.
    """
    global logger
    if logger is not None:
        return  # Skip reconfiguration if logger is already set

    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()

    # Ensure the log directory exists
    log_dir = os.path.dirname(log_path)
    if not os.path.exists(log_dir):
        if rank == 0:  # Only rank 0 creates the directory
            os.makedirs(log_dir, exist_ok=True)
            print(f"Log directory created: {log_dir}")
    comm.Barrier()  # Synchronize all ranks after creating the directory

    # Map custom log detail to logging levels
    log_level = CUSTOM_LOG_LEVELS.get(log_detail, logging.INFO)

    # Create a logger instance specific to the process rank
    logger = logging.getLogger(f"process_{rank}")
    logger.setLevel(log_level)

    # File handler for logging to a file
    log_path = log_path.replace('<rank>', str(rank).zfill(3))
    file_handler = logging.FileHandler(log_path, mode='w')
    file_formatter = logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
    file_handler.setFormatter(file_formatter)
    logger.addHandler(file_handler)

    # Add a stream handler for rank 0 to log to stdout
    if rank == 0:
        stream_handler = logging.StreamHandler(sys.stdout)
        stream_formatter = logging.Formatter("%(message)s")  # Simple format for stdout
        stream_handler.setFormatter(stream_formatter)
        stream_handler.flush = sys.stdout.flush  # Ensure flushing happens after every log
        logger.addHandler(stream_handler)

    logger.info(f"Logger configured for process {rank}. Log level: {logging.getLevelName(log_level)}")
    if rank == 0:
        logger.info("Process 0 will also log messages to stdout.")

def log_message(message: str, level: int = 1) -> None:
    """
    Logs a message at the specified custom log level, with tabulation depending on the level.

    Parameters
    ----------
    level : int
        The custom log level (1 = basic, 2 = extended, 3 = deep, 5 = debug, 7 = warning, 9 = critical).
    message : str
        The log message to be recorded.
    """
    if logger is None:
        raise RuntimeError("Logger is not configured. Call `configure_logger` first.")

    # Map custom level to Python's logging levels
    python_log_level = CUSTOM_LOG_LEVELS.get(level, logging.INFO)

    # Determine tabulation prefix based on the level of the message
    tab_prefix = ""
    if level == 2:
        tab_prefix = "\t"
    elif level == 3:
        tab_prefix = "\t\t"

    # Log messages at levels 7 and 9 independently of logger configuration
    formatted_message = f"{tab_prefix}{message}"
    if level == 7:  # WARNING level
        logger.warning(formatted_message)
        warn(formatted_message)
    elif level == 9:  # CRITICAL level
        sys.stderr.write(formatted_message)
        logger.critical(formatted_message)
    else:
        # For other levels, log with standard behavior
        if python_log_level == logging.INFO:
            logger.info(formatted_message)
        elif python_log_level == logging.DEBUG:
            logger.debug(formatted_message)
        else:
            logger.log(python_log_level, formatted_message)
    return None
