import logging
import sys
from enum import StrEnum
from pathlib import Path
from logging.handlers import RotatingFileHandler, TimedRotatingFileHandler
from datetime import datetime
import json


class LogLevel(StrEnum):
    DEBUG = "DEBUG"
    INFO = "INFO"
    WARNING = "WARNING"
    ERROR = "ERROR"
    CRITICAL = "CRITICAL"


class CustomFormatter(logging.Formatter):
    """Custom formatter with colors and detailed information"""
    
    # ANSI color codes
    COLORS = {
        'DEBUG': '\033[36m',      # Cyan
        'INFO': '\033[32m',       # Green
        'WARNING': '\033[33m',    # Yellow
        'ERROR': '\033[31m',      # Red
        'CRITICAL': '\033[35m',   # Magenta
        'RESET': '\033[0m'        # Reset
    }
    
    def format(self, record):
        # Add color to level name
        if hasattr(sys.stdout, 'isatty') and sys.stdout.isatty():
            levelname = record.levelname
            if levelname in self.COLORS:
                record.levelname = f"{self.COLORS[levelname]}{levelname}{self.COLORS['RESET']}"
        
        # Add extra context
        record.service = getattr(record, 'service', 'unknown')
        record.trace_id = getattr(record, 'trace_id', 'N/A')
        record.span_id = getattr(record, 'span_id', 'N/A')
        
        return super().format(record)


class JSONFormatter(logging.Formatter):
    """JSON formatter for structured logging"""
    
    def format(self, record):
        log_data = {
            'timestamp': datetime.utcnow().isoformat(),
            'level': record.levelname,
            'logger': record.name,
            'message': record.getMessage(),
            'module': record.module,
            'function': record.funcName,
            'line': record.lineno,
            'process': record.process,
            'thread': record.thread,
            'service': getattr(record, 'service', 'unknown'),
            'trace_id': getattr(record, 'trace_id', None),
            'span_id': getattr(record, 'span_id', None),
        }
        
        # Add exception info if present
        if record.exc_info:
            log_data['exception'] = self.formatException(record.exc_info)
        
        # Add extra fields
        if hasattr(record, 'extra_data'):
            log_data['extra'] = record.extra_data  # type: ignore
        
        return json.dumps(log_data)


def setup_logging(
    service_name: str = "devserver",
    log_level: str = "INFO",
    log_dir: str = "logs",
    enable_console: bool = True,
    enable_file: bool = True,
    enable_json: bool = True
) -> logging.Logger:
    """
    Setup comprehensive logging configuration
    
    Args:
        service_name: Name of the service for logging context
        log_level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
        log_dir: Directory to store log files
        enable_console: Enable console output
        enable_file: Enable file output
        enable_json: Enable JSON structured logging
    
    Returns:
        Configured logger instance
    """
    # Create logs directory
    log_path = Path(log_dir)
    log_path.mkdir(exist_ok=True)
    
    # Get root logger
    logger = logging.getLogger(service_name)
    logger.setLevel(getattr(logging, log_level.upper()))
    logger.handlers.clear()
    
    # Console Handler with colors
    if enable_console:
        console_handler = logging.StreamHandler(sys.stdout)
        console_handler.setLevel(logging.DEBUG)
        console_format = CustomFormatter(
            fmt='%(asctime)s | %(levelname)-8s | %(service)-12s | %(name)-20s | %(funcName)-25s:%(lineno)-4d | %(message)s',
            datefmt='%Y-%m-%d %H:%M:%S'
        )
        console_handler.setFormatter(console_format)
        logger.addHandler(console_handler)
    
    # File Handler (Rotating) - Text format
    if enable_file:
        file_handler = RotatingFileHandler(
            filename=log_path / f"{service_name}.log",
            maxBytes=10 * 1024 * 1024,  # 10MB
            backupCount=5,
            encoding='utf-8'
        )
        file_handler.setLevel(logging.DEBUG)
        file_format = logging.Formatter(
            fmt='%(asctime)s | %(levelname)-8s | %(service)-12s | %(name)-20s | %(module)-15s | %(funcName)-25s:%(lineno)-4d | TraceID:%(trace_id)s | %(message)s',
            datefmt='%Y-%m-%d %H:%M:%S'
        )
        file_handler.setFormatter(file_format)
        logger.addHandler(file_handler)
    
    # JSON File Handler (Daily rotation) - Structured logging
    if enable_json:
        json_handler = TimedRotatingFileHandler(
            filename=log_path / f"{service_name}_json.log",
            when='midnight',
            interval=1,
            backupCount=30,
            encoding='utf-8'
        )
        json_handler.setLevel(logging.INFO)
        json_handler.setFormatter(JSONFormatter())
        logger.addHandler(json_handler)
    
    # Error Handler (Separate file for errors)
    error_handler = RotatingFileHandler(
        filename=log_path / f"{service_name}_errors.log",
        maxBytes=5 * 1024 * 1024,  # 5MB
        backupCount=5,
        encoding='utf-8'
    )
    error_handler.setLevel(logging.ERROR)
    error_format = logging.Formatter(
        fmt='%(asctime)s | %(levelname)-8s | %(service)-12s | %(name)-20s | %(module)-15s | %(funcName)-25s:%(lineno)-4d | TraceID:%(trace_id)s\nMessage: %(message)s\n%(pathname)s\n',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    error_handler.setFormatter(error_format)
    logger.addHandler(error_handler)
    
    # Set default service name
    old_factory = logging.getLogRecordFactory()
    
    def record_factory(*args, **kwargs):
        record = old_factory(*args, **kwargs)
        record.service = service_name
        record.trace_id = 'N/A'
        record.span_id = 'N/A'
        return record
    
    logging.setLogRecordFactory(record_factory)
    
    logger.info(f"Logging initialized for service: {service_name}")
    logger.info(f"Log level: {log_level}")
    logger.info(f"Log directory: {log_path.absolute()}")
    
    return logger


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


def log_with_trace(logger: logging.Logger, level: str, message: str, trace_id: str | None = None, span_id: str | None = None, **kwargs):
    """
    Log a message with OpenTelemetry trace context
    
    Args:
        logger: Logger instance
        level: Log level (debug, info, warning, error, critical)
        message: Log message
        trace_id: OpenTelemetry trace ID
        span_id: OpenTelemetry span ID
        **kwargs: Additional context data
    """
    extra = {
        'trace_id': trace_id or 'N/A',
        'span_id': span_id or 'N/A',
        'extra_data': kwargs
    }
    
    log_func = getattr(logger, level.lower())
    log_func(message, extra=extra)