from __future__ import annotations

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

from yuxi.utils.logging_config import logger


class RetryDecision(StrEnum):
    RETRY = "retry"
    CIRCUIT_BREAK = "circuit_break"
    FALLBACK = "fallback"
    FAIL = "fail"


@dataclass
class RetryConfig:
    max_retries: int = 3
    base_delay: float = 1.0
    max_delay: float = 60.0
    backoff_factor: float = 2.0
    jitter: bool = True
    retryable_exceptions: tuple[type[Exception], ...] = (Exception,)
    total_timeout: float | None = None
    circuit_breaker: Any = None
    metrics_callback: Callable[[str, dict], None] | None = None


async def with_retry(
    func: Callable[[], Awaitable[Any]],
    config: RetryConfig | None = None,
) -> Any:
    cfg = config or RetryConfig()
    last_exception: Exception | None = None
    start_time = time.monotonic()

    for attempt in range(cfg.max_retries + 1):
        if cfg.total_timeout is not None and (time.monotonic() - start_time) >= cfg.total_timeout:
            raise last_exception or Exception("Retry total_timeout exceeded")

        if cfg.circuit_breaker is not None:
            cb = cfg.circuit_breaker
            if hasattr(cb, "state") and getattr(cb, "state", None) is not None:
                from yuxi.channels.infra.circuit_breaker import CircuitState

                if cb.state == CircuitState.OPEN:
                    if cfg.metrics_callback:
                        cfg.metrics_callback("circuit_breaker_skip", {"attempt": attempt})
                    raise last_exception or Exception("Circuit breaker is OPEN, skipping retry")

        try:
            result = await func()
            if cfg.metrics_callback:
                cfg.metrics_callback("retry_success", {"attempt": attempt})
            return result
        except cfg.retryable_exceptions as e:
            last_exception = e
            if attempt == cfg.max_retries:
                break
            if cfg.total_timeout is not None and (time.monotonic() - start_time) >= cfg.total_timeout:
                break
            delay = min(cfg.base_delay * (cfg.backoff_factor**attempt), cfg.max_delay)
            if cfg.jitter:
                delay *= random.uniform(0.5, 1.5)
            if cfg.metrics_callback:
                cfg.metrics_callback("retry_attempt", {"attempt": attempt + 1, "error": str(e), "delay": delay})
            logger.debug(f"[SDK/Retry] Attempt {attempt + 1}/{cfg.max_retries} failed: {e}, retrying in {delay:.2f}s")
            await asyncio.sleep(delay)

    if cfg.metrics_callback:
        cfg.metrics_callback("retry_exhausted", {"attempts": cfg.max_retries, "error": str(last_exception)})
    raise last_exception  # type: ignore[misc]