# -*- coding: utf-8 -*-
"""
下载一个视频 → 抽帧 → 用这些【真实帧】对 OCR 端点做压测。
比随机生成的图更贴近线上负载（真实字幕内容，OCR 干的活和实际一致）。

用法（项目根目录运行）:
  python temp/video_bench.py <视频...> --url http://IP:18081 [选项]

  <视频...>  一个或多个，cos://bucket/path/x.mp4 或本地 ./video/x.mp4
             多个视频时各抽一部分帧合并成图池(max-frames 在视频间均分)
  也可用 --list 文件 从文件读视频列表(每行一个)

选项:
  --url URL          OCR 端点(可多次，soak 只用第一个)
  --concurrency N    并发(默认 100；单台拐点约 80)
  --mode soak|sweep  soak=持续满速(默认)；sweep=自动加压找极限
  --fps F            抽帧频率(默认 1，压测用不需要 3fps，省时省内存)
  --max-frames N     最多用多少帧做图池(默认 2000)
  --interval N       soak 打印间隔秒(默认 5)
  --ffmpeg PATH      ffmpeg 路径(默认 ./bin/ffmpeg)
  --coscli PATH      coscli 路径(默认 ./coscli-linux-amd64)
  --cos-config PATH  coscli 配置(默认 .cos.yaml)
  --keep             跑完不删临时视频/帧
"""
import argparse
import glob
import os
import shutil
import subprocess
import sys
import tempfile

# 复用同目录 bench 脚本的压测逻辑
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import bench_ocr_endpoints as bench   # noqa: E402


def download_video(src, coscli, cos_config, dst):
    """cos:// 路径用 coscli 下载到 dst；本地路径直接返回。"""
    if not src.startswith("cos://"):
        if not os.path.exists(src):
            print(f"本地视频不存在: {src}")
            return None
        return src
    print(f"下载视频: {src}")
    r = subprocess.run([coscli, "-c", cos_config, "cp", src, dst],
                       capture_output=True, text=True)
    if r.returncode != 0 or not os.path.exists(dst):
        print(f"下载失败: {r.stderr.strip()[:300]}")
        return None
    print(f"  → {dst} ({os.path.getsize(dst)/1024/1024:.0f} MB)")
    return dst


def extract_frames(ffmpeg, video, frame_dir, fps):
    """用 ffmpeg 抽帧到 frame_dir，返回帧文件列表。"""
    os.makedirs(frame_dir, exist_ok=True)
    print(f"抽帧 (fps={fps}) ...")
    cmd = [ffmpeg, "-i", video, "-vf", f"fps={fps}", "-q:v", "2",
           os.path.join(frame_dir, "f_%06d.jpg"), "-loglevel", "error", "-y"]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        print(f"抽帧失败: {r.stderr.strip()[:400]}")
        return []
    frames = sorted(glob.glob(os.path.join(frame_dir, "*.jpg")))
    print(f"  抽出 {len(frames)} 帧")
    return frames


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("videos", nargs="*", default=[], help="一个或多个视频(cos://... 或本地路径)")
    ap.add_argument("--list", default=None, help="从文件读视频列表(每行一个，可代替直接传)")
    ap.add_argument("--url", action="append", default=None, help="OCR 端点(可多次)")
    ap.add_argument("--concurrency", type=int, default=100)
    ap.add_argument("--mode", choices=["soak", "sweep"], default="soak")
    ap.add_argument("--fps", type=float, default=1.0)
    ap.add_argument("--max-frames", type=int, default=2000)
    ap.add_argument("--interval", type=int, default=5)
    ap.add_argument("--ffmpeg", default="./bin/ffmpeg")
    ap.add_argument("--coscli", default="./coscli-linux-amd64")
    ap.add_argument("--cos-config", default=".cos.yaml")
    ap.add_argument("--keep", action="store_true")
    args = ap.parse_args()

    endpoints = bench.get_endpoints(args.url)
    if not endpoints:
        print("没有 OCR 端点，用 --url 指定")
        return

    # 收集视频列表：命令行参数 + --list 文件
    videos = list(args.videos)
    if args.list and os.path.exists(args.list):
        with open(args.list, encoding="utf-8") as f:
            videos += [ln.strip() for ln in f if ln.strip()]
    if not videos:
        print("没有视频，传一个或多个视频路径，或用 --list 文件")
        return
    # 每个视频均分帧额度，避免一个长视频占满图池
    per_video_cap = max(1, args.max_frames // len(videos))

    work = tempfile.mkdtemp(prefix="video_bench_")
    try:
        all_frames = []
        for i, src in enumerate(videos):
            print(f"\n[{i+1}/{len(videos)}] {src}")
            vlocal = os.path.join(work, f"video_{i}.mp4")
            fdir = os.path.join(work, f"frames_{i}")
            v = download_video(src, args.coscli, args.cos_config, vlocal)
            if not v:
                continue
            fr = extract_frames(args.ffmpeg, v, fdir, args.fps)
            if len(fr) > per_video_cap:
                fr = fr[:per_video_cap]
                print(f"  取前 {per_video_cap} 帧")
            all_frames += fr
            # 处理完即删该视频，省空间
            if v == vlocal:
                try:
                    os.remove(vlocal)
                except OSError:
                    pass

        if not all_frames:
            print("没抽到任何帧")
            return

        # 帧 → base64 图池
        print(f"\n加载 {len(all_frames)} 帧为图池 ...")
        pool = [bench.img2base64(f) for f in all_frames]
        print("=" * 60)
        print(f"图池: {len(pool)} 张真实帧(来自 {len(videos)} 个视频)  |  "
              f"端点: {endpoints[0]}  |  并发: {args.concurrency}")
        print("=" * 60)

        # 4) 压测
        if args.mode == "sweep":
            bench.sweep_endpoint(endpoints[0], pool[0], 10, 320, 6)
        else:
            bench.soak_endpoint(endpoints[0], pool, args.concurrency, args.interval)

    finally:
        if args.keep:
            print(f"\n临时文件保留在: {work}")
        else:
            shutil.rmtree(work, ignore_errors=True)


if __name__ == "__main__":
    main()