#!/usr/bin/env python
"""E2E-проверка TTFT-watchdog обёртки воркера (e-261014163808-4da3ec).

Три сценария на фейковом стриме:
 1. живой апстрим (контент сразу)   → все чанки, ни одного ретрая;
 2. висит до первого токена, оживёт на 2-й попытке → ретрай, ответ完整;
 3. висит всегда (3 попытки)        → TimeoutError с внятным текстом.

Плюс дым: сборка модели payload-ом возвращает прокси с рабочим stream().
Запуск: ./venv/bin/python tests/test_ttft_retry.py
"""
import asyncio
import os
import sys
import types

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from harness import _llm_worker as W  # noqa: E402

FAILS = []


def check(name, cond, extra=""):
    print(("PASS " if cond else "FAIL ") + name + (f" | {extra}" if extra and not cond else ""))
    if not cond:
        FAILS.append(name)


def chunk(ct):
    return {"chunk_type": ct}


def fake_stream_model(script):
    """script: список «попыток», каждая — список чанков или raise Hang."""
    calls = {"n": 0}

    class M:
        async def stream(self, *a, **k):
            i = calls["n"]
            calls["n"] += 1
            act = script[i]
            if act == "HANG":
                await asyncio.sleep(30)  # дольше watchdog-бюджета
                yield chunk("message_start")
                return
            for c in act:
                yield chunk(c)

    return M(), calls


async def run(model):
    wrapped = W._with_ttft_retry(model, "hour", 90)
    out = []
    async for c in wrapped.stream():
        out.append(W._chunk_type(c))
    return out


async def main():
    os.environ["SIMBIOT_LLM_TTFT_BUDGET"] = "0.5"  # быстрый watchdog для теста
    W._TTFT_FIRST_TOKEN_BUDGET = 0.5

    # 1. живой апстрим
    m, calls = fake_stream_model([
        ["message_start", "content_delta", "content_delta", "message_stop"]])
    out = await run(m)
    check("живой стрим проходит целиком",
          out == ["message_start", "content_delta", "content_delta", "message_stop"],
          str(out))
    check("живой стрим — 1 попытка", calls["n"] == 1, str(calls))

    # 2. виснет, оживает на 2-й
    m, calls = fake_stream_model([
        "HANG",
        ["message_start", "content_delta", "message_stop"]])
    out = await run(m)
    check("зависший вызов ретраится и отвечает",
          out == ["message_start", "content_delta", "message_stop"], str(out))
    check("ретрай сделал ровно 2 попытки", calls["n"] == 2, str(calls))

    # 3. виснет всегда
    m, calls = fake_stream_model(["HANG", "HANG", "HANG"])
    try:
        await run(m)
        check("глухой апстрим даёт TimeoutError", False, "не упал")
    except TimeoutError as e:
        check("глухой апстрим даёт TimeoutError", "завис до первого токена" in str(e), str(e))
    check("глухой апстрим — ровно 3 попытки", calls["n"] == 3, str(calls))

    # 4. делегирование атрибутов прокси
    m, _ = fake_stream_model([["content_delta"]])
    m.foo = "bar"
    wrapped = W._with_ttft_retry(m, "hour", 90)
    check("прокси делегирует атрибуты", getattr(wrapped, "foo", None) == "bar")

    print()
    print("ИТОГ:", "ВСЕ ОК" if not FAILS else f"ПРОВАЛЫ: {FAILS}")
    sys.exit(1 if FAILS else 0)


asyncio.run(main())