"""SourceAdmission: 一次尝试的准入编排,三条治理循环(chat/embedding/ocr)共用一份。

**收敛缘由(issue #14)**: 本模块的两个方法此前在 `middleware/retry.py`、
`embedding.py`、`ocr.py` 各存一份逐字复制(后两份是第一份的子集)。准入语义
一直在演进——issue #8 改过 stall 口径、M2.5 加过 AIMD pacer、issue #14 要加
熔断等待档——每演进一次就要三处同步,漏一处即行为分叉。三份复制正是库铁律
痛斥的那种模式(遥测"三项目 4 处复制"的教训),只不过这次发生在库内部。

**职责边界**: 只管"挑出一个可跑的源"与"一个都挑不出来时怎么办";一次尝试
本身(transport 调用、记账写回、逐次遥测)仍归各循环的 `_attempt`。

**共享而非持有**: `QuotaGate`/`BreakerGate`/`AdaptivePacer`/`SourceSelector` 由
调用方构造后传入**同一实例**——三处 `_attempt` 仍要用它们做记账写回与
`pacer.leave()`。pacer 尤其不能各建一个: 它有在途计数,分裂成两个计数器会让
`admit`/`enter` 与 `leave` 记到不同账上。`SourceCooldownMemo` 只被准入消费,
由本类独占。
"""

from __future__ import annotations

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

from loguru import logger

from polygateway.errors import AllSourcesExhausted, CircuitOpenError
from polygateway.sources import SourceCooldownMemo

# 两个准入策略键共用的值域;校验只此一处,不在各客户端重复
_POLICIES = frozenset({"wait", "fail_fast"})

if TYPE_CHECKING:
    from collections.abc import Callable

    from polygateway.middleware.breaker import BreakerGate
    from polygateway.middleware.ratelimit import QuotaGate
    from polygateway.middleware.retry import StallClock
    from polygateway.ports import GateDecision, Permit, SourceSelector
    from polygateway.sources import AdaptivePacer
    from polygateway.types import BackpressurePolicy, SourceConfig


async def settle_and_release(permit: Permit, actual: int) -> None:
    """finally 专用: settle 后必 release;失败降级 warning,绝不掩盖主异常/取消。

    三条循环的 `_attempt` 与本模块的准入拒绝路径共用这一份(此前三处逐字复制,
    仅 warning 文案不同)。
    """
    try:
        try:
            await permit.settle(actual)
        finally:
            await permit.release()
    except asyncio.CancelledError:
        raise
    except Exception as exc:
        logger.warning("permit 结算/释放失败(不掩盖主异常): {}", exc)


def _demote_call_failures(
    ordered: list[SourceConfig],
    attempt_fails: dict[str, int],
    health: Callable[[str], float] | None,
) -> list[SourceConfig]:
    """调用内降权(设计 §3.3/§3.36): 失败 ≥2 次且存在可信替代才让位。

    可信替代 = 某未失败候选 health ≥ 0.5 × 失败源 health——异构池里健康源
    偶发失败不该被推向已知坏源(第三轮教训: 期望成功率 83% vs 10%)。
    无健康视图(round_robin 等)保持无条件降权(冷启动保护)。

    `attempt_fails` 为空时恒等返回原列表对象——embedding/ocr 不维护调用内
    失败计数,故对它们这一步是零成本的空操作,无需在调用侧加分支。
    """
    demoted = [s for s in ordered if attempt_fails.get(s.name, 0) >= 2]
    if not demoted or len(demoted) == len(ordered):
        return ordered
    if health is None:
        return _move_to_tail(ordered, demoted)
    return _health_gated_reorder(ordered, demoted, attempt_fails, health)


def _move_to_tail(ordered: list[SourceConfig], demoted: list[SourceConfig]) -> list[SourceConfig]:
    """无健康视图: 无条件移尾(冷启动保护原语义)。"""
    names = {d.name for d in demoted}
    return [s for s in ordered if s.name not in names] + demoted


def _health_gated_reorder(
    ordered: list[SourceConfig],
    demoted: list[SourceConfig],
    attempt_fails: dict[str, int],
    health: Callable[[str], float],
) -> list[SourceConfig]:
    """健康门槛降权: 无可信替代则原地重试;有则插到可信替代之后。"""
    demoted = _credible_demotions(ordered, demoted, attempt_fails, health)
    if not demoted:
        return ordered
    names = {d.name for d in demoted}
    rest = [s for s in ordered if s.name not in names]
    return _insert_after_credible(rest, demoted, health)


def _insert_after_credible(
    rest: list[SourceConfig],
    demoted: list[SourceConfig],
    health: Callable[[str], float],
) -> list[SourceConfig]:
    """插入位置(第四轮教训): 被降权源排在可信替代之后、不可信源之前——
    可信替代被限流闸/熔断跳过时,下一候选是失败源本身而非垃圾源。"""
    bar = 0.5 * max(health(d.name) for d in demoted)
    credible = [s for s in rest if health(s.name) >= bar]
    junk = [s for s in rest if health(s.name) < bar]
    return credible + demoted + junk


def _credible_demotions(
    ordered: list[SourceConfig],
    demoted: list[SourceConfig],
    attempt_fails: dict[str, int],
    health: Callable[[str], float],
) -> list[SourceConfig]:
    """健康门槛过滤: 仅当存在"健康分 ≥ 失败源一半"的未失败候选,让位才有意义。"""
    alts = [o for o in ordered if attempt_fails.get(o.name, 0) < 2]
    return [s for s in demoted if any(health(o.name) >= 0.5 * health(s.name) for o in alts)]


class SourceAdmission:
    """准入编排器(CHS `governance.py:107-285` 同款);时钟/睡眠/随机全部注入。"""

    def __init__(
        self,
        *,
        scope: str,
        sources: list[SourceConfig],
        selector: SourceSelector,
        quota: QuotaGate,
        breaker: BreakerGate,
        backpressure: BackpressurePolicy,
        quota_full: str,
        circuit_open: str,
        memo: SourceCooldownMemo | None = None,
        pacer: AdaptivePacer | None = None,
        health_view: Callable[[str], float] | None = None,
        now: Callable[[], float] = time.monotonic,
        sleep: Callable[[float], object] = asyncio.sleep,
        rng: Callable[[], float] = random.random,
    ) -> None:
        for name, value in (("quota_full", quota_full), ("circuit_open", circuit_open)):
            if value not in _POLICIES:
                raise ValueError(f"{name} 必须是 wait|fail_fast: {value!r}")
        self._scope = scope
        self._sources = sources
        self._selector = selector
        self._quota = quota
        self._breaker = breaker
        self._bp = backpressure
        self._quota_full = quota_full
        self._circuit_open = circuit_open
        self._memo = memo or SourceCooldownMemo(now=now)
        self._pacer = pacer
        self._health_view = health_view
        self._now = now
        self._sleep = sleep
        self._rng = rng

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

    async def pick(
        self, reasons: dict[str, str], attempt_fails: dict[str, int]
    ) -> tuple[tuple[SourceConfig, Permit, GateDecision] | None, int]:
        """挑出第一个过闸的候选;返回 (选中三元组 | None, 熔断类拒绝计数)。"""
        stats = {s.name: await self._quota.stats(s) for s in self._sources}
        gate_rejections = 0
        ordered = _demote_call_failures(
            self._selector.order(self._sources, stats), attempt_fails, self._health_view
        )
        for cand in ordered:
            if self._memo.active(cand.name):
                # 冷却备忘跳过也计入拒绝数,保住 circuit_open 判据(CHS 同款)
                gate_rejections += 1
                reasons[cand.name] = "cooldown"
                continue
            if self._pacer is not None and not self._pacer.admit(cand.name):
                # AIMD 超限: 不计 gate_rejections → 走 quota-wait 排队,不误判熔断
                reasons.setdefault(cand.name, "adaptive_paced")
                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 settle_and_release(permit, 0)
            if entry.allowed:
                if self._pacer is not None:
                    self._pacer.enter(cand.name)
                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 settle_and_release(permit, 0)
        return None, gate_rejections

    # —— 背压与 stall 判死(CHS governance.py:270-285)——

    async def stalled(self, clock: StallClock) -> bool:
        """双条件 stall 判死(CHS governance.py:270-281): 本地累计等待与全局
        无进展**同时**超窗才判死——本地 monotonic 与后端时钟刻意不混用。

        本地一侧只计非生产性等待(issue #8,见 `StallClock`)。短路顺序有意为之:
        本地未超窗就不问后端,省一次 Redis 往返。
        """
        stall = self._bp.stall_window_s
        return clock.stalled_s() > stall and await self._quota.progress_age_s() > stall

    async def on_no_runnable(
        self, gate_rejections: int, reasons: dict[str, str], clock: StallClock
    ) -> None:
        """一个源都挑不出来时的处置: **按拒绝原因分派**到各自的策略。

        分派而非串行是硬要求(issue #14): 串行写法下 `circuit_open=wait` 不抛
        之后会径直掉进配额分支,`quota_full=fail_fast` 的调用方于是收到一个
        `reason=quota_exhausted` 的异常——而配额其实是满的,坏的是熔断门。
        """
        names = tuple(s.name for s in self._sources)
        if gate_rejections == len(self._sources):
            # 全部因熔断类原因(门开路 / 本地冷却备忘)被拒
            if self._circuit_open == "fail_fast":
                raise CircuitOpenError(
                    scope=self._scope,
                    retry_after_s=await self._breaker.retry_after_s(names),
                    per_source_reasons=reasons,
                )
            # wait: 保护作用完整保留(这一轮照样一个请求都不发),改变的只是
            # 调用方当场死还是排队等——多源可换源故 fail-fast 对,单源无源可换
            hint = await self._breaker.retry_after_s(names)
        else:
            # 至少一个源是被配额/AIMD 挡的,归 quota_full 管
            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,
                )
            hint = 0.0
        if await self.stalled(clock):
            raise AllSourcesExhausted(
                scope=self._scope,
                reason="stalled",
                retry_after_s=await self._breaker.retry_after_s(names),
                per_source_reasons=reasons,
            )
        nap = self._nap(hint, clock)
        if hint > 0:
            logger.info("熔断开路等待 {:.1f}s 后重试(scope={}, 原因={})", nap, self._scope, reasons)
        await self._sleep(nap)

    def _nap(self, hint: float, clock: StallClock) -> float:
        """本轮等待多久。**必须在 `stalled()` 判定之后调用**(预算可能已耗尽)。

        `hint > 0`(熔断开路有确定的冷却截止)时睡到那个时刻,而不是按
        `poll_interval` 空转——60 秒冷却用 10ms 轮询是 6000 次空转,内存后端
        只是查字典,Redis 后端则是 6000 次往返 × 每个在途调用。抖动**上**加
        而非缩放(既有 quota 路径是 `[0.5p, 1.0p]`): 对一个确定的截止时刻提前
        醒来必然被再拒一次,白跑一趟。

        两档都夹到剩余 stall 预算,故单次调用的最坏墙钟是 `stall_window_s`
        加一个 poll 间隔,不随 `max_cooldown_s` 漂移。多加的那一格是因为
        `stalled()` 判据是 `>` 而非 `>=`——恰好睡到窗口边界不判死,留这一格
        让下一轮必定判死。`hint == 0` 时整个式子退化为既有的 jitter 轮询。
        """
        jitter = self._bp.poll_interval_s * (0.5 + 0.5 * self._rng())
        budget = self._bp.stall_window_s - clock.stalled_s() + self._bp.poll_interval_s
        wait = hint + jitter if hint > 0 else jitter
        # 下界取 jitter 而非 poll_interval: 既有 quota 轮询是 [0.5p, 1.0p],用
        # poll_interval 兜底会把 rng→0 那半边抬上去。预算为负时(本地已超窗但
        # 全局仍在出餐,故 stalled() 不判死)靠它退回正常轮询节奏,不忙循环。
        return max(jitter, min(wait, budget))
