import asyncio
import time
import random
from typing import TypeVar, Callable, Any
from functools import wraps
import structlog
from ..core.exceptions import OpenAIError, RateLimitError

T = TypeVar('T')
logger = structlog.get_logger("app.utils.retry")


def retry_with_exponential_backoff(
    max_retries: int = 3,
    base_delay: float = 1.0,
    max_delay: float = 60.0,
    backoff_factor: float = 2.0,
    jitter: bool = True
):
    def decorator(func: Callable[..., T]) -> Callable[..., T]:
        @wraps(func)
        async def wrapper(*args, **kwargs) -> T:
            last_exception = None
            
            for attempt in range(max_retries + 1):
                try:
                    if attempt > 0:
                        delay = min(base_delay * (backoff_factor ** (attempt - 1)), max_delay)
                        if jitter:
                            delay *= (0.5 + random.random() * 0.5)
                        
                        logger.info(
                            "retry_attempt",
                            attempt=attempt,
                            max_retries=max_retries,
                            delay=delay,
                            function=func.__name__
                        )
                        await asyncio.sleep(delay)
                    
                    result = await func(*args, **kwargs)
                    
                    if attempt > 0:
                        logger.info(
                            "retry_success",
                            attempt=attempt,
                            function=func.__name__
                        )
                    
                    return result
                    
                except Exception as e:
                    last_exception = e
                    
                    if hasattr(e, 'status_code') and e.status_code == 429:
                        retry_after = getattr(e, 'retry_after', None)
                        if retry_after and attempt < max_retries:
                            logger.warning(
                                "rate_limit_retry",
                                attempt=attempt,
                                retry_after=retry_after,
                                function=func.__name__
                            )
                            await asyncio.sleep(retry_after)
                            continue
                    
                    if attempt == max_retries:
                        logger.error(
                            "retry_exhausted",
                            attempts=attempt + 1,
                            function=func.__name__,
                            error=str(e)
                        )
                        break
                    
                    logger.warning(
                        "retry_failure",
                        attempt=attempt,
                        function=func.__name__,
                        error=str(e)
                    )
            
            raise last_exception
        
        return wrapper
    return decorator


class RetryMetrics:
    def __init__(self):
        self.total_attempts = 0
        self.successful_attempts = 0
        self.failed_attempts = 0
        self.total_retry_time = 0.0
        
    def record_attempt(self, success: bool, retry_time: float = 0.0):
        self.total_attempts += 1
        if success:
            self.successful_attempts += 1
        else:
            self.failed_attempts += 1
        self.total_retry_time += retry_time
    
    def get_stats(self) -> dict:
        success_rate = self.successful_attempts / self.total_attempts if self.total_attempts > 0 else 0
        avg_retry_time = self.total_retry_time / self.total_attempts if self.total_attempts > 0 else 0
        
        return {
            "total_attempts": self.total_attempts,
            "successful_attempts": self.successful_attempts,
            "failed_attempts": self.failed_attempts,
            "success_rate": success_rate,
            "avg_retry_time": avg_retry_time,
            "total_retry_time": self.total_retry_time
        }