# -*- coding: utf-8 -*-
"""
压测每个 OCR 端点(IP)的实际 QPS，并把结果写成权重文件 config/ocr_weights.json，
供运行时按 QPS 加权分发请求(见 tool/Ocr/tencentcloud_ocr.py 的 OcrDispatcher)。

做法：对每个端点用固定并发、固定请求数发同一张样例图，
QPS = 成功请求数 / 墙钟耗时；同时报告平均时延和错误率。

本脚本完全自包含：单文件，不读 config、不 import 项目任何模块，拷过去就能跑。

用法:
  python bench_ocr_endpoints.py [样例图片] [--requests N] [--concurrency C] [--no-write]
  python bench_ocr_endpoints.py [样例图片] --url http://IP:PORT [--url http://IP2:PORT ...]

  样例图片不给时，自动在 ./composite、./split 下找一张 .jpg。
  端点：默认内置 Polaris 地址；要测别的 IP 用 --url(可多次)。
  --requests     每个端点发的总请求数(默认 60)
  --concurrency  每个端点的并发线程数(默认 10)
  --no-write     只打印，不写 ocr_weights.json
"""
import argparse
import base64
import glob
import itertools
import json
import os
import random
import threading
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

import requests

OCR_PATH = "/youtu/ocrapi/hpgeneralocr_v3"
OCR_PORT = 18081

# 内置的多台 OCR 服务 IP（自动负载均衡测试用）
DEFAULT_IPS = [
    "30.154.78.83"
]
DEFAULT_ENDPOINTS = [f"http://{ip}:{OCR_PORT}" for ip in DEFAULT_IPS]
WEIGHTS_FILE = "./ocr_weights.json"


def get_endpoints(cli_urls):
    """命令行 --url 优先，否则用内置默认端点。"""
    if cli_urls:
        return [u.rstrip("/") for u in cli_urls]
    return list(DEFAULT_ENDPOINTS)


def find_sample_image(given):
    if given:
        if os.path.exists(given):
            return given
        print(f"指定的图片不存在: {given}")
        return None
    for root in ("./composite", "./split"):
        hits = glob.glob(os.path.join(root, "**", "*.jpg"), recursive=True)
        if hits:
            return hits[0]
    return None


def img2base64(img_path):
    with open(img_path, "rb") as f:
        return base64.b64encode(f.read()).decode()


def ocr_request_once(base_url, image_b64, timeout=60):
    """向指定端点发一次 OCR 请求，返回原始响应文本；失败抛异常。"""
    url = base_url.rstrip("/") + OCR_PATH
    app_id = f"{int(time.time() * 1000)}{random.randint(100000, 999999)}"
    payload = {
        "app_id": app_id,
        "image": image_b64,
        "ocr_template": "OCR",
        "options": {"pdf_page_index": 0, "detect_text": True},
    }
    resp = requests.post(url, headers={"Content-Type": "application/json"},
                         json=payload, timeout=timeout)
    resp.raise_for_status()
    return resp.text


def _ok(raw):
    """判断一次响应是否算成功：有内容且不是 error。"""
    if not raw:
        return False
    try:
        data = json.loads(raw)
    except (json.JSONDecodeError, TypeError):
        return False
    if "TextDetections" in data:
        return True
    return data.get("error_code", 0) == 0


def bench_endpoint(base_url, image_b64, requests_n, concurrency):
    """对单个端点压测，返回统计 dict（含一个失败原因采样 sample_err）。"""
    latencies = []
    ok = 0
    err = 0
    sample_err = [None]   # 采样第一个失败原因

    def one_call():
        t0 = time.time()
        try:
            raw = ocr_request_once(base_url, image_b64, timeout=60)
            dt = time.time() - t0
            if _ok(raw):
                return (True, dt)
            # HTTP 200 但内容是错误：采样响应体片段
            if sample_err[0] is None:
                sample_err[0] = f"响应非成功: {raw[:200]}"
            return (False, dt)
        except Exception as e:
            if sample_err[0] is None:
                sample_err[0] = f"{type(e).__name__}: {e}"
            return (False, time.time() - t0)

    t_start = time.time()
    with ThreadPoolExecutor(max_workers=concurrency) as ex:
        futures = [ex.submit(one_call) for _ in range(requests_n)]
        for fut in as_completed(futures):
            success, dt = fut.result()
            latencies.append(dt)
            if success:
                ok += 1
            else:
                err += 1
    wall = time.time() - t_start

    return {
        "qps": ok / wall if wall > 0 else 0.0,
        "ok": ok,
        "err": err,
        "wall": wall,
        "avg_latency": sum(latencies) / len(latencies) if latencies else 0.0,
        "err_rate": err / requests_n if requests_n else 0.0,
        "sample_err": sample_err[0],
    }


def soak_endpoint(base_url, image_pool, concurrency, report_interval=5):
    """对单个端点持续满速压测：固定 concurrency 个线程不停发请求，
    每次从 image_pool 随机抽一张图(避免被服务端缓存)，每 report_interval 秒打印滚动 QPS，
    直到 Ctrl+C。image_pool 可以是单张或几千张。"""
    stop = threading.Event()
    lock = threading.Lock()
    st = {"ok": 0, "err": 0, "lat": 0.0, "n": 0}
    sample_err = [None]

    def worker():
        while not stop.is_set():
            img = random.choice(image_pool)        # 每次随机抽一张
            t0 = time.time()
            try:
                raw = ocr_request_once(base_url, img, timeout=60)
                success = _ok(raw)
                if not success and sample_err[0] is None:
                    sample_err[0] = f"响应非成功: {raw[:160]}"
            except Exception as e:
                success = False
                if sample_err[0] is None:
                    sample_err[0] = f"{type(e).__name__}: {e}"
            dt = time.time() - t0
            with lock:
                st["n"] += 1
                st["lat"] += dt
                if success:
                    st["ok"] += 1
                else:
                    st["err"] += 1

    threads = [threading.Thread(target=worker, daemon=True) for _ in range(concurrency)]
    for t in threads:
        t.start()

    print(f"持续满速压测 {base_url}   并发={concurrency}   每{report_interval}s 报一次   (Ctrl+C 停止)")
    print(f"{'用时s':>6} {'累计OK':>9} {'瞬时QPS':>9} {'平均QPS':>9} {'时延ms':>8} {'错误率':>7}")
    print("-" * 56)
    t_start = time.time()
    last_ok, last_t = 0, t_start
    try:
        while True:
            time.sleep(report_interval)
            now = time.time()
            with lock:
                ok, err, n, lat = st["ok"], st["err"], st["n"], st["lat"]
            inst = (ok - last_ok) / (now - last_t) if now > last_t else 0
            avg = ok / (now - t_start) if now > t_start else 0
            avg_lat = (lat / n * 1000) if n else 0
            er = err / n * 100 if n else 0
            print(f"{now-t_start:6.0f} {ok:>9} {inst:>9.1f} {avg:>9.1f} {avg_lat:>8.0f} {er:>6.0f}%", flush=True)
            last_ok, last_t = ok, now
    except KeyboardInterrupt:
        stop.set()
        el = time.time() - t_start
        with lock:
            ok, n = st["ok"], st["n"]
        print("-" * 56)
        print(f"停止。{el:.0f}s 内成功 {ok} / 共发 {n}，平均 {ok/el if el else 0:.1f} QPS")
        if sample_err[0]:
            print(f"失败采样: {sample_err[0]}")


def health_check(endpoints, image_b64):
    """给每个端点发一次请求，返回 (存活列表, 死亡列表)。"""
    alive, dead = [], []

    def _check(ep):
        try:
            return ep, _ok(ocr_request_once(ep, image_b64, timeout=15))
        except Exception:
            return ep, False

    with ThreadPoolExecutor(max_workers=min(32, len(endpoints))) as ex:
        for ep, ok in ex.map(_check, endpoints):
            (alive if ok else dead).append(ep)
    return alive, dead


def bench_aggregate(endpoints, image_b64, requests_n, concurrency):
    """把请求轮询分发到所有端点（模拟运行时自动负载均衡），测合并 QPS。"""
    rr = itertools.cycle(endpoints)
    rr_lock = threading.Lock()
    latencies = []
    ok = 0
    err = 0
    stat_lock = threading.Lock()

    def next_ep():
        with rr_lock:
            return next(rr)

    def one_call(_):
        ep = next_ep()
        t0 = time.time()
        try:
            raw = ocr_request_once(ep, image_b64, timeout=60)
            success = _ok(raw)
        except Exception:
            success = False
        dt = time.time() - t0
        with stat_lock:
            latencies.append(dt)
            return success

    t_start = time.time()
    with ThreadPoolExecutor(max_workers=concurrency) as ex:
        for success in ex.map(one_call, range(requests_n)):
            if success:
                ok += 1
            else:
                err += 1
    wall = time.time() - t_start

    return {
        "qps": ok / wall if wall > 0 else 0.0,
        "ok": ok,
        "err": err,
        "wall": wall,
        "avg_latency": sum(latencies) / len(latencies) if latencies else 0.0,
        "err_rate": err / requests_n if requests_n else 0.0,
    }


def run_aggregate(endpoints, image_b64, requests_n, concurrency):
    """先健康检查筛掉死端点，再在存活端点上做聚合负载测试。"""
    print(f"健康检查 {len(endpoints)} 个端点 ...", flush=True)
    alive, dead = health_check(endpoints, image_b64)
    print(f"  存活: {len(alive)}   死亡: {len(dead)}")
    if dead:
        for ep in dead:
            print(f"    ✗ {ep}")
    if not alive:
        print("没有存活端点，无法测试。")
        return
    print("=" * 64)
    print(f"聚合负载测试: {len(alive)} 个存活端点轮询分发，"
          f"总请求 {requests_n}，总并发 {concurrency}")
    s = bench_aggregate(alive, image_b64, requests_n, concurrency)
    print("-" * 64)
    print(f"  合并 QPS = {s['qps']:.2f}   成功={s['ok']}  失败={s['err']}"
          f"   平均时延={s['avg_latency']*1000:.0f}ms   错误率={s['err_rate']*100:.0f}%")
    if alive:
        print(f"  单端点平均 ≈ {s['qps']/len(alive):.2f} QPS"
              f"   （{len(alive)} 台合计 {s['qps']:.0f} QPS）")
    print("=" * 64)


def sweep_endpoint(base_url, image_b64, start, max_c, per_worker):
    """自动加压找极限 QPS：并发从 start 逐级翻倍到 max_c，
    QPS 不再随并发上升(或错误率升高)即判定到极限，停止。"""
    print(f"自动加压 {base_url}  (并发 {start}→{max_c} 逐级翻倍)")
    print(f"{'并发':>6} {'QPS':>9} {'时延ms':>9} {'错误率':>8}")
    print("-" * 36)
    best_qps, best_c = 0.0, 0
    prev_qps = 0.0
    c = start
    while c <= max_c:
        reqs = max(c * per_worker, 30)        # 每个并发至少发 per_worker 次，样本够稳
        s = bench_endpoint(base_url, image_b64, reqs, c)
        print(f"{c:>6} {s['qps']:>9.2f} {s['avg_latency']*1000:>9.0f} {s['err_rate']*100:>7.0f}%")
        if s["qps"] > best_qps:
            best_qps, best_c = s["qps"], c
        if s["err_rate"] > 0.1:
            print("  → 错误率升高，已超过服务承载，停止。")
            break
        if prev_qps > 0 and s["qps"] < prev_qps * 1.08:
            print("  → QPS 不再随并发上升，已到极限，停止。")
            break
        prev_qps = s["qps"]
        c *= 2
    print("-" * 36)
    print(f"极限 QPS ≈ {best_qps:.2f}  (最佳并发 {best_c})")
    return best_qps, best_c


def gen_random_images(n, size=(1280, 200)):
    """生成 n 张带随机文字的图片(让 OCR 真干活)，返回 base64 列表。需要 Pillow。"""
    try:
        import io
        import string
        from PIL import Image, ImageDraw
    except ImportError:
        print("生成随机图需要 Pillow(pip install pillow)；改用 --pool 指向真实图目录")
        raise
    chars = string.ascii_letters + string.digits + " ,.!?-"
    pool = []
    print(f"生成 {n} 张随机文字图...", flush=True)
    for i in range(n):
        bg = random.randint(200, 255)
        img = Image.new("RGB", size, (bg, bg, bg))
        d = ImageDraw.Draw(img)
        for row in range(random.randint(1, 3)):
            text = "".join(random.choice(chars) for _ in range(random.randint(10, 40)))
            d.text((20, 20 + row * 50), text, fill=(0, 0, 0))
        buf = io.BytesIO()
        img.save(buf, format="JPEG", quality=85)
        pool.append(base64.b64encode(buf.getvalue()).decode())
        if (i + 1) % 1000 == 0:
            print(f"  已生成 {i+1}/{n}", flush=True)
    return pool


def build_image_pool(args):
    """构建测试图池(base64 列表)：
    --pool DIR 用真实图目录(最贴近实际负载) > --gen N 随机生成 > 单张样例图。"""
    if getattr(args, "pool", None):
        files = (glob.glob(os.path.join(args.pool, "**", "*.jpg"), recursive=True) +
                 glob.glob(os.path.join(args.pool, "**", "*.png"), recursive=True))
        if not files:
            print(f"目录 {args.pool} 下没有图片")
            return []
        random.shuffle(files)
        cap = args.gen if args.gen else 5000
        files = files[:cap]
        print(f"从 {args.pool} 取 {len(files)} 张真实图作为图池")
        return [img2base64(f) for f in files]
    if getattr(args, "gen", 0):
        return gen_random_images(args.gen)
    sample = find_sample_image(args.image)
    if not sample:
        return []
    print(f"单张样例图: {sample}")
    return [img2base64(sample)]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("image", nargs="?", default=None)
    ap.add_argument("--url", action="append", default=None, help="OCR 端点(可多次)")
    ap.add_argument("--requests", type=int, default=60)
    ap.add_argument("--concurrency", type=int, default=10)
    ap.add_argument("--aggregate", action="store_true",
                    help="聚合负载测试:请求轮询分发到所有存活端点，测合并总 QPS")
    ap.add_argument("--sweep", action="store_true",
                    help="自动加压找极限 QPS(并发逐级翻倍直到 QPS 不再上升)")
    ap.add_argument("--soak", action="store_true",
                    help="单端点持续满速压测(固定并发不停发，实时滚动QPS，Ctrl+C 停)")
    ap.add_argument("--interval", type=int, default=5, help="soak 模式打印间隔秒")
    ap.add_argument("--gen", type=int, default=0,
                    help="随机生成 N 张带文字的图作为图池(每次请求随机抽，避免被缓存)")
    ap.add_argument("--pool", default=None,
                    help="用该目录下的真实图作为图池(最贴近实际负载，如 ./split/某视频)")
    ap.add_argument("--start-concurrency", type=int, default=10, help="sweep 起始并发")
    ap.add_argument("--max-concurrency", type=int, default=320, help="sweep 最大并发")
    ap.add_argument("--per-worker", type=int, default=6, help="sweep 每档每并发发多少次请求")
    ap.add_argument("--no-write", action="store_true")
    args = ap.parse_args()

    endpoints = get_endpoints(args.url)
    if not endpoints:
        print("没有可用的 OCR 端点")
        return

    # 构建图池：--pool 真实图目录 / --gen 随机生成 / 默认单张样例
    image_pool = build_image_pool(args)
    if not image_pool:
        print("没有可用图片：给一张 jpg 路径、或用 --gen N / --pool 目录。")
        return
    print(f"图池大小: {len(image_pool)} 张  |  端点数: {len(endpoints)}  并发: {args.concurrency}")
    print("=" * 64)

    image_b64 = image_pool[0]   # 非 soak 模式用单张(测纯吞吐)

    # ---- soak 模式：单端点持续满速压测，每次随机抽图，直到 Ctrl+C ----
    if args.soak:
        soak_endpoint(endpoints[0], image_pool, args.concurrency, args.interval)
        return

    # ---- aggregate 模式：所有端点轮询分发，测合并总 QPS ----
    if args.aggregate:
        run_aggregate(endpoints, image_b64, args.requests, args.concurrency)
        return

    # ---- sweep 模式：自动加压找极限，不写权重文件 ----
    if args.sweep:
        for i, ep in enumerate(endpoints, 1):
            print(f"\n[{i}/{len(endpoints)}]", end=" ")
            sweep_endpoint(ep, image_b64, args.start_concurrency,
                           args.max_concurrency, args.per_worker)
        return

    results = {}
    for i, ep in enumerate(endpoints, 1):
        print(f"[{i}/{len(endpoints)}] 压测 {ep} ...", flush=True)
        stat = bench_endpoint(ep, image_b64, args.requests, args.concurrency)
        results[ep] = stat
        print(f"    QPS={stat['qps']:.2f}  成功={stat['ok']}  失败={stat['err']}"
              f"  平均时延={stat['avg_latency']*1000:.0f}ms"
              f"  错误率={stat['err_rate']*100:.0f}%")
        if stat["err"] and stat.get("sample_err"):
            print(f"    失败原因(采样): {stat['sample_err']}")

    print("=" * 64)
    print("汇总(按 QPS 降序):")
    for ep, s in sorted(results.items(), key=lambda kv: kv[1]["qps"], reverse=True):
        print(f"  {ep:48s}  QPS={s['qps']:6.2f}  时延={s['avg_latency']*1000:5.0f}ms"
              f"  错误率={s['err_rate']*100:.0f}%")

    total_qps = sum(s["qps"] for s in results.values())
    print(f"\n合计可用 QPS ≈ {total_qps:.2f}")

    # 写权重文件：QPS 即权重；错误率过高(>50%)的端点权重清零(不参与分发)
    weights = {ep: (round(s["qps"], 3) if s["err_rate"] <= 0.5 else 0.0)
               for ep, s in results.items()}

    if args.no_write:
        print("\n--no-write：不写权重文件。")
        return

    os.makedirs(os.path.dirname(WEIGHTS_FILE), exist_ok=True)
    with open(WEIGHTS_FILE, "w", encoding="utf-8") as f:
        json.dump(weights, f, ensure_ascii=False, indent=2)
    print(f"\n已写入权重文件: {WEIGHTS_FILE}")
    print("运行时会按此权重把 OCR 请求加权分发到各端点（重启进程或新批次生效）。")


if __name__ == "__main__":
    main()