"""Logging utilities for the Trax platform.

This module provides utility functions for:
- Performance metrics logging
- Context management
- Correlation ID tracking
- Error context capture
"""

import asyncio
import contextvars
import functools
import time
import traceback
import uuid
from contextlib import asynccontextmanager, contextmanager
from typing import Any, Callable, Dict, Optional, TypeVar, cast

from . import get_logger

# Context variables for correlation tracking
correlation_id = contextvars.ContextVar('correlation_id', default=None)
operation_context = contextvars.ContextVar('operation_context', default={})

T = TypeVar('T')


class PerformanceMetrics:
    """Context manager for performance metrics logging."""
    
    def __init__(self, operation: str, logger_name: Optional[str] = None, threshold_ms: int = 1000):
        self.operation = operation
        self.logger = get_logger(logger_name or __name__)
        self.threshold_ms = threshold_ms
        self.start_time = None
        self.end_time = None
        self.duration_ms = None
    
    def __enter__(self):
        self.start_time = time.time()
        self.logger.debug(f"Starting operation: {self.operation}")
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.end_time = time.time()
        self.duration_ms = (self.end_time - self.start_time) * 1000
        
        if exc_type is not None:
            self.logger.error(
                f"Operation failed: {self.operation}",
                extra={
                    "operation": self.operation,
                    "duration_ms": self.duration_ms,
                    "error": str(exc_val),
                    "error_type": exc_type.__name__ if exc_type else None,
                    "correlation_id": correlation_id.get(),
                    "context": operation_context.get()
                }
            )
        else:
            if self.duration_ms >= self.threshold_ms:
                self.logger.warning(
                    f"Slow operation: {self.operation}",
                    extra={
                        "operation": self.operation,
                        "duration_ms": self.duration_ms,
                        "threshold_ms": self.threshold_ms,
                        "correlation_id": correlation_id.get(),
                        "context": operation_context.get()
                    }
                )
            else:
                self.logger.debug(
                    f"Operation completed: {self.operation}",
                    extra={
                        "operation": self.operation,
                        "duration_ms": self.duration_ms,
                        "correlation_id": correlation_id.get(),
                        "context": operation_context.get()
                    }
                )
    
    async def __aenter__(self):
        self.start_time = time.time()
        self.logger.debug(f"Starting async operation: {self.operation}")
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        self.end_time = time.time()
        self.duration_ms = (self.end_time - self.start_time) * 1000
        
        if exc_type is not None:
            self.logger.error(
                f"Async operation failed: {self.operation}",
                extra={
                    "operation": self.operation,
                    "duration_ms": self.duration_ms,
                    "error": str(exc_val),
                    "error_type": exc_type.__name__ if exc_type else None,
                    "correlation_id": correlation_id.get(),
                    "context": operation_context.get()
                }
            )
        else:
            if self.duration_ms >= self.threshold_ms:
                self.logger.warning(
                    f"Slow async operation: {self.operation}",
                    extra={
                        "operation": self.operation,
                        "duration_ms": self.duration_ms,
                        "threshold_ms": self.threshold_ms,
                        "correlation_id": correlation_id.get(),
                        "context": operation_context.get()
                    }
                )
            else:
                self.logger.debug(
                    f"Async operation completed: {self.operation}",
                    extra={
                        "operation": self.operation,
                        "duration_ms": self.duration_ms,
                        "correlation_id": correlation_id.get(),
                        "context": operation_context.get()
                    }
                )


@contextmanager
def correlation_context(corr_id: Optional[str] = None):
    """Context manager for correlation ID tracking."""
    if corr_id is None:
        corr_id = str(uuid.uuid4())
    
    token = correlation_id.set(corr_id)
    try:
        yield corr_id
    finally:
        correlation_id.reset(token)


@contextmanager
def operation_context_manager(context_data: Dict[str, Any]):
    """Context manager for operation context."""
    current_context = operation_context.get().copy()
    current_context.update(context_data)
    
    token = operation_context.set(current_context)
    try:
        yield current_context
    finally:
        operation_context.reset(token)


def log_function_call(func: Callable[..., T]) -> Callable[..., T]:
    """Decorator to log function calls with performance metrics."""
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        logger = get_logger(func.__module__)
        
        with PerformanceMetrics(f"{func.__module__}.{func.__name__}", logger_name=func.__module__):
            try:
                result = func(*args, **kwargs)
                logger.debug(
                    f"Function call completed: {func.__name__}",
                    extra={
                        "function": func.__name__,
                        "module": func.__module__,
                        "correlation_id": correlation_id.get(),
                        "context": operation_context.get()
                    }
                )
                return result
            except Exception as e:
                logger.error(
                    f"Function call failed: {func.__name__}",
                    extra={
                        "function": func.__name__,
                        "module": func.__module__,
                        "error": str(e),
                        "error_type": type(e).__name__,
                        "correlation_id": correlation_id.get(),
                        "context": operation_context.get()
                    }
                )
                raise
    
    return cast(Callable[..., T], wrapper)


def log_async_function_call(func: Callable[..., T]) -> Callable[..., T]:
    """Decorator to log async function calls with performance metrics."""
    @functools.wraps(func)
    async def wrapper(*args, **kwargs):
        logger = get_logger(func.__module__)
        
        async with PerformanceMetrics(f"{func.__module__}.{func.__name__}", logger_name=func.__module__):
            try:
                result = await func(*args, **kwargs)
                logger.debug(
                    f"Async function call completed: {func.__name__}",
                    extra={
                        "function": func.__name__,
                        "module": func.__module__,
                        "correlation_id": correlation_id.get(),
                        "context": operation_context.get()
                    }
                )
                return result
            except Exception as e:
                logger.error(
                    f"Async function call failed: {func.__name__}",
                    extra={
                        "function": func.__name__,
                        "module": func.__module__,
                        "error": str(e),
                        "error_type": type(e).__name__,
                        "correlation_id": correlation_id.get(),
                        "context": operation_context.get()
                    }
                )
                raise
    
    return cast(Callable[..., T], wrapper)


def capture_error_context(error: Exception, depth: int = 3) -> Dict[str, Any]:
    """Capture error context for logging."""
    tb = traceback.extract_tb(error.__traceback__)
    
    context = {
        "error_type": type(error).__name__,
        "error_message": str(error),
        "traceback": [
            {
                "filename": frame.filename,
                "line": frame.lineno,
                "function": frame.name,
                "line_content": frame.line
            }
            for frame in tb[-depth:] if frame
        ],
        "correlation_id": correlation_id.get(),
        "operation_context": operation_context.get()
    }
    
    return context


def log_error_with_context(error: Exception, message: str = "An error occurred", logger_name: Optional[str] = None, depth: int = 3):
    """Log an error with full context information."""
    logger = get_logger(logger_name or __name__)
    context = capture_error_context(error, depth)
    
    logger.error(message, extra=context)


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


def get_operation_context() -> Dict[str, Any]:
    """Get the current operation context."""
    return operation_context.get()


def set_correlation_id(corr_id: str) -> None:
    """Set the correlation ID for the current context."""
    correlation_id.set(corr_id)


def add_context_data(key: str, value: Any) -> None:
    """Add data to the current operation context."""
    current_context = operation_context.get().copy()
    current_context[key] = value
    operation_context.set(current_context)


def clear_context() -> None:
    """Clear the current operation context."""
    operation_context.set({})


# Convenience functions for common logging patterns
def log_operation_start(operation: str, **context_data):
    """Log the start of an operation."""
    logger = get_logger(__name__)
    logger.info(
        f"Starting operation: {operation}",
        extra={
            "operation": operation,
            "correlation_id": correlation_id.get(),
            "context": {**operation_context.get(), **context_data}
        }
    )


def log_operation_complete(operation: str, duration_ms: Optional[float] = None, **context_data):
    """Log the completion of an operation."""
    logger = get_logger(__name__)
    extra_data = {
        "operation": operation,
        "correlation_id": correlation_id.get(),
        "context": {**operation_context.get(), **context_data}
    }
    
    if duration_ms is not None:
        extra_data["duration_ms"] = duration_ms
    
    logger.info(f"Operation completed: {operation}", extra=extra_data)


def log_operation_failed(operation: str, error: Exception, **context_data):
    """Log the failure of an operation."""
    logger = get_logger(__name__)
    error_context = capture_error_context(error)
    
    logger.error(
        f"Operation failed: {operation}",
        extra={
            "operation": operation,
            "correlation_id": correlation_id.get(),
            "context": {**operation_context.get(), **context_data},
            **error_context
        }
    )
