tests/title_styler_regression/run.py

"""
title_styler 回归测试集 driver。

特性:
  - 不调 LLM (use_llm=False, template_name= 指定) → 完全 deterministic
  - 一次跑遍 PROBLEMATIC_CASES + FULL_TEMPLATE_TEST_CASES 30+ 个 case
  - 同时输出 SVG (styler 直接产物) 和 PNG (Selenium 渲染, overlap-fix 关闭)
  - PNG 用同一渲染管线 (modules.infographics_generator.screenshot_utils) → 跟最终产物一致
  - 输出按 run_name 分桶，方便 baseline / after_fix 对比

用法:
    PYTHONPATH=. python tests/title_styler_regression/run.py --run baseline
    # 改完代码再跑:
    PYTHONPATH=. python tests/title_styler_regression/run.py --run after_fix
    # 对比:
    PYTHONPATH=. python tests/title_styler_regression/run.py --diff baseline after_fix

输出:
    tests/title_styler_regression/results/<run>/<case_id>/{title.svg, title.png, meta.json}
    tests/title_styler_regression/results/<run>/summary.json
"""

import argparse
import json
import os
import re
import shutil
import sys
import time

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
sys.path.insert(0, ROOT)
os.chdir(ROOT)

from modules.title_styler.test_cases import PROBLEMATIC_CASES, FULL_TEMPLATE_TEST_CASES
from modules.title_styler.infographic_title_generator import generate_title
from modules.infographics_generator.screenshot_utils import get_driver, take_screenshot
from modules.infographics_generator.parse_utils import convert_svg_to_html

RESULTS_ROOT = os.path.join("tests", "title_styler_regression", "results")


def _safe_id(s: str) -> str:
    return re.sub(r"[^a-zA-Z0-9._-]+", "_", s)[:80]


def _font_size_histogram(svg_text: str) -> dict:
    hist = {}
    for s in re.findall(r"font-size[=:]\s*[\"']?(\d+(?:\.\d+)?)", svg_text):
        hist[s] = hist.get(s, 0) + 1
    return hist


def collect_cases():
    """合并去重；每个 case 用 (test_id + template_name + 前 30 char title) 作为 key。"""
    seen = set()
    out = []
    for src_name, src in [
        ("problematic", PROBLEMATIC_CASES),
        ("full", FULL_TEMPLATE_TEST_CASES),
    ]:
        for c in src:
            key = (c["test_id"], c["template_name"], c["title_segments"][:30])
            if key in seen:
                continue
            seen.add(key)
            c2 = dict(c)
            c2["__source__"] = src_name
            out.append(c2)
    return out


def render_one(driver, case, run_dir):
    case_id = _safe_id(f'{case["template_name"]}__{case["test_id"]}')
    case_dir = os.path.join(run_dir, case_id)
    os.makedirs(case_dir, exist_ok=True)

    t0 = time.time()
    results = generate_title(
        title=case["title_segments"],
        description=case["description"],
        primary_color=case["primary_color"],
        background_color=case["background_color"],
        template_name=case["template_name"],
        max_width=900,
        use_llm=False,
    )
    styler_ms = (time.time() - t0) * 1000

    if not results:
        print(f"  FAIL no-svg: {case_id}", flush=True)
        return {
            "case_id": case_id,
            "ok": False,
            "reason": "generate_title returned empty",
        }

    r = results[0]
    svg_text = r["svg"]
    svg_path = os.path.join(case_dir, "title.svg")
    with open(svg_path, "w", encoding="utf-8") as f:
        f.write(svg_text)

    html_path = os.path.join(case_dir, "render.html")
    convert_svg_to_html(svg_path, html_path)
    t1 = time.time()
    take_screenshot(driver, html_path, resolve_text_overlap=False)
    render_ms = (time.time() - t1) * 1000
    rendered_png = html_path.replace(".html", ".png")
    final_png = os.path.join(case_dir, "title.png")
    shutil.move(rendered_png, final_png)

    meta = {
        "case_id": case_id,
        "ok": True,
        "test_id": case["test_id"],
        "template_name": case["template_name"],
        "title_segments": case["title_segments"],
        "description": case["description"],
        "primary_color": case["primary_color"],
        "background_color": case["background_color"],
        "styler_ms": round(styler_ms, 1),
        "render_ms": round(render_ms, 1),
        "svg_width": r.get("width"),
        "svg_height": r.get("height"),
        "font_size_hist": _font_size_histogram(svg_text),
        "source_suite": case["__source__"],
    }
    with open(os.path.join(case_dir, "meta.json"), "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2, ensure_ascii=False)
    print(
        f"  OK  {case_id:60s}  {r['width']}x{r['height']}px  "
        f"styler={styler_ms:.0f}ms render={render_ms:.0f}ms",
        flush=True,
    )
    return meta


def run(run_name: str, only: str = None):
    cases = collect_cases()
    if only:
        cases = [c for c in cases if only in c["template_name"] or only in c["test_id"]]
    if not cases:
        print("no cases matched", flush=True)
        return

    run_dir = os.path.join(RESULTS_ROOT, run_name)
    if os.path.exists(run_dir):
        print(f"clearing {run_dir}", flush=True)
        shutil.rmtree(run_dir)
    os.makedirs(run_dir, exist_ok=True)

    print(f"=== run '{run_name}' on {len(cases)} cases ===", flush=True)
    driver = get_driver()
    metas = []
    for i, case in enumerate(cases, 1):
        print(f"[{i}/{len(cases)}]", flush=True)
        meta = render_one(driver, case, run_dir)
        metas.append(meta)
    driver.quit()

    summary = {
        "run_name": run_name,
        "n_cases": len(metas),
        "n_ok": sum(1 for m in metas if m.get("ok")),
        "n_fail": sum(1 for m in metas if not m.get("ok")),
        "cases": metas,
    }
    with open(os.path.join(run_dir, "summary.json"), "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2, ensure_ascii=False)
    print(
        f"\n=== run '{run_name}' done. ok={summary['n_ok']} fail={summary['n_fail']} ===",
        flush=True,
    )
    print(f"results in: {run_dir}/")


def diff(run_a: str, run_b: str):
    dir_a = os.path.join(RESULTS_ROOT, run_a)
    dir_b = os.path.join(RESULTS_ROOT, run_b)
    sa = json.load(open(os.path.join(dir_a, "summary.json")))
    sb = json.load(open(os.path.join(dir_b, "summary.json")))
    map_a = {c["case_id"]: c for c in sa["cases"]}
    map_b = {c["case_id"]: c for c in sb["cases"]}
    all_ids = sorted(set(map_a) | set(map_b))
    print(f"=== diff {run_a} vs {run_b} ===")
    print(
        f"{'case':70s}  {'svg WxH (a)':>18s}  {'svg WxH (b)':>18s}  {'dh':>5s}  {'fonts a/b'}"
    )
    for cid in all_ids:
        a, b = map_a.get(cid), map_b.get(cid)
        if a and a.get("ok") and b and b.get("ok"):
            wa, ha = a["svg_width"], a["svg_height"]
            wb, hb = b["svg_width"], b["svg_height"]
            dh = hb - ha
            fa = sorted(a["font_size_hist"].keys(), key=float)
            fb = sorted(b["font_size_hist"].keys(), key=float)
            flag = "" if dh == 0 else (" *" if abs(dh) >= 4 else "")
            print(
                f"{cid:70s}  {wa}x{ha:>4d}     {wb}x{hb:>4d}    {dh:+5d}  "
                f"{','.join(fa)} / {','.join(fb)}{flag}"
            )
        else:
            print(f"{cid:70s}  (missing or failed in one run)")


def main():
    ap = argparse.ArgumentParser()
    g = ap.add_mutually_exclusive_group(required=True)
    g.add_argument("--run", help="run a regression pass with this name (e.g. baseline, after_fix)")
    g.add_argument("--diff", nargs=2, metavar=("RUN_A", "RUN_B"), help="diff two runs")
    ap.add_argument("--only", help="filter cases whose template_name or test_id contains this substring")
    args = ap.parse_args()
    if args.run:
        run(args.run, only=args.only)
    else:
        diff(args.diff[0], args.diff[1])


if __name__ == "__main__":
    main()