"""一次性诊断：实测 SiliconFlow API 在不同并发下的 429 限流情况。

用于确定 llm_filter 的安全并发上限（run_011d01f19999 在 20/4 并发下均 429）。
"""
import json
import os
import threading
import time
import urllib.error
import urllib.request
from pathlib import Path

from dotenv import load_dotenv

load_dotenv(Path(".env"))

API_BASE = os.getenv("LLM_API_BASE", "https://api.siliconflow.cn/v1/chat/completions")
API_KEY = os.getenv("LLM_API_KEY", "")
MODEL = os.getenv("LLM_MODEL", "Qwen/Qwen3.5-35B-A3B")


def one_call(idx: int) -> str:
    """发送一次最小 LLM 请求，返回结果码：ok / 429 / 其他。"""
    body = {
        "model": MODEL,
        "messages": [
            {"role": "system", "content": "你是字幕质量过滤器。"},
            {"role": "user", "content": f"【目标】测试请求 {idx}"},
        ],
        "enable_thinking": False,
        "max_tokens": 8,
    }
    req = urllib.request.Request(
        API_BASE,
        data=json.dumps(body).encode("utf-8"),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {API_KEY}"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            resp.read()
            return "ok"
    except urllib.error.HTTPError as exc:
        # 429 时尝试读取响应体里的限流信息
        try:
            detail = exc.read().decode("utf-8")[:120]
        except Exception:
            detail = ""
        retry_after = exc.headers.get("Retry-After") if exc.headers else None
        return f"HTTP {exc.code} | Retry-After={retry_after} | {detail}"
    except Exception as exc:  # noqa: BLE001
        return f"{type(exc).__name__}: {exc}"


def probe(concurrency: int, total: int) -> list[str]:
    """以给定并发发送 total 个请求，返回结果码列表。"""
    results: list[str] = []
    lock = threading.Lock()
    index = 0

    def worker() -> None:
        nonlocal index
        while True:
            with lock:
                if index >= total:
                    return
                i = index
                index += 1
            with lock:
                results.append(one_call(i))

    threads = [threading.Thread(target=worker) for _ in range(concurrency)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    return results


if __name__ == "__main__":
    for conc in (1, 2, 4, 8):
        results = probe(conc, total=8)
        ok = sum(1 for r in results if r == "ok")
        other = [r for r in results if r != "ok"][:2]
        print(f"并发 {conc:2d}: 8 请求 -> ok {ok}/8 | 其余样例: {other}")
        time.sleep(2)  # 组间休息，避免误伤