# -*- coding: utf-8 -*-
"""
LLM 韧性套件：断路器 + 指数退避重试（Phase 6 AN）

借鉴 zhikuncode-main 的 ApiCircuitBreaker.java + RetryPolicy.java：
- 断路器三态机：CLOSED（正常） → OPEN（拒绝） → HALF_OPEN（探测） → CLOSED/HALF_OPEN
- 指数退避 + jitter：BASE_DELAY=500ms，最大 32s，抖动 0-25%
- 可重试错误分类：429/529/408/5xx/ECONNRESET 等

设计动机（abyss 当前状态）：
- src/abyss/api_client.py 仅有 RETRYABLE_CODES = {429, 500, 503} 简单集合
- 重试间隔固定 time.sleep(min(2 ** (attempt - 1), 8))，无 jitter
- 无断路器，连续失败时仍持续重试，可能加剧 API 雪崩
- 无 OPEN 态快速失败，浪费用户时间

与 Phase 3 I (LLM 韧性四件套) 的关系：本项提供熔断 + 重试的具体实现，
是 Phase 3 I 的前置基础。ApiKeyRotationManager / ModelDegradationChain 后续推进。
"""
import random
import threading
import time
from enum import Enum
from typing import Any, Callable, Optional, Set, Type

from openai import (
    APIError,
    APIConnectionError,
    APITimeoutError,
    InternalServerError,
    RateLimitError,
)


class CircuitState(Enum):
    """断路器三态枚举。

    - CLOSED：正常调用，失败计数达阈值（默认 5 次）转 OPEN
    - OPEN：直接拒绝请求（fast-fail），冷却时间后转 HALF_OPEN
    - HALF_OPEN：放一个探测请求，成功转 CLOSED，失败转 OPEN
    """

    CLOSED = "closed"
    OPEN = "open"
    HALF_OPEN = "half_open"


class CircuitOpenError(Exception):
    """断路器开启时抛出，表示请求被 fast-fail。

    调用方应向上抛此异常，让用户感知到"服务暂时不可用"，
    而非继续重试浪费时间。
    """

    def __init__(self, message: str = "API 断路器开启，请稍后重试"):
        super().__init__(message)


class ApiCircuitBreaker:
    """API 断路器：CLOSED → OPEN → HALF_OPEN → CLOSED 状态机。

    线程安全：所有状态变更使用 threading.Lock 保护。
    设计参考 Michael Nygard 的《Release It!》中的 Circuit Breaker 模式。

    使用示例：

        breaker = ApiCircuitBreaker(failure_threshold=5, recovery_timeout=60)

        if not breaker.allow_request():
            raise CircuitOpenError()

        try:
            result = api_call()
            breaker.record_success()
            return result
        except Exception as e:
            if classifier.is_retryable(e):
                breaker.record_failure()
            raise
    """

    def __init__(
        self,
        failure_threshold: int = 5,
        recovery_timeout: float = 60.0,
    ):
        """初始化断路器。

        Args:
            failure_threshold: 连续失败多少次后转 OPEN（默认 5）
            recovery_timeout: OPEN 态冷却时间（秒，默认 60）
        """
        self.failure_threshold = failure_threshold
        self.recovery_timeout = recovery_timeout

        self._state = CircuitState.CLOSED
        self._failure_count = 0
        self._last_failure_time: Optional[float] = None
        self._lock = threading.Lock()
        # HALF_OPEN 探测中标志：True 表示已放行一个探测请求，等待结果
        self._half_open_probe_in_flight = False

    @property
    def state(self) -> CircuitState:
        """当前状态（无锁读，CAS 保证可见性）。"""
        return self._state

    @property
    def failure_count(self) -> int:
        """当前失败计数（无锁读）。"""
        return self._failure_count

    def allow_request(self) -> bool:
        """是否允许请求通过。

        - CLOSED：始终允许
        - OPEN：冷却时间到则转 HALF_OPEN 并允许；否则拒绝
        - HALF_OPEN：仅允许第一个探测请求，后续拒绝直到结果回来

        Returns:
            True 表示允许调用，False 表示 fast-fail
        """
        with self._lock:
            if self._state == CircuitState.CLOSED:
                return True

            if self._state == CircuitState.OPEN:
                # 检查是否已过冷却期
                if self._last_failure_time is None:
                    # 防御性：不应该发生（OPEN 态必然有 last_failure_time）
                    self._state = CircuitState.HALF_OPEN
                    self._half_open_probe_in_flight = True
                    return True
                elapsed = time.time() - self._last_failure_time
                if elapsed >= self.recovery_timeout:
                    # 转 HALF_OPEN，放行探测请求
                    self._state = CircuitState.HALF_OPEN
                    self._half_open_probe_in_flight = True
                    return True
                return False  # fast-fail

            if self._state == CircuitState.HALF_OPEN:
                if self._half_open_probe_in_flight:
                    return False  # 探测中，拒绝其他请求
                # 探测已完成（不应到达此分支，但防御性放行）
                self._half_open_probe_in_flight = True
                return True

            return True  # 防御性

    def record_success(self) -> None:
        """记录一次成功调用。

        - CLOSED：重置失败计数
        - HALF_OPEN：关闭断路器，恢复正常
        - OPEN：不应该发生（success 仅在 allow_request 返回 True 后调用）
        """
        with self._lock:
            self._failure_count = 0
            self._half_open_probe_in_flight = False
            if self._state == CircuitState.HALF_OPEN:
                self._state = CircuitState.CLOSED
            # CLOSED 状态保持不变

    def record_failure(self) -> None:
        """记录一次失败调用。

        - CLOSED：失败计数 +1，达阈值转 OPEN
        - HALF_OPEN：失败，转回 OPEN
        - OPEN：刷新 last_failure_time（延长冷却期）
        """
        now = time.time()
        with self._lock:
            self._last_failure_time = now
            self._half_open_probe_in_flight = False

            if self._state == CircuitState.CLOSED:
                self._failure_count += 1
                if self._failure_count >= self.failure_threshold:
                    self._state = CircuitState.OPEN
            elif self._state == CircuitState.HALF_OPEN:
                # 探测失败，重新打开
                self._state = CircuitState.OPEN
                self._failure_count = self.failure_threshold  # 保留以备日志
            elif self._state == CircuitState.OPEN:
                # 已 OPEN，仅刷新 last_failure_time（已在上面赋值）
                pass

    def reset(self) -> None:
        """重置断路器到 CLOSED 初始状态（测试用）。"""
        with self._lock:
            self._state = CircuitState.CLOSED
            self._failure_count = 0
            self._last_failure_time = None
            self._half_open_probe_in_flight = False


class RetryPolicy:
    """指数退避 + jitter 重试策略。

    借鉴 zhikuncode-main RetryPolicy.java（104 行）：
    - BASE_DELAY = 500ms，每次翻倍
    - MAX_DELAY = 32s（封顶）
    - JITTER_RATIO = 0.25，抖动范围 [0, base * 0.25]
    - MAX_ATTEMPTS = 3（总尝试次数，含首次）

    使用示例：

        policy = RetryPolicy()
        for attempt in range(1, policy.max_attempts + 1):
            try:
                return api_call()
            except Exception as e:
                if not policy.should_retry(attempt) or not policy.is_retryable_exception(e):
                    raise
                time.sleep(policy.delay(attempt))
    """

    def __init__(
        self,
        base_delay: float = 0.5,
        max_delay: float = 32.0,
        jitter_ratio: float = 0.25,
        max_attempts: int = 3,
    ):
        """初始化重试策略。

        Args:
            base_delay: 首次重试基础延迟（秒）
            max_delay: 单次重试最大延迟（秒）
            jitter_ratio: 抖动比例 [0, 1]，0 表示无抖动
            max_attempts: 最大尝试次数（含首次调用）
        """
        self.BASE_DELAY = base_delay
        self.MAX_DELAY = max_delay
        self.JITTER_RATIO = jitter_ratio
        self.MAX_ATTEMPTS = max_attempts
        self._random = random.Random()

    # 兼容属性访问（与文档常量命名一致）
    @property
    def base_delay(self) -> float:
        return self.BASE_DELAY

    @property
    def max_delay(self) -> float:
        return self.MAX_DELAY

    @property
    def jitter_ratio(self) -> float:
        return self.JITTER_RATIO

    @property
    def max_attempts(self) -> int:
        return self.MAX_ATTEMPTS

    def delay(self, attempt: int) -> float:
        """计算第 attempt 次重试的延迟时间。

        Args:
            attempt: 当前尝试次数（1 表示首次重试，即第 2 次调用前）

        Returns:
            延迟秒数，含 jitter
        """
        # 指数退避：base * 2^(attempt-1)
        base = min(self.BASE_DELAY * (2 ** (attempt - 1)), self.MAX_DELAY)
        # jitter 在 [0, base * ratio] 范围内
        jitter = self._random.uniform(0, base * self.JITTER_RATIO)
        return base + jitter

    def should_retry(self, attempt: int) -> bool:
        """是否应该继续重试。

        Args:
            attempt: 当前已尝试次数（1 表示首次调用已失败）

        Returns:
            True 表示可以继续重试，False 表示已达上限
        """
        return attempt < self.MAX_ATTEMPTS

    def is_retryable_exception(self, exc: Exception) -> bool:
        """判断异常是否可重试（委托给 RetryableErrorClassifier）。"""
        return RetryableErrorClassifier.is_retryable(exc)


class RetryableErrorClassifier:
    """LLM 错误分类器：区分可重试与不可重试错误。

    借鉴 zhikuncode-main LlmErrorClassifier.java（~150 行）的核心分类逻辑：
    - 429 RateLimitError：可重试（需要 backoff）
    - 500/502/503 InternalServerError：可重试（服务端临时故障）
    - 529 Overloaded：可重试（服务过载，zhikuncode 529 源分类重试）
    - 408 RequestTimeout：可重试
    - 400/401/403/404：不可重试（客户端错误，重试无用）
    - APIConnectionError：可重试（网络抖动）
    - APITimeoutError：可重试（请求超时）
    """

    # 可重试 HTTP 状态码
    RETRYABLE_STATUS_CODES: Set[int] = {
        408,  # Request Timeout
        429,  # Too Many Requests
        500,  # Internal Server Error
        502,  # Bad Gateway
        503,  # Service Unavailable
        529,  # Server Overloaded (Anthropic 风格)
    }

    # 可重试异常类型
    RETRYABLE_EXCEPTION_TYPES = (
        APIConnectionError,
        APITimeoutError,
    )

    @classmethod
    def is_retryable(cls, exc: Exception) -> bool:
        """判断异常是否可重试。

        Args:
            exc: 待判断的异常

        Returns:
            True 表示可重试，False 表示快速失败不重试
        """
        # 1. 检查异常类型
        if isinstance(exc, cls.RETRYABLE_EXCEPTION_TYPES):
            return True

        # 2. 检查 status_code 属性（OpenAI SDK 的 APIStatusError 子类）
        status_code = getattr(exc, "status_code", None)
        if status_code is not None:
            try:
                return int(status_code) in cls.RETRYABLE_STATUS_CODES
            except (TypeError, ValueError):
                pass

        # 3. 显式处理 RateLimitError / InternalServerError（有 status_code）
        if isinstance(exc, (RateLimitError, InternalServerError)):
            return True

        return False


class LlmResilienceGate:
    """LLM 韧性网关：组合断路器 + 重试策略，提供统一的 call() 接口。

    使用示例：

        breaker = ApiCircuitBreaker()
        retry_policy = RetryPolicy()
        gate = LlmResilienceGate(breaker, retry_policy)

        try:
            result = gate.call(lambda: client.chat(messages, ...))
        except CircuitOpenError:
            # 断路器开启，提示用户稍后重试
            ...
        except Exception as e:
            # 重试耗尽后的最终异常
            ...
    """

    def __init__(
        self,
        breaker: ApiCircuitBreaker,
        retry_policy: RetryPolicy,
        classifier: Type[RetryableErrorClassifier] = RetryableErrorClassifier,
    ):
        self._breaker = breaker
        self._retry_policy = retry_policy
        self._classifier = classifier

    def call(self, func: Callable[[], Any]) -> Any:
        """执行函数，应用断路器 + 重试策略。

        Args:
            func: 无参数的可调用对象（已绑定了参数的 LLM 调用）

        Returns:
            函数返回值

        Raises:
            CircuitOpenError: 断路器开启，请求被 fast-fail
            Exception: 重试耗尽后的最终异常（原样抛出）
        """
        attempt = 0
        while True:
            attempt += 1

            # 1. 断路器检查
            if not self._breaker.allow_request():
                raise CircuitOpenError()

            # 2. 执行函数
            try:
                result = func()
                self._breaker.record_success()
                return result
            except Exception as exc:
                # 3. 判断是否可重试
                if not self._classifier.is_retryable(exc):
                    # 不可重试异常：不记录为断路器失败（避免污染计数）
                    # 但仍向上抛出
                    raise

                # 4. 可重试异常：记录断路器失败
                self._breaker.record_failure()

                # 5. 检查重试次数
                if not self._retry_policy.should_retry(attempt):
                    raise  # 重试耗尽，抛出最后异常

                # 6. 等待退避时间
                time.sleep(self._retry_policy.delay(attempt))
                # 循环继续下一次尝试
