"""并发调用多个下游（例如大模型 / 检索 / 第三方 API）：并发上限、单次超时、有限重试、部分失败、整体截止时间。

只用标准库。运行：python fanout.py；测试：python -m pytest -q test_fanout.py
"""
from __future__ import annotations

import asyncio
import random
import time
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Generic, TypeVar

T = TypeVar("T")


class RetryableError(Exception):
    """可以重试的错误（限流 429、服务暂时不可用 503 等）。"""


@dataclass
class Result(Generic[T]):
    key: str
    value: T | None = None
    error: str | None = None
    attempts: int = 0


async def call_with_retry(fn: Callable[[], Awaitable[T]], *, timeout: float, retries: int,
                          base_delay: float = 0.05) -> tuple[T, int]:
    """单次调用带超时；遇到超时或可重试错误时指数退避重试。返回 (结果, 尝试次数)。"""
    for attempt in range(1, retries + 2):
        try:
            async with asyncio.timeout(timeout):  # Python 3.11+
                return await fn(), attempt
        except (TimeoutError, RetryableError):
            if attempt > retries:
                raise
            await asyncio.sleep(base_delay * 2 ** (attempt - 1) * (1 + random.random() / 2))
    raise AssertionError("不会执行到这里")


async def fan_out(jobs: dict[str, Callable[[], Awaitable[T]]], *, concurrency: int, timeout: float,
                  retries: int = 1, deadline: float | None = None) -> list[Result[T]]:
    """并发执行所有任务，最多 concurrency 个同时进行；单个失败不影响其他任务。
    deadline：整体截止时间（秒），到时未完成的任务被取消并标记为超时。"""
    sem = asyncio.Semaphore(concurrency)
    results = {k: Result[T](key=k) for k in jobs}

    async def run(key: str, fn: Callable[[], Awaitable[T]]) -> None:
        async def counted() -> T:  # 记录尝试次数（失败时也能看到试了几次）
            results[key].attempts += 1
            return await fn()

        async with sem:
            try:
                results[key].value, _ = await call_with_retry(counted, timeout=timeout, retries=retries)
            except asyncio.CancelledError:
                results[key].error = "整体超时，已取消"
                raise  # 取消必须继续向上传播
            except Exception as e:  # 单个任务失败只记录，不影响其他任务
                results[key].error = f"{type(e).__name__}: {e}" if str(e) else type(e).__name__

    tasks = [asyncio.create_task(run(k, fn)) for k, fn in jobs.items()]
    done, pending = await asyncio.wait(tasks, timeout=deadline)
    for t in pending:
        t.cancel()
    await asyncio.gather(*pending, return_exceptions=True)  # 等待取消完成，避免任务泄漏
    return [results[k] for k in jobs]


# ---------------- 演示 ----------------
def fake_llm(name: str, delay: float, fail_times: int = 0) -> Callable[[], Awaitable[str]]:
    state = {"left": fail_times}

    async def call() -> str:
        await asyncio.sleep(delay)
        if state["left"] > 0:
            state["left"] -= 1
            raise RetryableError("429 Too Many Requests")
        return f"{name} 的回答"

    return call


async def demo() -> None:
    jobs = {
        "摘要": fake_llm("摘要", 0.10),
        "分类": fake_llm("分类", 0.10, fail_times=1),  # 第一次被限流，重试后成功
        "翻译": fake_llm("翻译", 2.00),                # 太慢，单次超时
        "关键词": fake_llm("关键词", 0.05),
    }
    start = time.perf_counter()
    for r in await fan_out(jobs, concurrency=2, timeout=0.5, retries=1):
        print(f"{r.key}: {r.value or r.error}（尝试 {r.attempts} 次）")
    print(f"总耗时 {time.perf_counter() - start:.2f}s")


if __name__ == "__main__":
    asyncio.run(demo())