"""
Logging utilities for Toxo platform.

This module provides structured logging capabilities with support for:
- Multiple log levels and formats
- Performance monitoring and metrics
- Structured JSON logging for production
- Development-friendly console logging
- Log aggregation and analysis
"""

import os
import sys
import json
import time
from pathlib import Path
from typing import Dict, Any, Optional, Union
from datetime import datetime
from functools import wraps
import logging
import logging.handlers
from contextlib import contextmanager

from loguru import logger as loguru_logger


class ToxoLogger:
    """
    Custom logger for Toxo platform with enhanced capabilities.
    """
    
    def __init__(
        self,
        name: str,
        level: str = "INFO",
        log_dir: Optional[Union[str, Path]] = None,
        structured: bool = False,
        enable_performance_logging: bool = True
    ):
        """
        Initialize Toxo logger.
        
        Args:
            name: Logger name
            level: Log level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
            log_dir: Directory for log files
            structured: Enable structured JSON logging
            enable_performance_logging: Enable performance metrics logging
        """
        self.name = name
        self.level = level.upper()
        self.log_dir = Path(log_dir) if log_dir else Path("logs")
        self.structured = structured
        self.enable_performance_logging = enable_performance_logging
        
        # Create log directory
        self.log_dir.mkdir(exist_ok=True)
        
        # Initialize loggers
        self._setup_loguru()
        self._setup_standard_logger()
        
        # Performance tracking
        self.performance_metrics = {}
        
    def _setup_loguru(self):
        """Setup loguru logger with custom formatting."""
        # Remove default handler
        loguru_logger.remove()
        
        # Console handler with colors for development
        if not self.structured:
            format_string = (
                "<green>{time:YYYY-MM-DD HH:mm:ss}</green> | "
                "<level>{level: <8}</level> | "
                "<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> | "
                "<level>{message}</level>"
            )
            loguru_logger.add(
                sys.stderr,
                format=format_string,
                level=self.level,
                colorize=True,
                backtrace=True,
                diagnose=True
            )
        else:
            # Structured JSON logging for production
            def json_formatter(record):
                return json.dumps({
                    "timestamp": record["time"].isoformat(),
                    "level": record["level"].name,
                    "logger": record["name"],
                    "module": record["module"],
                    "function": record["function"],
                    "line": record["line"],
                    "message": record["message"],
                    "extra": record.get("extra", {})
                })
            
            loguru_logger.add(
                sys.stdout,
                format=json_formatter,
                level=self.level,
                serialize=False
            )
        
        # File handlers
        main_log_file = self.log_dir / f"{self.name}.log"
        error_log_file = self.log_dir / f"{self.name}_error.log"
        
        # Main log file with rotation
        loguru_logger.add(
            main_log_file,
            rotation="10 MB",
            retention="7 days",
            compression="gz",
            level=self.level,
            format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} | {message}",
            backtrace=True,
            diagnose=True
        )
        
        # Error-only log file
        loguru_logger.add(
            error_log_file,
            rotation="5 MB",
            retention="30 days",
            compression="gz",
            level="ERROR",
            format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} | {message}",
            backtrace=True,
            diagnose=True
        )
        
        # Performance log file if enabled
        if self.enable_performance_logging:
            perf_log_file = self.log_dir / f"{self.name}_performance.log"
            loguru_logger.add(
                perf_log_file,
                rotation="10 MB",
                retention="3 days",
                level="DEBUG",
                filter=lambda record: "performance" in record.get("extra", {}),
                format="{time:YYYY-MM-DD HH:mm:ss} | PERF | {extra[performance]}",
            )
    
    def _setup_standard_logger(self):
        """Setup standard Python logger for compatibility."""
        self.logger = logging.getLogger(self.name)
        self.logger.setLevel(getattr(logging, self.level))
        
        # Prevent duplicate logs
        if not self.logger.handlers:
            handler = logging.StreamHandler()
            formatter = logging.Formatter(
                '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
            )
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)
    
    def debug(self, message: str, **kwargs):
        """Log debug message."""
        loguru_logger.bind(**kwargs).debug(message)
    
    def info(self, message: str, **kwargs):
        """Log info message."""
        loguru_logger.bind(**kwargs).info(message)
    
    def warning(self, message: str, **kwargs):
        """Log warning message."""
        loguru_logger.bind(**kwargs).warning(message)
    
    def error(self, message: str, **kwargs):
        """Log error message."""
        loguru_logger.bind(**kwargs).error(message)
    
    def critical(self, message: str, **kwargs):
        """Log critical message."""
        loguru_logger.bind(**kwargs).critical(message)
    
    def log_performance(self, operation: str, duration: float, **metrics):
        """
        Log performance metrics.
        
        Args:
            operation: Name of the operation
            duration: Operation duration in seconds
            **metrics: Additional performance metrics
        """
        if not self.enable_performance_logging:
            return
        
        perf_data = {
            "operation": operation,
            "duration": duration,
            "timestamp": datetime.now().isoformat(),
            **metrics
        }
        
        # Store in memory for aggregation
        if operation not in self.performance_metrics:
            self.performance_metrics[operation] = []
        self.performance_metrics[operation].append(perf_data)
        
        # Log to file
        loguru_logger.bind(performance=json.dumps(perf_data)).debug("Performance metric")
    
    def log_training_metrics(
        self,
        epoch: int,
        loss: float,
        accuracy: Optional[float] = None,
        **metrics
    ):
        """
        Log training-specific metrics.
        
        Args:
            epoch: Training epoch
            loss: Training loss
            accuracy: Training accuracy
            **metrics: Additional training metrics
        """
        training_data = {
            "type": "training",
            "epoch": epoch,
            "loss": loss,
            "accuracy": accuracy,
            "timestamp": datetime.now().isoformat(),
            **metrics
        }
        
        self.info(
            f"Training Epoch {epoch}: Loss={loss:.4f}, Accuracy={accuracy:.4f if accuracy else 'N/A'}",
            training_metrics=training_data
        )
    
    def log_api_call(
        self,
        endpoint: str,
        method: str,
        duration: float,
        status_code: int,
        tokens_used: Optional[int] = None
    ):
        """
        Log API call metrics.
        
        Args:
            endpoint: API endpoint
            method: HTTP method
            duration: Request duration
            status_code: HTTP status code
            tokens_used: Number of tokens used
        """
        api_data = {
            "type": "api_call",
            "endpoint": endpoint,
            "method": method,
            "duration": duration,
            "status_code": status_code,
            "tokens_used": tokens_used,
            "timestamp": datetime.now().isoformat()
        }
        
        level = "info" if 200 <= status_code < 300 else "warning"
        getattr(self, level)(
            f"API {method} {endpoint}: {status_code} ({duration:.3f}s)",
            api_metrics=api_data
        )
    
    def get_performance_summary(self, operation: Optional[str] = None) -> Dict[str, Any]:
        """
        Get performance metrics summary.
        
        Args:
            operation: Specific operation to summarize (None for all)
            
        Returns:
            Performance summary
        """
        if operation:
            metrics = self.performance_metrics.get(operation, [])
            operations = {operation: metrics}
        else:
            operations = self.performance_metrics
        
        summary = {}
        for op_name, op_metrics in operations.items():
            if not op_metrics:
                continue
            
            durations = [m["duration"] for m in op_metrics]
            summary[op_name] = {
                "count": len(durations),
                "avg_duration": sum(durations) / len(durations),
                "min_duration": min(durations),
                "max_duration": max(durations),
                "total_duration": sum(durations)
            }
        
        return summary


class PerformanceTimer:
    """Context manager for timing operations."""
    
    def __init__(self, logger: ToxoLogger, operation: str, **metrics):
        """
        Initialize performance timer.
        
        Args:
            logger: Toxo logger instance
            operation: Operation name
            **metrics: Additional metrics to log
        """
        self.logger = logger
        self.operation = operation
        self.metrics = metrics
        self.start_time = None
        self.end_time = None
    
    def __enter__(self):
        """Start timing."""
        self.start_time = time.time()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """End timing and log performance."""
        self.end_time = time.time()
        duration = self.end_time - self.start_time
        
        # Add exception info if any
        if exc_type:
            self.metrics["exception"] = str(exc_type.__name__)
            self.metrics["success"] = False
        else:
            self.metrics["success"] = True
        
        self.logger.log_performance(self.operation, duration, **self.metrics)


def performance_monitor(operation: str, logger: Optional[ToxoLogger] = None):
    """
    Decorator for monitoring function performance.
    
    Args:
        operation: Operation name
        logger: Logger instance (uses default if None)
    """
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            nonlocal logger
            if logger is None:
                logger = get_logger(func.__module__)
            
            with PerformanceTimer(logger, operation):
                return func(*args, **kwargs)
        return wrapper
    return decorator


# Global logger instances
_loggers: Dict[str, ToxoLogger] = {}


def get_logger(
    name: str,
    level: str = "INFO",
    log_dir: Optional[Union[str, Path]] = None,
    structured: bool = None,
    enable_performance_logging: bool = True
) -> ToxoLogger:
    """
    Get or create a logger instance.
    
    Args:
        name: Logger name
        level: Log level
        log_dir: Log directory
        structured: Enable structured logging (auto-detect if None)
        enable_performance_logging: Enable performance logging
        
    Returns:
        ToxoLogger instance
    """
    # Auto-detect structured logging based on environment
    if structured is None:
        structured = os.getenv("TOXO_STRUCTURED_LOGGING", "false").lower() == "true"
    
    # Use environment variables for defaults
    level = os.getenv("TOXO_LOG_LEVEL", level)
    if log_dir is None:
        log_dir = os.getenv("TOXO_LOG_DIR", "logs")
    
    logger_key = f"{name}_{level}_{log_dir}_{structured}"
    
    if logger_key not in _loggers:
        _loggers[logger_key] = ToxoLogger(
            name=name,
            level=level,
            log_dir=log_dir,
            structured=structured,
            enable_performance_logging=enable_performance_logging
        )
    
    return _loggers[logger_key]


def setup_logging(
    level: str = "INFO",
    log_dir: Optional[Union[str, Path]] = None,
    structured: bool = None
):
    """
    Setup global logging configuration.
    
    Args:
        level: Global log level
        log_dir: Log directory
        structured: Enable structured logging
    """
    # Set environment variables for default logger creation
    os.environ["TOXO_LOG_LEVEL"] = level
    if log_dir:
        os.environ["TOXO_LOG_DIR"] = str(log_dir)
    if structured is not None:
        os.environ["TOXO_STRUCTURED_LOGGING"] = str(structured).lower()


@contextmanager
def log_context(logger: ToxoLogger, **context):
    """
    Context manager for adding context to all log messages.
    
    Args:
        logger: Logger instance
        **context: Context variables to add
    """
    # This is a simplified implementation
    # In practice, you might want to use contextvars or similar
    original_bind = loguru_logger.bind
    
    def new_bind(**kwargs):
        return original_bind(**{**context, **kwargs})
    
    loguru_logger.bind = new_bind
    try:
        yield
    finally:
        loguru_logger.bind = original_bind


# Convenience function for common use case
def get_default_logger() -> ToxoLogger:
    """Get default Toxo logger."""
    return get_logger("toxo") 