"""
Logging and logging utilities for InferenceKit.
Handles structured logging and log management.
"""

import logging
import logging.handlers
from typing import Dict, Any, Optional, Union, List
from datetime import datetime
import json
import sys


class Logger:
    """Main logger class for structured logging."""

    LOG_LEVELS = {
        "DEBUG": logging.DEBUG,
        "INFO": logging.INFO,
        "WARNING": logging.WARNING,
        "ERROR": logging.ERROR,
        "CRITICAL": logging.CRITICAL,
    }

    def __init__(
        self,
        level: str = "INFO",
        format: str = "json",
        log_file: Optional[str] = None,
        max_log_size: int = 10 * 1024 * 1024,  # 10MB
        backup_count: int = 5,
        **config
    ):
        self.level = level.upper()
        self.format = format
        self.log_file = log_file
        self.max_log_size = max_log_size
        self.backup_count = backup_count
        self.config = config

        # Initialize logger
        self.logger = logging.getLogger("inferencekit")
        self.logger.setLevel(self.LOG_LEVELS.get(self.level, logging.INFO))
        self.logger.propagate = False

        # Clear existing handlers
        self.logger.handlers.clear()

        # Add handlers
        self._add_console_handler()
        if self.log_file:
            self._add_file_handler()

    def _add_console_handler(self) -> None:
        """Add console handler for logging."""
        console_handler = logging.StreamHandler(sys.stdout)
        console_handler.setLevel(self.LOG_LEVELS.get(self.level, logging.INFO))
        console_handler.setFormatter(self._get_formatter())
        self.logger.addHandler(console_handler)

    def _add_file_handler(self) -> None:
        """Add file handler for logging."""
        file_handler = logging.handlers.RotatingFileHandler(
            self.log_file,
            maxBytes=self.max_log_size,
            backupCount=self.backup_count,
            encoding="utf-8"
        )
        file_handler.setLevel(self.LOG_LEVELS.get(self.level, logging.INFO))
        file_handler.setFormatter(self._get_formatter())
        self.logger.addHandler(file_handler)

    def _get_formatter(self) -> logging.Formatter:
        """Get appropriate formatter based on format setting."""
        if self.format == "json":
            return JsonFormatter()
        elif self.format == "structured":
            return StructuredFormatter()
        else:
            return logging.Formatter(
                fmt="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
                datefmt="%Y-%m-%d %H:%M:%S"
            )

    def debug(self, message: str, extra: Optional[Dict[str, Any]] = None) -> None:
        """Log a debug message."""
        self.logger.debug(message, extra=extra)

    def info(self, message: str, extra: Optional[Dict[str, Any]] = None) -> None:
        """Log an info message."""
        self.logger.info(message, extra=extra)

    def warning(self, message: str, extra: Optional[Dict[str, Any]] = None) -> None:
        """Log a warning message."""
        self.logger.warning(message, extra=extra)

    def error(self, message: str, extra: Optional[Dict[str, Any]] = None) -> None:
        """Log an error message."""
        self.logger.error(message, extra=extra)

    def critical(self, message: str, extra: Optional[Dict[str, Any]] = None) -> None:
        """Log a critical message."""
        self.logger.critical(message, extra=extra)

    def log(self, level: str, message: str, extra: Optional[Dict[str, Any]] = None) -> None:
        """Log a message at a specific level."""
        level_upper = level.upper()
        if level_upper in self.LOG_LEVELS:
            self.logger.log(self.LOG_LEVELS[level_upper], message, extra=extra)
        else:
            self.logger.error(f"Invalid log level: {level}")

    def exception(self, message: str, extra: Optional[Dict[str, Any]] = None) -> None:
        """Log an exception with traceback."""
        self.logger.exception(message, extra=extra)

    def getChild(self, name: str) -> logging.Logger:
        """Get a child logger."""
        return self.logger.getChild(name)

    def setLevel(self, level: str) -> None:
        """Set the logging level."""
        self.level = level.upper()
        self.logger.setLevel(self.LOG_LEVELS.get(self.level, logging.INFO))

    def getLevel(self) -> str:
        """Get the current logging level."""
        return self.level

    def hasHandlers(self) -> bool:
        """Check if logger has handlers."""
        return len(self.logger.handlers) > 0

    def getEffectiveLevel(self) -> int:
        """Get the effective logging level."""
        return self.logger.getEffectiveLevel()

    def isEnabledFor(self, level: str) -> bool:
        """Check if a specific level is enabled."""
        level_upper = level.upper()
        if level_upper in self.LOG_LEVELS:
            return self.logger.isEnabledFor(self.LOG_LEVELS[level_upper])
        return False


class JsonFormatter(logging.Formatter):
    """JSON formatter for structured logging."""

    def format(self, record: logging.LogRecord) -> str:
        """Format log record as JSON."""
        log_data = {
            "timestamp": datetime.utcnow().isoformat(),
            "level": record.levelname,
            "logger": record.name,
            "message": record.getMessage(),
            "module": record.module,
            "function": record.funcName,
            "line": record.lineno,
        }

        # Add extra attributes
        if hasattr(record, 'extra') and record.extra:
            log_data.update(record.extra)

        # Add exception info if present
        if record.exc_info:
            log_data["exception"] = self.formatException(record.exc_info)

        return json.dumps(log_data)


class StructuredFormatter(logging.Formatter):
    """Structured text formatter for logging."""

    def format(self, record: logging.LogRecord) -> str:
        """Format log record as structured text."""
        base_format = f"[{record.levelname}] {record.name}: {record.getMessage}"

        # Add extra attributes
        extra_info = []
        if hasattr(record, 'extra') and record.extra:
            for key, value in record.extra.items():
                extra_info.append(f"{key}={value}")

        if extra_info:
            return f"{base_format} ({', '.join(extra_info)})"
        return base_format


class LogFilter:
    """Filter for log messages based on various criteria."""

    def __init__(self, min_level: str = "INFO", excluded_modules: Optional[List[str]] = None):
        self.min_level = min_level.upper()
        self.excluded_modules = excluded_modules or []

    def filter(self, record: logging.LogRecord) -> bool:
        """Filter log records."""
        # Check level
        if record.levelname < self.min_level:
            return False

        # Check excluded modules
        if record.module in self.excluded_modules:
            return False

        return True


class LogHandler:
    """Custom log handler for specialized logging."""

    def __init__(
        self,
        handler_type: str = "stream",
        level: str = "INFO",
        formatter: Optional[logging.Formatter] = None,
        **config
    ):
        self.handler_type = handler_type
        self.level = level.upper()
        self.formatter = formatter
        self.config = config

        # Create appropriate handler
        if self.handler_type == "stream":
            self.handler = logging.StreamHandler(sys.stdout)
        elif self.handler_type == "file":
            self.handler = logging.FileHandler(
                config.get("filename", "inferencekit.log"),
                encoding="utf-8"
            )
        elif self.handler_type == "rotating_file":
            self.handler = logging.handlers.RotatingFileHandler(
                config.get("filename", "inferencekit.log"),
                maxBytes=config.get("max_bytes", 10 * 1024 * 1024),
                backupCount=config.get("backup_count", 5),
                encoding="utf-8"
            )
        else:
            raise ValueError(f"Unknown handler type: {self.handler_type}")

        self.handler.setLevel(self.LOG_LEVELS.get(self.level, logging.INFO))
        if self.formatter:
            self.handler.setFormatter(self.formatter)

    def get_handler(self) -> logging.Handler:
        """Get the configured handler."""
        return self.handler


class LogContext:
    """Context manager for logging context."""

    def __init__(self, logger: Logger, context_data: Dict[str, Any]):
        self.logger = logger
        self.context_data = context_data

    def __enter__(self):
        # Add context data to logger
        self.original_extra = getattr(self.logger.logger, 'extra', {})
        self.logger.logger.extra = self.context_data
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # Restore original extra data
        self.logger.logger.extra = self.original_extra

    def log(self, level: str, message: str) -> None:
        """Log a message with context."""
        self.logger.log(level, message, extra=self.context_data)