"""
智能重试执行器
支持指数退避
"""
import asyncio
import random
from typing import Optional, Dict, Any, Callable, Awaitable, Union, TypeVar, List
from datetime import datetime
from loguru import logger

from .types import RetryConfig, RetryResult

T = TypeVar('T')


class RetryExecutor:
    """智能重试执行器，支持指数退避"""

    def __init__(self, config: Optional[RetryConfig] = None):
        self.config = config or RetryConfig()
        self._stats = {"total": 0, "successful": 0, "failed": 0, "retries": 0}

    async def execute(
        self,
        func: Union[Callable[[], Awaitable[T]], Callable[[], T]],
        **override
    ) -> RetryResult[T]:
        config = RetryConfig(
            max_retries=override.get("max_retries", self.config.max_retries),
            base_delay=override.get("base_delay", self.config.base_delay),
            max_delay=override.get("max_delay", self.config.max_delay),
            exponential_base=override.get("exponential_base", self.config.exponential_base),
            jitter=override.get("jitter", self.config.jitter),
            timeout=override.get("timeout", self.config.timeout),
            retry_on=override.get("retry_on", self.config.retry_on),
        )
        result = RetryResult[T]()
        self._stats["total"] += 1

        while result.attempts < config.max_retries + 1:
            result.attempts += 1
            try:
                if asyncio.iscoroutinefunction(func):
                    value = await (
                        asyncio.wait_for(func(), timeout=config.timeout)
                        if config.timeout else func()
                    )
                else:
                    value = func()
                result.success = True
                result.result = value
                result.end_time = datetime.now()
                self._stats["successful"] += 1
                return result
            except asyncio.CancelledError:
                result.end_time = datetime.now()
                self._stats["failed"] += 1
                raise
            except Exception as e:
                result.errors.append(e)
                result.error = e
                if result.attempts >= config.max_retries + 1:
                    break
                delay = self._calculate_delay(result.attempts, config)
                result.total_delay += delay
                self._stats["retries"] += 1
                await asyncio.sleep(delay)

        result.success = False
        result.end_time = datetime.now()
        self._stats["failed"] += 1
        return result

    def _calculate_delay(self, attempt: int, config: RetryConfig) -> float:
        delay = config.base_delay * (config.exponential_base ** (attempt - 1))
        delay = min(delay, config.max_delay)
        if config.jitter:
            delay += delay * config.jitter_range * (random.random() - 0.5)
        return delay

    def get_stats(self) -> Dict[str, Any]:
        return self._stats.copy()

    def reset_stats(self):
        """重置统计数据"""
        self._stats = {"total": 0, "successful": 0, "failed": 0, "retries": 0}

    async def execute_with_callback(
        self,
        func: Union[Callable[[], Awaitable[T]], Callable[[], T]],
        on_success: Optional[Callable[[T], Awaitable[None]]] = None,
        on_failure: Optional[Callable[[Exception], Awaitable[None]]] = None,
        on_retry: Optional[Callable[[int, Exception], Awaitable[None]]] = None,
        **override
    ) -> RetryResult[T]:
        """带回调的执行"""
        result = await self.execute(func, **override)

        if result.success and on_success:
            await on_success(result.result)
        elif not result.success and on_failure:
            await on_failure(result.error)
        elif result.attempts > 1 and on_retry:
            await on_retry(result.attempts, result.error)

        return result

    async def execute_batch(
        self,
        funcs: List[Union[Callable[[], Awaitable[T]], Callable[[], T]]],
        **override
    ) -> List[RetryResult[T]]:
        """批量执行"""
        results = []
        for func in funcs:
            result = await self.execute(func, **override)
            results.append(result)
        return results


def with_retry(
    max_retries: int = 3,
    base_delay: float = 1.0,
    **retry_options
):
    """重试装饰器"""
    def decorator(func):
        async def wrapper(*args, **kwargs):
            executor = RetryExecutor(RetryConfig(
                max_retries=max_retries,
                base_delay=base_delay,
                **retry_options
            ))

            async def wrapped():
                return await func(*args, **kwargs)

            result = await executor.execute(wrapped)
            if result.success:
                return result.result
            raise result.error
        return wrapper
    return decorator
