"""OcrClient 治理循环测试(M3 设计 §5): 换源/熔断口径/stall/取消/G1 契约。

循环与 EmbeddingClient 同构(设计 §2.A 有限重复裁决);本文件的
retry_exhausted/circuit_open/stalled 三组断言即设计 §6 ③ 的 G1 契约钉。
"""

import asyncio

import pytest

from polygateway.backends.memory.breaker import InMemoryGate
from polygateway.backends.memory.limiter import InMemoryLimiter
from polygateway.errors import (
    AllSourcesExhausted,
    CircuitOpenError,
    RequestRejectedError,
    ResultInvalidError,
    SourceDeadError,
    TransientError,
)
from polygateway.ocr import OcrClient
from polygateway.types import (
    BackpressurePolicy,
    BreakerConfig,
    GlobalLimits,
    OcrLayoutElement,
    OcrLayoutTransportResult,
    OcrTextTransportResult,
    RetryPolicy,
    SourceConfig,
)

_BREAKER = BreakerConfig(fail_threshold=3, cooldown_s=60.0, probe_ttl_s=120.0)
_NO_GLOBAL = GlobalLimits(max_concurrency=0, rpm=0, tpm=0)
_TEXT_OK = OcrTextTransportResult(text="LINE-1", raw={"task_type": "text"})
_LAYOUT_OK = OcrLayoutTransportResult(
    elements=[OcrLayoutElement(type="table", bbox=(41.0, 48.0, 218.0, 282.0), page_index=0)],
    page_sizes=[(759.0, 540.0)],
    raw={"success": True},
)


def _src(**overrides):
    base = {
        "name": "m1",
        "provider": "monkey",
        "base_url": "http://10.77.0.20:7866",
        "api_key": "none",
        "model": "monkey-ocr",
        "timeout_s": 120.0,
    }
    base.update(overrides)
    return SourceConfig(**base)


class ScriptedOcrTransport:
    """按脚本响应: Exception / "text" / "layout" / "hang";两方法共用一份脚本。"""

    def __init__(self, script):
        self.script = list(script)
        self.calls = []

    async def _next(self, method, source, call_id):
        self.calls.append((method, source.name, call_id))
        action = self.script.pop(0)
        if isinstance(action, Exception):
            raise action
        if action == "hang":
            await asyncio.Event().wait()
        return _TEXT_OK if action == "text" else _LAYOUT_OK

    async def recognize_text(self, *, image, source, call_id):
        return await self._next("text", source, call_id)

    async def parse_layout(self, *, image, source, call_id):
        return await self._next("layout", source, call_id)

    async def check_health(self, *, source):
        raise NotImplementedError


class StaticSelector:
    def order(self, sources, stats):
        return list(sources)


class RecordingSelector(StaticSelector):
    def __init__(self):
        self.outcomes = []

    def record_outcome(self, source_name, ok):
        self.outcomes.append((source_name, ok))

    def health(self, source_name):
        return 1.0


class RecordingGate:
    """InMemoryGate 包装: 录 record_success 的 count_attempt 与 release_probe。"""

    def __init__(self, inner):
        self._inner = inner
        self.successes = []
        self.probe_releases = 0

    async def try_enter(self, source_name, owner):
        return await self._inner.try_enter(source_name, owner)

    async def record_success(self, entry, *, count_attempt=True):
        self.successes.append((entry.source_name, count_attempt))
        return await self._inner.record_success(entry, count_attempt=count_attempt)

    async def record_failure(self, entry, reason, force_open):
        return await self._inner.record_failure(entry, reason, force_open)

    async def release_probe(self, entry):
        self.probe_releases += 1
        return await self._inner.release_probe(entry)

    async def retry_after_s(self, sources):
        return await self._inner.retry_after_s(sources)


class _MemoryRecorder:
    def __init__(self):
        self.rows = []

    async def record_llm_call(self, **fields):
        self.rows.append(fields)


def _client(sources, script, **overrides):
    limiter = InMemoryLimiter(
        scope="ocr",
        sources={s.name: s for s in sources},
        global_limits=_NO_GLOBAL,
        lease_ttl_s=100.0,
    )
    gate = RecordingGate(InMemoryGate(config=_BREAKER))
    kwargs = {
        "scope": "ocr",
        "sources": sources,
        "selector": RecordingSelector(),
        "limiter": limiter,
        "breaker": gate,
        "transport": ScriptedOcrTransport(script),
        "retry": RetryPolicy(max_attempts=3, backoff_base_s=0.001, backoff_max_s=0.01),
        "backpressure": BackpressurePolicy(stall_window_s=300.0, poll_interval_s=0.001),
    }
    kwargs.update(overrides)
    return OcrClient(**kwargs), limiter, gate


class TestSuccessPaths:
    async def test_text_result_fields(self):
        client, _, gate = _client([_src()], ["text"])
        r = await client.recognize_text(b"jpg")
        assert r.text == "LINE-1"
        assert r.source_name == "m1"
        assert r.usage.prompt_tokens == 0 and r.usage.completion_tokens == 0
        assert r.latency_ms >= 0 and r.call_id and r.raw["task_type"] == "text"
        assert gate.successes == [("m1", True)]
        assert client._selector.outcomes == [("m1", True)]

    async def test_layout_result_fields(self):
        client, _, _ = _client([_src()], ["layout"])
        r = await client.parse_layout(b"jpg")
        assert r.elements[0].type == "table" and r.page_sizes == [(759.0, 540.0)]

    async def test_input_validation(self):
        client, _, _ = _client([_src()], [])
        with pytest.raises(TypeError):
            await client.recognize_text("not-bytes")
        with pytest.raises(ValueError):
            await client.parse_layout(b"")


class TestFailover:
    async def test_transient_retries_with_backoff(self):
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        client, _, _ = _client(
            [_src()], [TransientError("boom", status_code=500), "text"], sleep=fake_sleep
        )
        r = await client.recognize_text(b"jpg")
        assert r.text == "LINE-1"
        assert any(s > 0 for s in sleeps)  # Transient 退避后重试
        assert client._selector.outcomes == [("m1", False), ("m1", True)]

    async def test_source_dead_switches_immediately(self):
        sleeps = []

        async def fake_sleep(delay):
            sleeps.append(delay)

        s1, s2 = _src(name="m1"), _src(name="m2", base_url="http://10.77.0.20:7867")
        client, _, _ = _client(
            [s1, s2], [SourceDeadError("401", status_code=401), "text"], sleep=fake_sleep
        )
        r = await client.recognize_text(b"jpg")
        assert r.source_name == "m2"
        assert not [s for s in sleeps if s > 0.0005]  # dead 立即换源不退避

    async def test_retry_exhausted_carries_g1_fields(self):
        client, _, _ = _client(
            [_src()],
            [TransientError("1", status_code=500)] * 3,
        )
        with pytest.raises(AllSourcesExhausted) as ei:
            await client.recognize_text(b"jpg")
        exc = ei.value
        assert exc.reason == "retry_exhausted"
        assert exc.per_source_reasons == {"m1": "network_error"}  # G1: 逐源原因
        assert exc.retry_after_s > 0  # G1: 可延期重投

    async def test_circuit_open_carries_g1_fields(self):
        client, _, _ = _client([_src()], [SourceDeadError("401", status_code=401)])
        with pytest.raises(CircuitOpenError) as ei:
            await client.recognize_text(b"jpg")
        exc = ei.value
        assert exc.per_source_reasons["m1"] == "circuit_open"
        assert exc.retry_after_s > 0


class TestTerminalOutcomes:
    async def test_result_invalid_counts_no_attempt_and_no_feed(self):
        client, _, gate = _client([_src()], [ResultInvalidError("bad zip")])
        with pytest.raises(ResultInvalidError):
            await client.parse_layout(b"jpg")
        assert gate.successes == [("m1", False)]  # 记成功但不入失败率窗
        assert client._selector.outcomes == []  # 坏结果 ≠ 坏服务,不喂健康

    async def test_rejected_with_status_counts_no_attempt(self):
        client, _, gate = _client([_src()], [RequestRejectedError("parse failed", status_code=200)])
        with pytest.raises(RequestRejectedError):
            await client.parse_layout(b"jpg")
        assert gate.successes == [("m1", False)]

    async def test_rejected_without_status_no_gate_success(self):
        client, _, gate = _client([_src()], [RequestRejectedError("local refuse")])
        with pytest.raises(RequestRejectedError):
            await client.recognize_text(b"jpg")
        assert gate.successes == []  # 服务未响应: 不记成功(非探针也不归还)


class TestBackpressure:
    async def test_fail_fast_when_quota_full(self):
        src = _src(max_concurrency=1)
        client, limiter, _ = _client([src], ["text"], quota_full="fail_fast")
        permit = await limiter.acquire("m1", 0)  # 占满唯一并发位
        try:
            with pytest.raises(AllSourcesExhausted) as ei:
                await client.recognize_text(b"jpg")
            assert ei.value.reason == "quota_exhausted"
        finally:
            await permit.settle(0)
            await permit.release()

    async def test_wait_until_permit_freed(self):
        src = _src(max_concurrency=1)
        client, limiter, _ = _client([src], ["text"])
        permit = await limiter.acquire("m1", 0)

        async def free_later():
            await asyncio.sleep(0.01)
            await permit.settle(0)
            await permit.release()

        release_task = asyncio.create_task(free_later())
        r = await client.recognize_text(b"jpg")
        await release_task
        assert r.text == "LINE-1"  # quota_full=wait: 等到许可释放而非报错

    async def test_stalled_when_no_progress(self):
        src = _src(max_concurrency=1)
        client, limiter, _ = _client(
            [src],
            ["text"],
            backpressure=BackpressurePolicy(stall_window_s=0.01, poll_interval_s=0.001),
        )
        permit = await limiter.acquire("m1", 0)  # 永不释放且全局无进展
        try:
            with pytest.raises(AllSourcesExhausted) as ei:
                await client.recognize_text(b"jpg")
            assert ei.value.reason == "stalled"
            assert ei.value.per_source_reasons["m1"] == "rate_limited"
        finally:
            await permit.settle(0)
            await permit.release()


class FakeClock:
    def __init__(self, start=1000.0):
        self.t = start

    def __call__(self):
        return self.t

    def advance(self, seconds):
        self.t += seconds


class TestCancellation:
    async def test_cancel_during_probe_returns_probe(self):
        # 探针取消归还(铁律分支;verifier I3): 开路→冷却过后半开探针→
        # transport 挂起中取消 → release_probe 必须发生,后续可再探
        clock = FakeClock()
        inner = InMemoryGate(config=_BREAKER, now=clock)  # 门与 client 共用注入钟
        gate = RecordingGate(inner)
        client, _, _ = _client([_src()], ["hang"], now=clock, breaker=gate)
        entry = await inner.try_enter("m1", "setup")
        await inner.record_failure(entry, "source_dead", True)  # force_open → OPEN
        clock.advance(61.0)  # 越过 cooldown_s=60 → HALF_OPEN
        task = asyncio.create_task(client.recognize_text(b"jpg"))
        await asyncio.sleep(0.05)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        assert gate.probe_releases >= 1  # 探针已归还,不悬挂
        probe = await inner.try_enter("m1", "again")
        assert probe.allowed and probe.is_probe  # 可再探 = 未悬挂的行为证据

    async def test_cancel_during_backoff_sleep_propagates(self):
        async def cancelling_sleep(delay):
            raise asyncio.CancelledError()

        client, limiter, _ = _client(
            [_src()], [TransientError("boom", status_code=500)], sleep=cancelling_sleep
        )
        with pytest.raises(asyncio.CancelledError):
            await client.recognize_text(b"jpg")
        stats = await limiter.source_stats("m1")
        assert stats.inflight == 0  # 退避期取消: permit 已在 finally 释放

    async def test_cancel_during_transport_releases_permit(self):
        src = _src(max_concurrency=1)
        client, limiter, _ = _client([src], ["hang"])
        task = asyncio.create_task(client.recognize_text(b"jpg"))
        await asyncio.sleep(0.01)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task
        stats = await limiter.source_stats("m1")
        assert stats.inflight == 0  # permit 在 finally 释放


class TestCheckHealth:
    class _HealthTransport(ScriptedOcrTransport):
        def __init__(self, mapping):
            super().__init__([])
            self.mapping = mapping

        async def check_health(self, *, source):
            result = self.mapping[source.name]
            if result == "hang":
                await asyncio.Event().wait()
            if isinstance(result, Exception):
                raise result
            return result

    async def test_per_source_dict(self):
        s1, s2 = _src(name="m1"), _src(name="m2", base_url="http://10.77.0.20:7867")
        client, _, _ = _client([s1, s2], [])
        client._transport = self._HealthTransport({"m1": True, "m2": False})
        assert await client.check_health() == {"m1": True, "m2": False}

    async def test_cancellation_passes_through(self):
        client, _, _ = _client([_src()], [])
        client._transport = self._HealthTransport({"m1": "hang"})
        task = asyncio.create_task(client.check_health())
        await asyncio.sleep(0.01)
        task.cancel()
        with pytest.raises(asyncio.CancelledError):
            await task


class TestTelemetry:
    async def test_success_and_failure_recorded_without_image_bytes(self):
        recorder = _MemoryRecorder()
        client, _, _ = _client(
            [_src()],
            [TransientError("boom", status_code=500), "text"],
            telemetry=recorder,
        )
        await client.recognize_text(b"RAW-IMAGE-BYTES")
        assert len(recorder.rows) == 2  # 失败尝试与成功尝试均必录
        for row in recorder.rows:
            assert "<ocr:text image_bytes=15>" in row["messages"]
            assert "RAW-IMAGE-BYTES" not in row["messages"]  # 图像字节绝不入库
        assert recorder.rows[0]["error"].startswith("TransientError:")  # 类名前缀口径
        assert recorder.rows[1]["error"] is None
        assert recorder.rows[1]["prompt_tokens"] == 0


class TestAssembly:
    _ENV = {
        "OCR__MONKEY__1__BASE_URL": "http://10.77.0.20:7866",
        "OCR__MONKEY__1__API_KEY": "none",
        "OCR__MONKEY__1__MODEL": "monkey-ocr",
        "OCR__MONKEY__1__TIMEOUT_S": "120",
        "LLM_MAX_RETRIES": "3",
        "LLM_RETRY_BASE_DELAY": "2.0",
        "LLM_RETRY_MAX_DELAY": "30.0",
        "LLM_CIRCUIT_BREAKER_THRESHOLD": "5",
        "LLM_CIRCUIT_BREAKER_COOLDOWN": "60",
        "PGW_CACHE_BACKEND": "none",
        "PGW_TELEMETRY_BACKEND": "none",
    }

    async def test_from_env_assembles(self):
        client = OcrClient.from_env("OCR", env=dict(self._ENV))
        try:
            assert client._scope == "ocr"
        finally:
            await client.aclose()

    async def test_non_monkey_provider_rejected(self):
        env = {k.replace("MONKEY", "GLM"): v for k, v in self._ENV.items()}
        with pytest.raises(ValueError, match="monkey"):
            OcrClient.from_env("OCR", env=env)

    async def test_aclose_idempotent(self):
        client = OcrClient.from_env("OCR", env=dict(self._ENV))
        await client.aclose()
        await client.aclose()