#!/usr/bin/env python3
"""
ERNIE-Image use_pe 对比基准测试
用法: python benchmark_pe.py
从 prompt.txt 读取 100 条 prompt，分别测试 use_pe=False 和 use_pe=True 两种模式下
全部请求正常返回的总耗时与平均耗时。
"""

import requests
import argparse
import json
import time
import base64
import sys
from pathlib import Path


def load_prompts(path: str) -> list[str]:
    p = Path(path)
    if not p.exists():
        print(f"错误: 找不到 prompt 文件 {path}")
        sys.exit(1)
    prompts = [line.strip() for line in p.read_text(encoding="utf-8").splitlines() if line.strip()]
    if not prompts:
        print(f"错误: {path} 中没有有效的 prompt")
        sys.exit(1)
    print(f"已加载 {len(prompts)} 条 prompt")
    return prompts


def send_request(url: str, payload: dict, timeout: int) -> tuple[bool, float]:
    """发送单次请求，返回 (是否成功, 耗时秒数)"""
    start = time.time()
    try:
        resp = requests.post(url, json=payload, timeout=timeout)
        elapsed = time.time() - start
        if resp.status_code == 200:
            return True, elapsed
        print(f"  [失败] status={resp.status_code}, prompt={payload['prompt'][:50]}...")
        return False, elapsed
    except requests.exceptions.Timeout:
        elapsed = time.time() - start
        print(f"  [超时] prompt={payload['prompt'][:50]}...")
        return False, elapsed
    except Exception as e:
        elapsed = time.time() - start
        print(f"  [异常] {e}, prompt={payload['prompt'][:50]}...")
        return False, elapsed


def run_benchmark(
    prompts: list[str],
    server_url: str,
    use_pe: bool,
    steps: int,
    width: int,
    height: int,
    guidance_scale: float,
    seed: int,
    timeout: int,
    concurrency: int,
) -> dict:
    """执行一轮基准测试，返回汇总结果。"""
    url = f"{server_url.rstrip('/')}/v1/images/generations"
    label = "use_pe=True" if use_pe else "use_pe=False"
    print(f"\n{'='*60}")
    print(f"开始测试: {label}  |  共 {len(prompts)} 条, 并发={concurrency}")
    print(f"{'='*60}")

    from concurrent.futures import ThreadPoolExecutor, as_completed

    success_count = 0
    fail_count = 0
    total_time = 0.0
    timings: list[float] = []

    def task(prompt: str, idx: int):
        payload = {
            "prompt": prompt,
            "size": f"{width}x{height}",
            "num_inference_steps": steps,
            "guidance_scale": guidance_scale,
            "seed": seed,
            "negative_prompt": " ",
            "response_format": "b64_json",
            "n": 1,
            "use_pe": use_pe,
        }
        ok, elapsed = send_request(url, payload, timeout)
        return idx, ok, elapsed

    wall_start = time.time()

    with ThreadPoolExecutor(max_workers=concurrency) as pool:
        futures = {pool.submit(task, p, i): i for i, p in enumerate(prompts)}
        for future in as_completed(futures):
            idx, ok, elapsed = future.result()
            if ok:
                success_count += 1
                timings.append(elapsed)
            else:
                fail_count += 1
            total_time += elapsed
            # 进度
            done = success_count + fail_count
            if done % 10 == 0 or done == len(prompts):
                print(f"  进度: {done}/{len(prompts)}  成功={success_count}  失败={fail_count}")

    wall_elapsed = time.time() - wall_start

    result = {
        "label": label,
        "use_pe": use_pe,
        "total_prompts": len(prompts),
        "success": success_count,
        "fail": fail_count,
        "wall_time": wall_elapsed,
        "avg_success_time": sum(timings) / len(timings) if timings else 0,
        "min_time": min(timings) if timings else 0,
        "max_time": max(timings) if timings else 0,
        "sum_request_time": sum(timings),
    }
    return result


def print_result(r: dict):
    print(f"\n--- {r['label']} 结果 ---")
    print(f"  总请求数:   {r['total_prompts']}")
    print(f"  成功:       {r['success']}")
    print(f"  失败:       {r['fail']}")
    print(f"  实际耗时(Wall Time): {r['wall_time']:.2f}s")
    print(f"  成功请求累计耗时:     {r['sum_request_time']:.2f}s")
    print(f"  成功请求平均耗时:     {r['avg_success_time']:.2f}s")
    print(f"  最快:       {r['min_time']:.2f}s")
    print(f"  最慢:       {r['max_time']:.2f}s")


def main():
    parser = argparse.ArgumentParser(description="ERNIE-Image use_pe 对比基准测试")
    parser.add_argument("--prompt-file", type=str, default="prompt.txt", help="prompt 文件路径")
    parser.add_argument("--server-url", type=str, default="http://0.0.0.0:8842", help="服务地址")
    parser.add_argument("--steps", type=int, default=50, help="推理步数")
    parser.add_argument("--width", type=int, default=1024, help="图像宽度")
    parser.add_argument("--height", type=int, default=1024, help="图像高度")
    parser.add_argument("--guidance-scale", type=float, default=5.0, help="Guidance scale")
    parser.add_argument("--seed", type=int, default=43, help="随机种子")
    parser.add_argument("--timeout", type=int, default=600, help="单次请求超时(秒)")
    parser.add_argument("--concurrency", type=int, default=1, help="并发数")
    args = parser.parse_args()

    prompts = load_prompts(args.prompt_file)

    # 先跑 use_pe=True
    result_true = run_benchmark(
        prompts, args.server_url, use_pe=True,
        steps=args.steps, width=args.width, height=args.height,
        guidance_scale=args.guidance_scale, seed=args.seed,
        timeout=args.timeout, concurrency=args.concurrency,
    )
    print_result(result_true)

    # 再跑 use_pe=False
    result_false = run_benchmark(
        prompts, args.server_url, use_pe=False,
        steps=args.steps, width=args.width, height=args.height,
        guidance_scale=args.guidance_scale, seed=args.seed,
        timeout=args.timeout, concurrency=args.concurrency,
    )
    print_result(result_false)

    # 汇总对比
    print(f"\n{'='*60}")
    print("对比汇总")
    print(f"{'='*60}")
    print(f"{'指标':<30} {'use_pe=False':>15} {'use_pe=True':>15}")
    print(f"{'-'*60}")
    print(f"{'成功请求数':<30} {result_false['success']:>15} {result_true['success']:>15}")
    print(f"{'失败请求数':<30} {result_false['fail']:>15} {result_true['fail']:>15}")
    print(f"{'实际耗时(Wall Time)':<30} {result_false['wall_time']:>14.2f}s {result_true['wall_time']:>14.2f}s")
    print(f"{'成功请求平均耗时':<30} {result_false['avg_success_time']:>14.2f}s {result_true['avg_success_time']:>14.2f}s")
    print(f"{'最快请求':<30} {result_false['min_time']:>14.2f}s {result_true['min_time']:>14.2f}s")
    print(f"{'最慢请求':<30} {result_false['max_time']:>14.2f}s {result_true['max_time']:>14.2f}s")

    # 保存结果到 JSON
    output = {
        "params": {
            "server_url": args.server_url,
            "steps": args.steps,
            "size": f"{args.width}x{args.height}",
            "guidance_scale": args.guidance_scale,
            "seed": args.seed,
            "timeout": args.timeout,
            "concurrency": args.concurrency,
        },
        "use_pe_false": result_false,
        "use_pe_true": result_true,
    }
    out_file = Path(f"benchmark_result_{time.strftime('%Y%m%d_%H%M%S')}.json")
    out_file.write_text(json.dumps(output, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"\n结果已保存: {out_file}")


if __name__ == "__main__":
    main()