import asyncio
import time

from fanout import RetryableError, fake_llm, fan_out


def run(coro):
    return asyncio.run(coro)


def test_concurrency_limit():
    active, peak = 0, 0

    def job():
        async def call():
            nonlocal active, peak
            active += 1
            peak = max(peak, active)
            await asyncio.sleep(0.05)
            active -= 1
            return "ok"
        return call

    res = run(fan_out({str(i): job() for i in range(10)}, concurrency=3, timeout=1))
    assert peak == 3 and all(r.value == "ok" for r in res)


def test_retry_then_success():
    [r] = run(fan_out({"a": fake_llm("a", 0.01, fail_times=1)}, concurrency=1, timeout=1, retries=1))
    assert r.value == "a 的回答" and r.attempts == 2


def test_partial_failure():
    async def boom():
        raise ValueError("参数错误")  # 不可重试的错误：只尝试一次

    res = run(fan_out({"ok": fake_llm("ok", 0.01), "bad": boom, "slow": fake_llm("slow", 1)},
                      concurrency=3, timeout=0.1, retries=1))
    by = {r.key: r for r in res}
    assert by["ok"].value == "ok 的回答"
    assert by["bad"].error == "ValueError: 参数错误"
    assert by["slow"].error == "TimeoutError"


def test_retry_exhausted():
    [r] = run(fan_out({"x": fake_llm("x", 0.01, fail_times=5)}, concurrency=1, timeout=1, retries=2))
    assert r.value is None and "RetryableError" in r.error


def test_overall_deadline_cancels():
    start = time.perf_counter()
    res = run(fan_out({"fast": fake_llm("fast", 0.01), "slow": fake_llm("slow", 5)},
                      concurrency=2, timeout=10, deadline=0.2))
    assert time.perf_counter() - start < 1
    by = {r.key: r for r in res}
    assert by["fast"].value and by["slow"].error == "整体超时，已取消"


def test_retryable_error_is_exception():
    assert issubclass(RetryableError, Exception)