#!/usr/bin/env python3
"""
GENESIS STANDARDIZED LOGGING SYSTEM
=====================================
Unified logging with structured output, multiple handlers, and correlation IDs.

Features:
    - Structured JSON logging
    - Correlation ID tracking
    - Multiple output handlers
    - Log level filtering
    - Context enrichment
    - Performance metrics

Usage:
    logger = get_logger("module_name")
    logger.info("Message", extra={"user": "test"})

    with logger.context(request_id="123"):
        logger.info("Within context")
"""

import json
import logging
import sys
import threading
import time
import traceback
import uuid
from contextlib import contextmanager
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Any, Optional, Callable, Union
from enum import Enum
import queue


class LogLevel(Enum):
    """Log levels."""
    DEBUG = 10
    INFO = 20
    WARNING = 30
    ERROR = 40
    CRITICAL = 50


@dataclass
class LogRecord:
    """A structured log record."""
    timestamp: str
    level: str
    logger_name: str
    message: str
    correlation_id: Optional[str] = None
    context: Dict[str, Any] = field(default_factory=dict)
    exception: Optional[str] = None
    source_file: Optional[str] = None
    source_line: Optional[int] = None
    function_name: Optional[str] = None
    duration_ms: Optional[float] = None

    def to_dict(self) -> Dict:
        result = {
            "timestamp": self.timestamp,
            "level": self.level,
            "logger": self.logger_name,
            "message": self.message
        }
        if self.correlation_id:
            result["correlation_id"] = self.correlation_id
        if self.context:
            result["context"] = self.context
        if self.exception:
            result["exception"] = self.exception
        if self.source_file:
            result["source"] = {
                "file": self.source_file,
                "line": self.source_line,
                "function": self.function_name
            }
        if self.duration_ms is not None:
            result["duration_ms"] = self.duration_ms
        return result

    def to_json(self) -> str:
        return json.dumps(self.to_dict())

    def to_text(self) -> str:
        parts = [
            f"[{self.timestamp}]",
            f"[{self.level:8}]",
            f"[{self.logger_name}]"
        ]
        if self.correlation_id:
            parts.append(f"[{self.correlation_id[:8]}]")
        parts.append(self.message)

        if self.context:
            ctx_str = " ".join(f"{k}={v}" for k, v in self.context.items())
            parts.append(f"| {ctx_str}")

        if self.duration_ms is not None:
            parts.append(f"({self.duration_ms:.2f}ms)")

        return " ".join(parts)


class LogHandler:
    """Base class for log handlers."""

    def __init__(self, level: LogLevel = LogLevel.DEBUG):
        self.level = level

    def handle(self, record: LogRecord):
        raise NotImplementedError

    def should_handle(self, level: str) -> bool:
        level_value = LogLevel[level].value
        return level_value >= self.level.value


class ConsoleHandler(LogHandler):
    """Handler that outputs to console."""

    def __init__(self, level: LogLevel = LogLevel.DEBUG, json_format: bool = False, colorize: bool = True):
        super().__init__(level)
        self.json_format = json_format
        self.colorize = colorize

        # ANSI color codes
        self.colors = {
            "DEBUG": "\033[36m",     # Cyan
            "INFO": "\033[32m",      # Green
            "WARNING": "\033[33m",   # Yellow
            "ERROR": "\033[31m",     # Red
            "CRITICAL": "\033[35m",  # Magenta
            "RESET": "\033[0m"
        }

    def handle(self, record: LogRecord):
        if not self.should_handle(record.level):
            return

        if self.json_format:
            output = record.to_json()
        else:
            output = record.to_text()
            if self.colorize:
                color = self.colors.get(record.level, "")
                reset = self.colors["RESET"]
                output = f"{color}{output}{reset}"

        print(output, file=sys.stderr)

        if record.exception:
            print(record.exception, file=sys.stderr)


class FileHandler(LogHandler):
    """Handler that outputs to file."""

    def __init__(
        self,
        file_path: Path,
        level: LogLevel = LogLevel.DEBUG,
        json_format: bool = True,
        max_size_mb: int = 10,
        backup_count: int = 5
    ):
        super().__init__(level)
        self.file_path = Path(file_path)
        self.json_format = json_format
        self.max_size = max_size_mb * 1024 * 1024
        self.backup_count = backup_count
        self._lock = threading.Lock()

        self.file_path.parent.mkdir(parents=True, exist_ok=True)

    def handle(self, record: LogRecord):
        if not self.should_handle(record.level):
            return

        with self._lock:
            # Check if rotation needed
            if self.file_path.exists() and self.file_path.stat().st_size > self.max_size:
                self._rotate()

            output = record.to_json() if self.json_format else record.to_text()

            with open(self.file_path, 'a', encoding='utf-8') as f:
                f.write(output + '\n')

    def _rotate(self):
        """Rotate log files."""
        for i in range(self.backup_count - 1, 0, -1):
            src = self.file_path.with_suffix(f".{i}.log")
            dst = self.file_path.with_suffix(f".{i+1}.log")
            if src.exists():
                src.rename(dst)

        if self.file_path.exists():
            self.file_path.rename(self.file_path.with_suffix(".1.log"))


class AsyncHandler(LogHandler):
    """Handler that processes logs asynchronously."""

    def __init__(self, wrapped_handler: LogHandler, buffer_size: int = 1000):
        super().__init__(wrapped_handler.level)
        self.wrapped = wrapped_handler
        self._queue: queue.Queue = queue.Queue(maxsize=buffer_size)
        self._running = True
        self._worker = threading.Thread(target=self._process_loop, daemon=True)
        self._worker.start()

    def handle(self, record: LogRecord):
        if not self.should_handle(record.level):
            return
        try:
            self._queue.put_nowait(record)
        except queue.Full:
            pass  # Drop log if queue is full

    def _process_loop(self):
        while self._running:
            try:
                record = self._queue.get(timeout=0.5)
                self.wrapped.handle(record)
            except queue.Empty:
                continue

    def flush(self):
        """Wait for queue to empty."""
        while not self._queue.empty():
            time.sleep(0.1)

    def stop(self):
        self._running = False
        self._worker.join(timeout=2)


class CallbackHandler(LogHandler):
    """Handler that calls a callback function."""

    def __init__(self, callback: Callable[[LogRecord], None], level: LogLevel = LogLevel.DEBUG):
        super().__init__(level)
        self.callback = callback

    def handle(self, record: LogRecord):
        if not self.should_handle(record.level):
            return
        try:
            self.callback(record)
        except Exception:
            pass


# Thread-local storage for context
_context_storage = threading.local()


def get_context() -> Dict[str, Any]:
    """Get current thread context."""
    if not hasattr(_context_storage, 'context'):
        _context_storage.context = {}
    return _context_storage.context


def set_context(key: str, value: Any):
    """Set a context value."""
    get_context()[key] = value


def clear_context():
    """Clear thread context."""
    _context_storage.context = {}


class GenesisLogger:
    """
    Structured logger for Genesis.
    """

    def __init__(self, name: str, handlers: List[LogHandler] = None):
        self.name = name
        self.handlers = handlers or []
        self._timers: Dict[str, float] = {}

    def add_handler(self, handler: LogHandler):
        """Add a log handler."""
        self.handlers.append(handler)

    def _log(
        self,
        level: str,
        message: str,
        extra: Dict[str, Any] = None,
        exc_info: bool = False
    ):
        """Internal logging method."""
        # Get caller info
        import inspect
        frame = inspect.currentframe()
        caller_frame = frame.f_back.f_back if frame else None

        source_file = None
        source_line = None
        function_name = None

        if caller_frame:
            source_file = caller_frame.f_code.co_filename
            source_line = caller_frame.f_lineno
            function_name = caller_frame.f_code.co_name

        # Merge context
        context = get_context().copy()
        if extra:
            context.update(extra)

        # Get exception info
        exception = None
        if exc_info:
            exception = traceback.format_exc()

        # Get correlation ID from context
        correlation_id = context.pop('correlation_id', None) or context.pop('request_id', None)

        record = LogRecord(
            timestamp=datetime.now().isoformat(),
            level=level,
            logger_name=self.name,
            message=message,
            correlation_id=correlation_id,
            context=context,
            exception=exception,
            source_file=source_file,
            source_line=source_line,
            function_name=function_name
        )

        for handler in self.handlers:
            try:
                handler.handle(record)
            except Exception:
                pass

    def debug(self, message: str, **kwargs):
        self._log("DEBUG", message, kwargs.get('extra'))

    def info(self, message: str, **kwargs):
        self._log("INFO", message, kwargs.get('extra'))

    def warning(self, message: str, **kwargs):
        self._log("WARNING", message, kwargs.get('extra'))

    def error(self, message: str, exc_info: bool = False, **kwargs):
        self._log("ERROR", message, kwargs.get('extra'), exc_info)

    def critical(self, message: str, exc_info: bool = False, **kwargs):
        self._log("CRITICAL", message, kwargs.get('extra'), exc_info)

    def exception(self, message: str, **kwargs):
        """Log with exception info."""
        self.error(message, exc_info=True, **kwargs)

    @contextmanager
    def context(self, **kwargs):
        """Context manager for adding context to logs."""
        old_context = get_context().copy()
        get_context().update(kwargs)
        try:
            yield
        finally:
            _context_storage.context = old_context

    def start_timer(self, name: str):
        """Start a named timer."""
        self._timers[name] = time.time()

    def stop_timer(self, name: str, message: str = None, level: str = "DEBUG"):
        """Stop timer and log duration."""
        if name not in self._timers:
            return

        duration = (time.time() - self._timers[name]) * 1000
        del self._timers[name]

        record = LogRecord(
            timestamp=datetime.now().isoformat(),
            level=level,
            logger_name=self.name,
            message=message or f"Timer '{name}' completed",
            correlation_id=get_context().get('correlation_id'),
            context=get_context().copy(),
            duration_ms=duration
        )

        for handler in self.handlers:
            handler.handle(record)

    @contextmanager
    def timed(self, operation: str, level: str = "DEBUG"):
        """Context manager for timing operations."""
        start = time.time()
        try:
            yield
        finally:
            duration = (time.time() - start) * 1000
            record = LogRecord(
                timestamp=datetime.now().isoformat(),
                level=level,
                logger_name=self.name,
                message=f"{operation} completed",
                correlation_id=get_context().get('correlation_id'),
                context=get_context().copy(),
                duration_ms=duration
            )
            for handler in self.handlers:
                handler.handle(record)


# Global logger registry
_loggers: Dict[str, GenesisLogger] = {}
_global_handlers: List[LogHandler] = []
_initialized = False


def configure(
    level: LogLevel = LogLevel.INFO,
    json_format: bool = False,
    log_file: Path = None,
    colorize: bool = True
):
    """Configure global logging."""
    global _global_handlers, _initialized

    _global_handlers = []

    # Console handler
    console = ConsoleHandler(level=level, json_format=json_format, colorize=colorize)
    _global_handlers.append(console)

    # File handler
    if log_file:
        file_handler = FileHandler(log_file, level=level, json_format=True)
        _global_handlers.append(file_handler)

    _initialized = True


def get_logger(name: str) -> GenesisLogger:
    """Get or create a logger by name."""
    global _initialized

    if not _initialized:
        configure()

    if name not in _loggers:
        logger = GenesisLogger(name, _global_handlers.copy())
        _loggers[name] = logger

    return _loggers[name]


def with_correlation_id(correlation_id: str = None):
    """Decorator to add correlation ID to all logs in a function."""
    def decorator(func):
        def wrapper(*args, **kwargs):
            cid = correlation_id or str(uuid.uuid4())[:8]
            set_context('correlation_id', cid)
            try:
                return func(*args, **kwargs)
            finally:
                get_context().pop('correlation_id', None)
        return wrapper
    return decorator


def main():
    """CLI and demo for genesis logger."""
    import argparse
    parser = argparse.ArgumentParser(description="Genesis Logger")
    parser.add_argument("command", choices=["demo", "test"])
    parser.add_argument("--json", action="store_true", help="JSON output")
    parser.add_argument("--file", help="Log file path")
    args = parser.parse_args()

    if args.command == "demo":
        # Configure logging
        configure(
            level=LogLevel.DEBUG,
            json_format=args.json,
            log_file=Path(args.file) if args.file else None
        )

        logger = get_logger("demo")

        print("Genesis Logger Demo")
        print("=" * 40)

        # Basic logging
        logger.debug("Debug message")
        logger.info("Info message")
        logger.warning("Warning message")
        logger.error("Error message")

        # With context
        with logger.context(user="test_user", request_id="req-123"):
            logger.info("Request started")
            logger.info("Processing")
            logger.info("Request completed")

        # With extra data
        logger.info("Task completed", extra={"task_id": "task-456", "duration": 150})

        # Timing
        with logger.timed("Database query"):
            time.sleep(0.1)

        # Exception logging
        try:
            raise ValueError("Test exception")
        except Exception:
            logger.exception("An error occurred")

        print("\nDemo complete!")

    elif args.command == "test":
        configure(level=LogLevel.DEBUG)
        logger = get_logger("test")

        # Test all levels
        for level in ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"]:
            logger._log(level, f"Test {level.lower()} message")

        print("Test complete!")


if __name__ == "__main__":
    main()
