# ================================================================
# shared/monitoring/decorators.py
# ================================================================
"""
Decorators for monitoring and metrics
"""

import time
import logging
import functools
from typing import Callable, Any, Optional
from .metrics import get_metrics_collector

logger = logging.getLogger(__name__)

def track_execution_time(metric_name: str, service_name: str, 
                        labels: Optional[dict] = None):
    """Decorator to track function execution time"""
    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs) -> Any:
            start_time = time.time()
            try:
                result = await func(*args, **kwargs)
                duration = time.time() - start_time
                
                metrics_collector = get_metrics_collector(service_name)
                if metrics_collector:
                    metrics_collector.observe_histogram(metric_name, duration, labels)
                
                return result
            except Exception as e:
                duration = time.time() - start_time
                logger.error(f"Function {func.__name__} failed after {duration:.2f}s: {e}")
                raise
        
        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs) -> Any:
            start_time = time.time()
            try:
                result = func(*args, **kwargs)
                duration = time.time() - start_time
                
                metrics_collector = get_metrics_collector(service_name)
                if metrics_collector:
                    metrics_collector.observe_histogram(metric_name, duration, labels)
                
                return result
            except Exception as e:
                duration = time.time() - start_time
                logger.error(f"Function {func.__name__} failed after {duration:.2f}s: {e}")
                raise
        
        # Return appropriate wrapper based on function type
        import asyncio
        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        else:
            return sync_wrapper
    
    return decorator


def count_calls(metric_name: str, service_name: str, 
                labels: Optional[dict] = None):
    """Decorator to count function calls"""
    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs) -> Any:
            metrics_collector = get_metrics_collector(service_name)
            if metrics_collector:
                metrics_collector.increment_counter(metric_name, labels=labels)
            return await func(*args, **kwargs)
        
        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs) -> Any:
            metrics_collector = get_metrics_collector(service_name)
            if metrics_collector:
                metrics_collector.increment_counter(metric_name, labels=labels)
            return func(*args, **kwargs)
        
        # Return appropriate wrapper based on function type
        import asyncio
        if asyncio.iscoroutinefunction(func):
            return async_wrapper
        else:
            return sync_wrapper
    
    return decorator