"""三层流式活性超时看门狗：给任意异步迭代器施加首产出/产出间/总时长三层超时保护。

领域无关、不打日志、不触网（Humble Object）：仅以 asyncio.timeout 包裹取值，
便于注入假异步迭代器做单测。超时抛 StreamLivenessTimeout（携 kind/elapsed/是否已出首项），
由 adapter 边界翻译为领域错误并埋点。

参考: CHSAnalyzer2 app/providers/streaming.py
"""

from __future__ import annotations

import asyncio
import contextlib
import time
from typing import TYPE_CHECKING, TypeVar

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

_T = TypeVar("_T")


class StreamLivenessTimeout(Exception):  # noqa: N818 — 规格要求名称，非 Error 后缀
    """流活性超时异常。

    属性:
        kind: 超时类型，取值 {"ttft", "inter_token", "total"}。
        elapsed_s: 自流开始到触发超时的秒数。
        first_token_seen: 触发时是否已产出过首项。
    """

    def __init__(self, kind: str, elapsed_s: float, first_token_seen: bool) -> None:
        self.kind = kind
        self.elapsed_s = elapsed_s
        self.first_token_seen = first_token_seen
        super().__init__(f"流活性超时({kind}, elapsed={elapsed_s:.1f}s)")


async def _anext_within(
    it: AsyncIterator[_T],
    timeout_s: float,
    *,
    kind: str,
    start: float,
    first: bool,
) -> _T:
    """限时取下一项；本层 deadline 触发抛 StreamLivenessTimeout(kind)，上游自抛的 TimeoutError 原样上抛。

    参数:
        it: 底层异步迭代器。
        timeout_s: 本次取项的超时预算（秒）。
        kind: 本层超时归类（ttft/inter_token/total），仅在本层 deadline 触发时使用。
        start: 流开始时刻（time.monotonic），用于计算 elapsed_s。
        first: 是否尚未产出过首项（用于 StreamLivenessTimeout.first_token_seen）。

    返回:
        source 的下一项。

    异常:
        StreamLivenessTimeout: 本层 asyncio.timeout deadline 触发（cm.expired()）。
        StopAsyncIteration: 底层迭代器耗尽（由调用方处理收尾）。
    """
    try:
        async with asyncio.timeout(timeout_s) as cm:
            return await it.__anext__()
    except TimeoutError:
        if not cm.expired():
            raise  # 上游自抛的 TimeoutError，非本层 deadline，原样上抛
        raise StreamLivenessTimeout(kind, time.monotonic() - start, not first) from None


async def stream_with_liveness_timeouts(
    source: AsyncIterator[_T],
    *,
    ttft_s: float,
    inter_token_s: float,
    total_s: float,
) -> AsyncIterator[_T]:
    """逐项产出 source，并施加三层活性超时。

    参数:
        source: 待包裹的异步迭代器（如 VLM 流式分片）。
        ttft_s: 首项产出超时（time-to-first-token，秒）。
        inter_token_s: 相邻两项之间的产出超时（秒）。
        total_s: 整条流的总时长超时（秒）。

    产出:
        与 source 逐项一致的元素。

    异常:
        StreamLivenessTimeout: 首项超 ttft_s、后续项间超 inter_token_s 或整体超 total_s。

    关键实现:
        超时**只包裹单次 __anext__**，绝不包裹 yield——否则总时长 deadline 会在生成器挂起
        （消费者处理已产出项）期间继续计时，可能取消消费者的 await 而非由本生成器稳定抛出
        StreamLivenessTimeout（取消泄漏）。总时长改由「每轮取 min(本层预算, 剩余总额) 钳制」
        实现：钳到剩余总额（clamped）时若超时即判 total，否则判 ttft/inter_token。用
        asyncio.timeout(...) 的 cm.expired() 区分本层 deadline 与上游自抛的 TimeoutError。
        finally 中调用底层迭代器的 aclose（若有）以释放上游资源。
    """
    it = source.__aiter__()
    start = time.monotonic()
    deadline = start + total_s
    first = True
    try:
        while True:
            remaining_total = deadline - time.monotonic()
            if remaining_total <= 0:
                raise StreamLivenessTimeout("total", time.monotonic() - start, not first)
            budget = ttft_s if first else inter_token_s
            # 剩余总额比本层预算更紧 → 先到的是 total
            clamped = remaining_total <= budget
            kind = "total" if clamped else ("ttft" if first else "inter_token")
            try:
                item = await _anext_within(
                    it,
                    remaining_total if clamped else budget,
                    kind=kind,
                    start=start,
                    first=first,
                )
            except StopAsyncIteration:
                return
            first = False
            yield item
    finally:
        # 收尾关闭上游迭代器；清理失败不得掩盖主异常，故显式抑制
        aclose = getattr(it, "aclose", None)
        if aclose is not None:
            with contextlib.suppress(Exception):
                await aclose()