"""
Logger Setup Utility.

Ref: revisions.txt Section 1
"""

import logging
import logging.handlers
import sys
import os
from typing import Dict, Any

# Conditional import for colorlog
try:
    import colorlog
    COLORLOG_AVAILABLE = True
except ImportError:
    COLORLOG_AVAILABLE = False

# Assuming IOManager is in the same directory or accessible via path
try:
    from .io_manager import IOManager
except ImportError:
    # Fallback if run as script or structure changes
    IOManager = None 

# Define default log format strings (as fallbacks)
DEFAULT_LOG_FORMAT_CONSOLE = '%(log_color)s%(levelname)-8s%(reset)s | %(name)-12s | %(message)s'
DEFAULT_LOG_FORMAT_FILE = '%(asctime)s | %(levelname)-8s | %(name)-15s | %(filename)s:%(lineno)d | %(message)s'
DEFAULT_LOG_DATE_FORMAT = '%Y-%m-%d %H:%M:%S'

def setup_logger(cfg: Dict[str, Any], run_id: str, io: IOManager) -> logging.Logger:
    """
    Configures the root logger with console and rotating file handlers based on config.

    Args:
        cfg (Dict[str, Any]): The pipeline configuration dictionary.
        run_id (str): The unique run identifier.
        io (IOManager): The IOManager instance for getting log file path.

    Returns:
        logging.Logger: The configured root logger instance.
    """
    # --- Get Logging Config --- #
    log_cfg = cfg.get('logging', {})
    output_cfg = cfg.get('output', {}) # Still need this for base log level

    console_log_level_str = log_cfg.get('console_level', 'INFO').upper()
    file_log_level_str = log_cfg.get('file_level', 'DEBUG').upper()
    log_to_file = log_cfg.get('log_to_file', True)
    log_format_console = log_cfg.get('log_format_console', DEFAULT_LOG_FORMAT_CONSOLE)
    log_format_file = log_cfg.get('log_format', DEFAULT_LOG_FORMAT_FILE) # Use 'log_format' for file
    log_date_format = log_cfg.get('log_date_format', DEFAULT_LOG_DATE_FORMAT)
    log_file_max_bytes = log_cfg.get('log_file_max_bytes', 10*1024*1024) # Default 10MB
    log_file_backup_count = log_cfg.get('log_file_backup_count', 5) # Default 5 backups
    # --- End Logging Config --- #

    console_log_level = getattr(logging, console_log_level_str, logging.INFO)
    file_log_level = getattr(logging, file_log_level_str, logging.DEBUG)

    root_logger = logging.getLogger() # Get the root logger
    # Set root level to the lowest of the handler levels to ensure messages pass through
    root_logger.setLevel(min(console_log_level, file_log_level))

    # Remove existing handlers to avoid duplication if called multiple times
    for handler in root_logger.handlers[:]:
        root_logger.removeHandler(handler)

    # --- Console Handler --- #
    if COLORLOG_AVAILABLE:
        cformat = colorlog.ColoredFormatter(
            log_format_console, # Use format from config
            datefmt=log_date_format, # Use date format from config
            reset=True,
            log_colors={
                'DEBUG':    'cyan',
                'INFO':     'green',
                'WARNING':  'yellow',
                'ERROR':    'red',
                'CRITICAL': 'red,bg_white',
            },
            secondary_log_colors={},
            style='%'
        )
        console_handler = colorlog.StreamHandler(sys.stdout)
        console_handler.setFormatter(cformat)
    else:
        # Use file format from config if colorlog not available (or user specified non-color format)
        formatter = logging.Formatter(log_format_file, datefmt=log_date_format)
        console_handler = logging.StreamHandler(sys.stdout)
        console_handler.setFormatter(formatter)
        
    console_handler.setLevel(console_log_level) # Console respects the config level
    root_logger.addHandler(console_handler)

    # --- Rotating File Handler (if enabled) --- #
    if log_to_file and io is not None:
        try:
            # Use run_id in log filename
            log_file_path = io.path('logs', f'pipeline_{run_id}', suffix='.log')
            # Use RotatingFileHandler with params from config
            file_handler = logging.handlers.RotatingFileHandler(
                log_file_path, 
                maxBytes=log_file_max_bytes, 
                backupCount=log_file_backup_count, 
                encoding='utf-8'
            )
            file_formatter = logging.Formatter(log_format_file, datefmt=log_date_format)
            file_handler.setFormatter(file_formatter)
            file_handler.setLevel(file_log_level) # File respects its config level
            root_logger.addHandler(file_handler)
            logging.info(f"File logging ({file_log_level_str} level) configured at: {log_file_path}")
        except Exception as e:
             logging.error(f"Failed to configure file logging: {e}", exc_info=True)
    elif log_to_file and io is None:
        logging.warning("File logging enabled in config, but IOManager not provided. Cannot configure file handler.")
    else:
         logging.info("File logging disabled in config.")

    # --- Set TensorFlow Log Level --- #
    # Quieter TF logs by default
    tf_log_level = os.environ.get('TF_CPP_MIN_LOG_LEVEL', '2') # Default to ERROR
    os.environ['TF_CPP_MIN_LOG_LEVEL'] = tf_log_level 
    # Also set Python TF logger level (optional, affects tf.get_logger())
    tf_logger = logging.getLogger('tensorflow')
    if tf_log_level == '0': # ALL
        tf_logger.setLevel(logging.DEBUG)
    elif tf_log_level == '1': # INFO
         tf_logger.setLevel(logging.INFO)
    elif tf_log_level == '2': # WARNING
         tf_logger.setLevel(logging.WARNING)
    else: # 3 = ERROR
         tf_logger.setLevel(logging.ERROR)
    logging.info(f"TensorFlow logging level set based on TF_CPP_MIN_LOG_LEVEL={tf_log_level}")

    logging.info(f"Root logger configured. Console level: {console_log_level_str}, File level: {file_log_level_str}")
    return root_logger

# Example usage:
if __name__ == '__main__':
    mock_config = {
        'base_dirs': {'results': 'temp_results', 'models': 'temp_models', 'logs': 'temp_logs'},
        'output': {'log_level': 'INFO'},
        'logging': {
            'console_level': 'INFO',
            'file_level': 'DEBUG',
            'log_to_file': True,
            'log_format_console': '%(log_color)s%(levelname)-8s%(reset)s | %(name)-12s | %(message)s',
            'log_format': '%(asctime)s | %(levelname)-8s | %(name)-15s | %(filename)s:%(lineno)d | %(message)s',
            'log_date_format': '%Y-%m-%d %H:%M:%S',
            'log_file_max_bytes': 10*1024*1024,
            'log_file_backup_count': 5
        }
    }
    mock_run_id = "20250418_113000_logtest"
    
    # Need a mock IOManager for the example
    if IOManager is None:
         print("Mocking IOManager as it couldn't be imported.")
         class MockIOManager:
             def __init__(self, cfg, run_id):
                 self.run_logs_dir = os.path.join('temp_logs', run_id)
                 os.makedirs(self.run_logs_dir, exist_ok=True)
             def path(self, section, name, suffix):
                 return os.path.join(self.run_logs_dir, f"{name}{suffix}")
         io_manager = MockIOManager(mock_config, mock_run_id)
    else:
        if not os.path.exists('temp_logs'): os.makedirs('temp_logs')
        io_manager = IOManager(mock_config, mock_run_id)

    logger_instance = setup_logger(mock_config, mock_run_id, io_manager)
    
    logger_instance.debug("This is a debug message (should only go to file).")
    logger_instance.info("This is an info message.")
    logger_instance.warning("This is a warning message.")
    logger_instance.error("This is an error message.")

    print(f"Check log file in: {io_manager.run_logs_dir}")

    # Clean up
    import shutil
    if os.path.exists('temp_logs'): shutil.rmtree('temp_logs') 