from __future__ import annotations

import asyncio
import logging
import random
import time
from dataclasses import dataclass
from typing import Any
from collections.abc import Awaitable, Callable

logger = logging.getLogger(__name__)

BackoffFn = Callable[[int], float]


@dataclass
class BackoffConfig:
    base_delay: float = 5.0
    max_delay: float = 300.0
    exponent: float = 2.0
    jitter: bool = True
    jitter_factor: float = 0.1
    max_retries: int = 10

    def compute_delay(self, attempt: int) -> float:
        delay = self.base_delay * (self.exponent**attempt)
        delay = min(delay, self.max_delay)
        if self.jitter:
            jitter_amount = delay * self.jitter_factor
            delay += random.uniform(-jitter_amount, jitter_amount)
            delay = max(delay, 0.01)
        return delay


class ErrorBackoff:
    """错误退避控制器 — 管理重试逻辑与指数退避。

    使用按 key 分锁策略，不同 key 之间的退避操作互不阻塞，
    避免粗粒度单锁成为高并发场景下的全局瓶颈。
    """

    def __init__(self, config: BackoffConfig | None = None):
        self.config = config or BackoffConfig()
        self._attempts: dict[str, int] = {}
        self._last_error: dict[str, float] = {}
        self._locks: dict[str, asyncio.Lock] = {}
        self._global_lock = asyncio.Lock()

    def _get_lock(self, key: str) -> asyncio.Lock:
        """获取指定 key 的独立锁，必要时惰性创建。"""
        lock = self._locks.get(key)
        if lock is None:
            lock = asyncio.Lock()
            self._locks[key] = lock
        return lock

    async def execute(
        self,
        key: str,
        operation: Callable[[], Awaitable[Any]],
        *,
        on_retry: Callable[[int, Exception], Awaitable[None]] | None = None,
    ) -> Any:
        from yuxi.channel.errors import _infer_severity
        from yuxi.channel.protocols import ErrorSeverity

        lock = self._get_lock(key)
        async with lock:
            attempt = self._attempts.get(key, 0)

            for i in range(self.config.max_retries + 1):
                try:
                    result = await operation()
                    self._attempts[key] = 0
                    self._last_error.pop(key, None)
                    return result
                except Exception as e:
                    self._attempts[key] = attempt + i + 1
                    self._last_error[key] = time.time()

                    if _infer_severity(e) == ErrorSeverity.FATAL:
                        logger.error("Operation '%s' failed with fatal error: %s", key, e)
                        raise

                    if i >= self.config.max_retries:
                        logger.error("Operation '%s' failed after %d retries: %s", key, self.config.max_retries, e)
                        raise

                    delay = self.config.compute_delay(attempt + i)
                    logger.warning(
                        "Operation '%s' failed (attempt %d/%d), retrying in %.2fs: %s",
                        key,
                        i + 1,
                        self.config.max_retries + 1,
                        delay,
                        e,
                    )

                    if on_retry:
                        await on_retry(i, e)

                    await asyncio.sleep(delay)

    def get_attempts(self, key: str) -> int:
        return self._attempts.get(key, 0)

    def reset(self, key: str) -> None:
        self._attempts.pop(key, None)
        self._last_error.pop(key, None)
        self._locks.pop(key, None)

    def is_backing_off(self, key: str, cooldown_seconds: float = 300.0) -> bool:
        last_error = self._last_error.get(key)
        if last_error is None:
            return False
        return time.time() - last_error < cooldown_seconds

    def get_backoff_remaining(self, key: str, cooldown_seconds: float = 300.0) -> float:
        last_error = self._last_error.get(key)
        if last_error is None:
            return 0.0
        remaining = cooldown_seconds - (time.time() - last_error)
        return max(remaining, 0.0)