"""HealthAwareSelector 单测(M2.5 设计 §3.2): EWMA×在途 P2C,地板探索。"""

import pytest

from polygateway.ports import OutcomeAwareSelector
from polygateway.sources import HealthAwareSelector, RoundRobinSelector
from polygateway.types import SourceConfig, SourceStats


def _src(name: str) -> SourceConfig:
    return SourceConfig(
        name=name,
        provider="p",
        base_url="https://gw.example/v1",
        api_key="sk-x",
        model="m",
        timeout_s=60.0,
    )


def _stats(**inflight: int) -> dict[str, SourceStats]:
    return {name: SourceStats(inflight=n, rpm_used=0, tpm_used=0) for name, n in inflight.items()}


class TestHealthAwareSelector:
    def test_implements_outcome_protocol(self):
        assert isinstance(HealthAwareSelector(), OutcomeAwareSelector)
        assert not isinstance(RoundRobinSelector(), OutcomeAwareSelector)

    def test_unhealthy_source_demoted(self):
        # s2 连续失败 → EWMA 塌陷 → 排序永远在健康源之后(rng 定值消除 P2C 随机性)
        sel = HealthAwareSelector(rng=lambda: 0.0)
        sources = [_src("s1"), _src("s2")]
        for _ in range(10):
            sel.record_outcome("s2", ok=False)
        order = sel.order(sources, _stats(s1=0, s2=0))
        assert [s.name for s in order] == ["s1", "s2"]

    def test_floor_keeps_exploration_possible(self):
        # 地板 0.05: 塌陷源分数不归零——EWMA 十次失败后仍 > 0,P2C 采样到时可胜平局
        sel = HealthAwareSelector(rng=lambda: 0.0)
        for _ in range(50):
            sel.record_outcome("s2", ok=False)
        assert sel._score("s2", _stats(s2=0)) == pytest.approx(0.05)

    def test_ewma_climbs_back_on_recovery(self):
        # 恢复源连续成功,EWMA α=0.2 自然爬升(即天然 slow-start)
        sel = HealthAwareSelector(rng=lambda: 0.0)
        for _ in range(50):
            sel.record_outcome("s2", ok=False)
        low = sel._score("s2", _stats(s2=0))
        for _ in range(10):
            sel.record_outcome("s2", ok=True)
        high = sel._score("s2", _stats(s2=0))
        assert high > 0.85 > low

    def test_inflight_suppresses_score(self):
        sel = HealthAwareSelector(rng=lambda: 0.0)
        assert sel._score("s1", _stats(s1=0)) == pytest.approx(1.0)
        assert sel._score("s1", _stats(s1=3)) == pytest.approx(0.25)

    def test_p2c_head_randomized_rest_by_score(self):
        # rng 驱动 P2C 取样: 三源同分时头名由 rng 决定;其余按分数降序
        sources = [_src("s1"), _src("s2"), _src("s3")]
        sel = HealthAwareSelector(rng=iter([0.9, 0.0]).__next__)  # 采样 s3 与 s1 比分
        for _ in range(5):
            sel.record_outcome("s3", ok=False)  # s3 塌陷
        order = sel.order(sources, _stats(s1=0, s2=0, s3=0))
        assert order[0].name == "s1"  # 两候选中 s1 胜出
        assert order[-1].name == "s3"  # 塌陷源垫底

    def test_missing_stats_defaults_to_zero_inflight(self):
        sel = HealthAwareSelector(rng=lambda: 0.0)
        assert [s.name for s in sel.order([_src("s1")], {})] == ["s1"]