import logging
import sys
import platform
import os
from pathlib import Path

class SafeFormatter(logging.Formatter):
    """Custom formatter that safely handles non-ASCII characters"""
    
    def format(self, record):
        # Handle message string safely
        if hasattr(record, 'msg') and record.msg is not None:
            if isinstance(record.msg, str):
                # Strip non-ASCII characters to prevent encoding errors
                record.msg = record.msg.encode("ascii", "ignore").decode()
            elif isinstance(record.msg, bytes):
                # Handle bytes objects
                record.msg = record.msg.decode("utf-8", "ignore")
        
        # Handle args tuple if present
        if hasattr(record, 'args') and record.args:
            safe_args = []
            for arg in record.args:
                if isinstance(arg, str):
                    safe_args.append(arg.encode("ascii", "ignore").decode())
                elif isinstance(arg, bytes):
                    safe_args.append(arg.decode("utf-8", "ignore"))
                else:
                    safe_args.append(str(arg))
            record.args = tuple(safe_args)
        
        # Handle exc_text if present
        if hasattr(record, 'exc_text') and record.exc_text:
            if isinstance(record.exc_text, str):
                record.exc_text = record.exc_text.encode("ascii", "ignore").decode()
        
        return super().format(record)

class SafeStreamHandler(logging.StreamHandler):
    """Stream handler that forces UTF-8 encoding where supported"""
    def __init__(self, stream=None):
        super().__init__(stream)
        if platform.system() == "Windows":
            # Force UTF-8 encoding on Windows
            if hasattr(stream, 'reconfigure'):
                try:
                    stream.reconfigure(encoding='utf-8', errors='ignore')
                except:
                    pass

def setup_safe_logging(log_level=logging.INFO, log_file='logs/safe_logging.log'):
    """Setup logging with SafeFormatter and UTF-8 encoding with enhanced persistence
    
    Args:
        log_level: Logging level (default: INFO)
        log_file: Path to log file (default: logs/safe_logging.log)
    """
    # Ensure logs directory exists
    log_path = Path(log_file)
    log_path.parent.mkdir(parents=True, exist_ok=True)
    
    # Clear existing handlers to avoid duplicates
    root_logger = logging.getLogger()
    for handler in root_logger.handlers[:]:
        root_logger.removeHandler(handler)
    
    # Create handlers with proper encoding
    handlers = []
    
    # Console handler with safe UTF-8 handling
    console_handler = SafeStreamHandler(sys.stdout)
    console_handler.setFormatter(SafeFormatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    ))
    handlers.append(console_handler)
    
    # File handler with UTF-8 encoding and error handling - ENHANCED for persistence
    try:
        encoding_kwargs = {
            "encoding": "utf-8",
            "errors": "ignore" if platform.system() == "Windows" else "backslashreplace"
        }
        
        # Use rotating file handler to prevent huge log files
        from logging.handlers import RotatingFileHandler
        file_handler = RotatingFileHandler(
            log_file, 
            maxBytes=10*1024*1024,  # 10MB max file size
            backupCount=5,  # Keep 5 backup files
            **encoding_kwargs
        )
        file_handler.setFormatter(SafeFormatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        ))
        
        # Force immediate flush for critical logs
        class FlushingHandler(RotatingFileHandler):
            def emit(self, record):
                super().emit(record)
                self.flush()  # Force flush after each log
        
        # Replace with flushing handler for critical systems
        file_handler = FlushingHandler(
            log_file,
            maxBytes=10*1024*1024,
            backupCount=5,
            **encoding_kwargs
        )
        file_handler.setFormatter(SafeFormatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        ))
        
        handlers.append(file_handler)
    except (OSError, IOError) as e:
        # If file handler fails, just use console handler
        print(f"Warning: Could not create log file {log_file}: {e}", file=sys.stderr)
    
    # Configure root logger
    logging.basicConfig(
        level=log_level,
        handlers=handlers,
        force=True  # Force reconfiguration
    )
    
    # Apply SafeFormatter to all existing loggers
    safe_formatter = SafeFormatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    for logger_name in logging.Logger.manager.loggerDict:
        logger = logging.getLogger(logger_name)
        for handler in logger.handlers:
            handler.setFormatter(safe_formatter)
    
    # Set up signal handlers for graceful shutdown and log flushing
    import signal
    import atexit
    
    def flush_all_logs():
        """Flush all log handlers"""
        for handler in logging.getLogger().handlers:
            if hasattr(handler, 'flush'):
                handler.flush()
        # Force logging shutdown
        logging.shutdown()
    
    def signal_handler(signum, frame):
        """Handle shutdown signals"""
        print(f"Received signal {signum}, flushing logs...")
        flush_all_logs()
        sys.exit(0)
    
    # Register signal handlers (Windows compatible)
    if platform.system() == "Windows":
        signal.signal(signal.SIGTERM, signal_handler)
        signal.signal(signal.SIGINT, signal_handler)
    else:
        signal.signal(signal.SIGTERM, signal_handler)
        signal.signal(signal.SIGINT, signal_handler)
        signal.signal(signal.SIGHUP, signal_handler)
    
    # Register atexit handler for normal shutdown
    atexit.register(flush_all_logs)

def setup_training_logger(log_level=logging.INFO, log_file='logs/training.log'):
    """Setup a separate training logger that writes to training.log
    
    Args:
        log_level: Logging level (default: INFO)
        log_file: Path to training log file (default: logs/training.log)
    
    Returns:
        logging.Logger: The training logger instance
    """
    # Ensure logs directory exists
    log_path = Path(log_file)
    log_path.parent.mkdir(parents=True, exist_ok=True)
    
    # Create training logger
    training_logger = logging.getLogger('training')
    training_logger.setLevel(log_level)
    
    # Clear existing handlers to avoid duplicates
    for handler in training_logger.handlers[:]:
        training_logger.removeHandler(handler)
    
    # Create file handler for training logs
    try:
        encoding_kwargs = {
            "encoding": "utf-8",
            "errors": "ignore" if platform.system() == "Windows" else "backslashreplace"
        }
        
        from logging.handlers import RotatingFileHandler
        file_handler = RotatingFileHandler(
            log_file,
            maxBytes=10*1024*1024,  # 10MB max file size
            backupCount=5,  # Keep 5 backup files
            **encoding_kwargs
        )
        file_handler.setFormatter(SafeFormatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        ))
        
        # Force immediate flush for training logs
        class FlushingHandler(RotatingFileHandler):
            def emit(self, record):
                super().emit(record)
                self.flush()  # Force flush after each log
        
        file_handler = FlushingHandler(
            log_file,
            maxBytes=10*1024*1024,
            backupCount=5,
            **encoding_kwargs
        )
        file_handler.setFormatter(SafeFormatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        ))
        
        training_logger.addHandler(file_handler)
        
    except (OSError, IOError) as e:
        print(f"Warning: Could not create training log file {log_file}: {e}", file=sys.stderr)
    
    # Prevent propagation to root logger to avoid duplicate logs
    training_logger.propagate = False
    
    return training_logger