"""RetryMW: 重试循环拥有"尝试"的全部编排(D13 自研;蓝本 CHS governance.py:107-268)。

每次尝试 = 选源(跳过冷却源)→ 该源熔断门 → 限流 permit → transport;
换源、退避、逐次遥测、熔断写回、permit 结算全部在循环内。作为洋葱的
终端(compose 的 terminal)被调用: `await retry_mw(request)`。

例外说明(P7): 本模块 import httpx 仅用于失败原因归类(CHS 同款判定),
httpx 是库的核心依赖而非实现层内部件,不违反"middleware 只依赖端口"。
"""

from __future__ import annotations

import asyncio
import random
import time
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING

import httpx
from loguru import logger

from polygateway.errors import (
    AllSourcesExhausted,
    CircuitOpenError,
    PolyGatewayError,
    RequestRejectedError,
    ResultInvalidError,
    SourceDeadError,
    TransientError,
)
from polygateway.middleware.breaker import BreakerGate
from polygateway.middleware.ratelimit import QuotaGate
from polygateway.sources import SourceCooldownMemo
from polygateway.streaming import StreamLivenessTimeout
from polygateway.types import LLMResponse

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    from polygateway.ports import (
        GateDecision,
        Permit,
        ProviderGate,
        RateLimiter,
        SourceSelector,
        Transport,
    )
    from polygateway.types import (
        BackpressurePolicy,
        ChatRequest,
        RetryPolicy,
        SourceConfig,
        TransportResult,
    )


def _failure_reason(exc: PolyGatewayError) -> str:
    """失败原因归类(CHS governance.py:169 同款)。"""
    if isinstance(exc, SourceDeadError):
        return "source_dead"
    if exc.status_code == 429:
        return "rate_limited"
    if isinstance(exc.__cause__, httpx.TimeoutException | StreamLivenessTimeout):
        return "timeout"
    return "network_error"


@dataclass(frozen=True)
class _Failed:
    """一次可重试失败的结果(SourceDead 立即换源,Transient 退避)。"""

    exc: PolyGatewayError
    immediate: bool


class RetryMW:
    """尝试编排器;时钟/睡眠/随机全部注入,纯确定性可测(P6)。"""

    def __init__(
        self,
        *,
        scope: str,
        sources: list[SourceConfig],
        selector: SourceSelector,
        limiter: RateLimiter,
        gate: ProviderGate,
        transport: Transport,
        retry: RetryPolicy,
        backpressure: BackpressurePolicy,
        quota_full: str = "wait",
        cooldown_memo: SourceCooldownMemo | None = None,
        emitter: object | None = None,
        now: Callable[[], float] = time.monotonic,
        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
        rng: Callable[[], float] = random.random,
    ) -> None:
        if quota_full not in ("wait", "fail_fast"):
            raise ValueError(f"quota_full 必须是 wait|fail_fast: {quota_full!r}")
        self._scope = scope
        self._sources = list(sources)
        self._selector = selector
        self._quota = QuotaGate(limiter)
        self._breaker = BreakerGate(gate)
        self._transport = transport
        self._retry = retry
        self._bp = backpressure
        self._quota_full = quota_full
        self._memo = cooldown_memo or SourceCooldownMemo(now=now)
        self._emitter = emitter
        self._now = now
        self._sleep = sleep
        self._rng = rng

    async def __call__(self, request: ChatRequest) -> LLMResponse:
        """执行治理调用;scope 级失败按 §6.1 携结构化字段上抛。"""
        if not self._sources:
            raise AllSourcesExhausted(scope=self._scope, reason="no_sources", retry_after_s=0.0)
        fails = 0
        reasons: dict[str, str] = {}
        while True:
            picked, gate_rejections = await self._pick_runnable(reasons)
            if picked is None:
                await self._on_no_runnable(gate_rejections, reasons)
                continue
            outcome = await self._attempt(request, *picked, reasons)
            if isinstance(outcome, LLMResponse):
                return outcome
            fails += 1
            if fails >= self._retry.max_attempts:
                raise AllSourcesExhausted(
                    scope=self._scope,
                    reason="retry_exhausted",
                    retry_after_s=self._retry.backoff_base_s,
                    per_source_reasons=reasons,
                ) from outcome.exc
            if not outcome.immediate:
                await self._sleep(self._backoff_delay(fails, outcome.exc))

    # —— 选源与准入(CHS _pick_runnable 120-167)——

    async def _pick_runnable(
        self, reasons: dict[str, str]
    ) -> tuple[tuple[SourceConfig, Permit, GateDecision] | None, int]:
        stats = {s.name: await self._quota.stats(s) for s in self._sources}
        gate_rejections = 0
        for cand in self._selector.order(self._sources, stats):
            if self._memo.active(cand.name):
                # 冷却备忘跳过也计入拒绝数,保住 circuit_open 判据(CHS 同款)
                gate_rejections += 1
                reasons[cand.name] = "cooldown"
                continue
            permit = await self._quota.try_acquire(cand)
            if permit is None:
                reasons.setdefault(cand.name, "rate_limited")
                continue
            entry = None
            try:
                entry = await self._breaker.try_enter(cand, uuid.uuid4().hex)
            finally:
                # try_enter 未归还 entry(异常/取消)→ 释放已占 permit,不吞任何异常
                if entry is None:
                    await self._settle_and_release(permit, 0)
            if entry.allowed:
                return (cand, permit, entry), gate_rejections
            gate_rejections += 1
            reasons[cand.name] = "circuit_open"
            # 开路源本地记冷却,避免每轮白烧 RPM 探测(CHS governance.py:107)
            self._memo.set_until(cand.name, self._now() + entry.retry_after_s)
            await self._settle_and_release(permit, 0)
        return None, gate_rejections

    async def _on_no_runnable(self, gate_rejections: int, reasons: dict[str, str]) -> None:
        if gate_rejections == len(self._sources):
            names = tuple(s.name for s in self._sources)
            raise CircuitOpenError(
                scope=self._scope,
                retry_after_s=await self._breaker.retry_after_s(names),
                per_source_reasons=reasons,
            )
        if self._quota_full == "fail_fast":
            raise AllSourcesExhausted(
                scope=self._scope,
                reason="quota_exhausted",
                retry_after_s=self._bp.poll_interval_s,
                per_source_reasons=reasons,
            )
        await self._sleep(self._bp.poll_interval_s)

    # —— 单次尝试(CHS run 200-268)——

    async def _attempt(
        self,
        request: ChatRequest,
        source: SourceConfig,
        permit: Permit,
        entry: GateDecision,
        reasons: dict[str, str],
    ) -> LLMResponse | _Failed:
        call_id = str(uuid.uuid4())
        started = self._now()
        actual = 0
        try:
            result = await self._transport.complete(
                messages=request.messages,
                source=source,
                stream=request.stream,
                overlay=request.overlay,
                call_id=call_id,
            )
            actual = result.prompt_tokens + result.completion_tokens
            await self._breaker.record_success(entry)
            await self._quota.mark_progress()
            response = self._build_response(source, result, call_id, started)
            await self._emit(request, source, call_id, started, response=response)
            return response
        except RequestRejectedError as exc:
            await self._on_rejected(exc, source, entry)
            await self._emit(request, source, call_id, started, error=exc)
            raise
        except ResultInvalidError as exc:
            # 坏结果 ≠ 坏服务: 熔断记成功,异常上抛消耗业务失败预算(§6.3)
            await self._breaker.record_success(entry)
            await self._emit(request, source, call_id, started, error=exc)
            raise
        except asyncio.CancelledError:
            if entry.is_probe:
                await self._breaker.release_probe(entry)
            await self._emit(request, source, call_id, started, error="cancelled")
            raise
        except (SourceDeadError, TransientError) as exc:
            dead = isinstance(exc, SourceDeadError)
            reason = _failure_reason(exc)
            reasons[source.name] = reason
            await self._breaker.record_failure(entry, reason, dead)
            if not dead:
                actual = source.est_tokens  # 保守: 失败请求可能已被网关计费(CHS 同款)
            await self._emit(request, source, call_id, started, error=exc)
            return _Failed(exc, immediate=dead)
        finally:
            await self._settle_and_release(permit, actual)

    async def _on_rejected(
        self, exc: RequestRejectedError, source: SourceConfig, entry: GateDecision
    ) -> None:
        provider_responded = exc.source_name == source.name and exc.status_code is not None
        if provider_responded:
            await self._breaker.record_success(entry)  # 网关健康地拒了坏请求
        elif entry.is_probe:
            await self._breaker.release_probe(entry)

    # —— 辅助 ——

    def _backoff_delay(self, fails: int, exc: PolyGatewayError) -> float:
        """指数退避+jitter,与 Retry-After 提示取大(ARCH §7.2;VT jitter 系数)。"""
        base = min(self._retry.backoff_base_s * (2 ** (fails - 1)), self._retry.backoff_max_s)
        delay = base * (0.5 + self._rng())
        retry_after = getattr(exc, "retry_after_s", None) or 0.0
        return max(delay, retry_after)

    def _build_response(
        self, source: SourceConfig, result: TransportResult, call_id: str, started: float
    ) -> LLMResponse:
        return LLMResponse(
            content=result.content,
            thinking=result.thinking,
            model=source.model,
            provider=source.provider,
            prompt_tokens=result.prompt_tokens,
            completion_tokens=result.completion_tokens,
            latency_ms=int((self._now() - started) * 1000),
            ttft_ms=result.ttft_ms,
            max_inter_token_ms=result.max_inter_token_ms,
            cache_hit=False,
            call_id=call_id,
            source_name=source.name,
            cost=None,
            usage_source=result.usage_source,
        )

    async def _settle_and_release(self, permit: Permit, actual: int) -> None:
        """finally 专用: settle 后必 release;失败降级 warning,绝不掩盖主异常/取消。"""
        try:
            try:
                await permit.settle(actual)
            finally:
                await permit.release()
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            logger.warning("permit 结算/释放失败(不掩盖主异常): {}", exc)

    async def _emit(
        self,
        request: ChatRequest,
        source: SourceConfig,
        call_id: str,
        started: float,
        *,
        response: LLMResponse | None = None,
        error: object | None = None,
    ) -> None:
        """逐次遥测(经注入的单一 Emitter);遥测失败不得影响调用(铁律)。"""
        if self._emitter is None:
            return
        try:
            await self._emitter.emit_attempt(
                request=request,
                source=source,
                call_id=call_id,
                latency_ms=int((self._now() - started) * 1000),
                response=response,
                error=None if error is None else str(error),
            )
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            logger.warning("逐次遥测记录失败(降级不冒泡): {}", exc)