"""
Logging utilities for the COBY system.
"""

import logging
import logging.handlers
import sys
import uuid
from pathlib import Path
from typing import Optional
from contextvars import ContextVar

# Context variable for correlation ID
correlation_id: ContextVar[Optional[str]] = ContextVar('correlation_id', default=None)


class CorrelationFilter(logging.Filter):
    """Add correlation ID to log records"""
    
    def filter(self, record):
        record.correlation_id = correlation_id.get() or 'N/A'
        return True


class COBYFormatter(logging.Formatter):
    """Custom formatter with correlation ID support"""
    
    def __init__(self, include_correlation_id: bool = True):
        self.include_correlation_id = include_correlation_id
        
        if include_correlation_id:
            fmt = '%(asctime)s - %(name)s - %(levelname)s - [%(correlation_id)s] - %(message)s'
        else:
            fmt = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            
        super().__init__(fmt, datefmt='%Y-%m-%d %H:%M:%S')


def setup_logging(
    level: str = 'INFO',
    log_file: Optional[str] = None,
    max_file_size: int = 100,  # MB
    backup_count: int = 5,
    enable_correlation_id: bool = True,
    console_output: bool = True
) -> None:
    """
    Set up logging configuration for the COBY system.
    
    Args:
        level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
        log_file: Path to log file (None = no file logging)
        max_file_size: Maximum log file size in MB
        backup_count: Number of backup files to keep
        enable_correlation_id: Whether to include correlation IDs in logs
        console_output: Whether to output logs to console
    """
    # Convert string level to logging constant
    numeric_level = getattr(logging, level.upper(), logging.INFO)
    
    # Create root logger
    root_logger = logging.getLogger()
    root_logger.setLevel(numeric_level)
    
    # Clear existing handlers
    root_logger.handlers.clear()
    
    # Create formatter
    formatter = COBYFormatter(include_correlation_id=enable_correlation_id)
    
    # Add correlation filter if enabled
    correlation_filter = CorrelationFilter() if enable_correlation_id else None
    
    # Console handler
    if console_output:
        console_handler = logging.StreamHandler(sys.stdout)
        console_handler.setLevel(numeric_level)
        console_handler.setFormatter(formatter)
        if correlation_filter:
            console_handler.addFilter(correlation_filter)
        root_logger.addHandler(console_handler)
    
    # File handler
    if log_file:
        # Create log directory if it doesn't exist
        log_path = Path(log_file)
        log_path.parent.mkdir(parents=True, exist_ok=True)
        
        # Rotating file handler
        file_handler = logging.handlers.RotatingFileHandler(
            log_file,
            maxBytes=max_file_size * 1024 * 1024,  # Convert MB to bytes
            backupCount=backup_count
        )
        file_handler.setLevel(numeric_level)
        file_handler.setFormatter(formatter)
        if correlation_filter:
            file_handler.addFilter(correlation_filter)
        root_logger.addHandler(file_handler)
    
    # Set specific logger levels
    logging.getLogger('websockets').setLevel(logging.WARNING)
    logging.getLogger('urllib3').setLevel(logging.WARNING)
    logging.getLogger('requests').setLevel(logging.WARNING)


def get_logger(name: str) -> logging.Logger:
    """
    Get a logger instance with the specified name.
    
    Args:
        name: Logger name (typically __name__)
        
    Returns:
        logging.Logger: Logger instance
    """
    return logging.getLogger(name)


def set_correlation_id(corr_id: Optional[str] = None) -> str:
    """
    Set correlation ID for current context.
    
    Args:
        corr_id: Correlation ID (generates UUID if None)
        
    Returns:
        str: The correlation ID that was set
    """
    if corr_id is None:
        corr_id = str(uuid.uuid4())[:8]  # Short UUID
    
    correlation_id.set(corr_id)
    return corr_id


def get_correlation_id() -> Optional[str]:
    """
    Get current correlation ID.
    
    Returns:
        str: Current correlation ID or None
    """
    return correlation_id.get()


def clear_correlation_id() -> None:
    """Clear correlation ID from current context."""
    correlation_id.set(None)