from __future__ import annotations

import asyncio
import random

import httpx
from common_logging import get_logger
logger = get_logger(__name__)



_NO_RETRY_STATUS = frozenset({400, 401, 403, 404, 405, 406, 410, 422, 451})
_RETRY_STATUS = frozenset({429, 500, 502, 503, 504})

class RetryPolicy:

    def __init__(self, max_retries: int=3, base_delay: float=2.0, max_delay: float=60.0, jitter: float=0.3):
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.max_delay = max_delay
        self.jitter = jitter

    @staticmethod
    def should_retry_status(status_code: int) -> bool:
        if status_code in _NO_RETRY_STATUS:
            return False
        if status_code in _RETRY_STATUS:
            return True
        return status_code >= 500

    @staticmethod
    def should_retry_error(exc: BaseException) -> bool:
        return isinstance(exc, httpx.TimeoutException | httpx.ConnectError | httpx.ReadError | httpx.RemoteProtocolError | ConnectionError | TimeoutError)

    def compute_delay(self, attempt: int, status_code: int | None=None) -> float:
        base = self.base_delay * 3 if status_code == 429 else self.base_delay
        delay = min(base * 2 ** (attempt - 1), self.max_delay)
        delay *= 1 + random.uniform(-self.jitter, self.jitter)
        return max(delay, 0.1)

    async def execute(self, request_fn, *args, **kwargs) -> httpx.Response:
        last_exc: BaseException | None = None
        attempt = 0
        while attempt <= self.max_retries:
            try:
                resp: httpx.Response = await request_fn(*args, **kwargs)
                if resp.status_code < 400:
                    return resp
                if not self.should_retry_status(resp.status_code):
                    logger.warning(f"[retry] 不可重试状态码 {resp.status_code}，终止: {(args[0] if args else '')}")
                    resp.raise_for_status()
                if attempt >= self.max_retries:
                    logger.error(f'[retry] 达到最大重试 ({self.max_retries}) status={resp.status_code}')
                    resp.raise_for_status()
                delay = self.compute_delay(attempt + 1, resp.status_code)
                logger.warning(f'[retry] status={resp.status_code} attempt={attempt + 1}/{self.max_retries} 等待 {delay:.1f}s')
                await asyncio.sleep(delay)
                attempt += 1
            except httpx.HTTPStatusError:
                raise
            except Exception as exc:
                if not self.should_retry_error(exc):
                    logger.error(f'[retry] 不可重试异常 {type(exc).__name__}: {exc}')
                    raise
                last_exc = exc
                if attempt >= self.max_retries:
                    logger.error(f'[retry] 达到最大重试 ({self.max_retries}) error={type(exc).__name__}: {exc}')
                    raise
                delay = self.compute_delay(attempt + 1)
                logger.warning(f'[retry] {type(exc).__name__} attempt={attempt + 1}/{self.max_retries} 等待 {delay:.1f}s')
                await asyncio.sleep(delay)
                attempt += 1
        if last_exc:
            raise last_exc
        raise RuntimeError('retry policy internal error')
default_retry_policy = RetryPolicy(max_retries=3, base_delay=2.0, max_delay=60.0, jitter=0.3)