scripts/run_quality_check.py

"""
Per-template quality-check driver.

Reads a plan JSON (produced by pick_smoke_templates.py), runs
infographics_generator.process(input, output, "", "", chart_name) for
each (template, data) pair in a process pool, and incrementally writes:

  - <output-dir>/_tasks.csv       per-task report (one row per (tpl, data))
  - <output-dir>/_summary.csv     per-template rollup, worst-first
  - <output-dir>/PROGRESS.md      live status, ETA, per-tpl table

Resume is supported: if --resume is passed and _tasks.csv already has rows,
those (chart_name, input) pairs are skipped on this run.

Usage:
    python scripts/run_quality_check.py \\
        --plan scripts/_smoke_templates.json \\
        --output-dir output/quality_check \\
        --threads 8
"""
import argparse
import json
import os
import sys
import time
import traceback
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from scripts._progress import ProgressTracker


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--plan", default="scripts/_smoke_templates.json")
    p.add_argument("--data-dir", nargs="+", default=None,
                   help="Override data dir(s) (defaults to plan.data_dirs / data_dir)")
    p.add_argument("--output-dir", default="output/quality_check")
    p.add_argument("--threads", type=int, default=8)
    p.add_argument("--chrome-path", default="/usr/bin/google-chrome",
                   help="Chrome executable for puppeteer (sets PUPPETEER_EXECUTABLE_PATH)")
    p.add_argument("--limit", type=int, default=None,
                   help="Optional cap on total tasks (for ultra-quick smoke)")
    p.add_argument("--resume", action="store_true",
                   help="Read existing _tasks.csv and skip already-done jobs")
    p.add_argument("--match-csv", default="scripts/_template_match.csv",
                   help="Used to list skipped templates in PROGRESS.md")
    p.add_argument("--min-data-for-skipped", type=int, default=5,
                   help="Show skipped templates whose match count is below this")
    return p.parse_args()


def _worker_init(chrome_path: str):
    # Each worker process inherits its own env; set puppeteer's chrome here.
    os.environ.setdefault("PUPPETEER_EXECUTABLE_PATH", chrome_path)
    # Reduce torch / faiss threading noise inside workers.
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("MKL_NUM_THREADS", "1")


# FALLBACK_MARKER mirrors what html_to_svg.py writes when puppeteer cannot
# extract a real SVG from the DOM. Same string used by the loader to flag it.
FALLBACK_MARKER = "This is a fallback SVG using a PNG screenshot"


def _is_fallback_svg(svg_path: str) -> bool:
    if not os.path.exists(svg_path):
        return False
    with open(svg_path, "r", encoding="utf-8", errors="ignore") as fh:
        head = fh.read(4096)
    return FALLBACK_MARKER in head


def _count_svg_elements(svg_path: str):
    """Cheap structural metric: rough count of drawing/text tags.

    Doesn't fully parse; just regex-counts common shape and text tags.
    Good enough as a sanity signal for "is this SVG basically empty?".
    """
    if not os.path.exists(svg_path):
        return 0, 0
    import re
    with open(svg_path, "r", encoding="utf-8", errors="ignore") as fh:
        content = fh.read()
    shape_tags = ("path", "rect", "circle", "line", "polygon", "polyline", "ellipse", "use")
    n_shapes = sum(len(re.findall(rf"<{t}\b", content)) for t in shape_tags)
    n_text = len(re.findall(r"<text\b", content))
    return n_shapes, n_text


def _run_one(args_tuple):
    """Run one (template, data_file) job. Returns a dict of metrics."""
    chart_name, input_path, data_key, template_out_dir, base_url, api_key = args_tuple

    # Lazy import inside the worker (after worker_init sets env).
    from modules.infographics_generator.infographics_generator import process

    input_basename = os.path.basename(input_path)
    # process() expects output like "<dir>/<basename>" - it will create
    # a "<timestamp>_<chart_name>_<basename_stem>/" subfolder next to it
    # AND write a "<timestamp>_<chart_name>_<basename_stem>.svg" SVG.
    output_path = os.path.join(template_out_dir, input_basename)

    t0 = time.time()
    ok = False
    err_msg = ""
    try:
        ok = process(
            input=input_path,
            output=output_path,
            base_url=base_url,
            api_key=api_key,
            chart_name=chart_name,
        )
    except BaseException as e:
        err_msg = f"{type(e).__name__}: {e}"
        tb = traceback.format_exc()
        err_msg = (err_msg + " | " + tb.splitlines()[-1])[:300]
    elapsed = time.time() - t0

    # Find the final SVG that was produced (most-recent matching file).
    stem = os.path.splitext(input_basename)[0]
    final_svg = None
    final_svg_size = 0
    chart_svg_path = None
    chart_svg_is_fallback = False
    n_paths = 0  # actually n_shapes; legacy local name
    n_text = 0
    if os.path.isdir(template_out_dir):
        candidates = [
            f for f in os.listdir(template_out_dir)
            if f.endswith(f"_{chart_name}_{stem}.svg")
        ]
        if candidates:
            candidates.sort()
            final_svg = os.path.join(template_out_dir, candidates[-1])
            final_svg_size = os.path.getsize(final_svg)
            n_paths, n_text = _count_svg_elements(final_svg)

        # The per-chart subfolder also contains the raw chart SVG.
        sub_candidates = [
            d for d in os.listdir(template_out_dir)
            if os.path.isdir(os.path.join(template_out_dir, d))
            and d.endswith(f"_{chart_name}_{stem}")
        ]
        if sub_candidates:
            sub_candidates.sort()
            chart_subdir = os.path.join(template_out_dir, sub_candidates[-1])
            chart_svg_path = os.path.join(chart_subdir, "chart.svg")
            chart_svg_is_fallback = _is_fallback_svg(chart_svg_path)

    return {
        "chart_name": chart_name,
        # Use the plan-level key (may include "<dir>/<filename>") so resume
        # matches across multi-pool plans. Falls back to bare basename when
        # the caller didn't provide one.
        "input": data_key or input_basename,
        "ok": bool(ok),
        "elapsed_s": round(elapsed, 2),
        "final_svg": final_svg or "",
        "final_svg_bytes": final_svg_size,
        "chart_svg": chart_svg_path or "",
        "chart_svg_fallback_png": chart_svg_is_fallback,
        "n_shapes": n_paths,
        "n_text": n_text,
        "err": err_msg,
    }


def main():
    args = parse_args()
    with open(args.plan, "r") as fh:
        plan = json.load(fh)

    # Resolve data dirs: CLI override > plan.data_dirs > legacy plan.data_dir.
    if args.data_dir:
        data_dirs = [Path(p) for p in args.data_dir]
    elif "data_dirs" in plan:
        data_dirs = [Path(p) for p in plan["data_dirs"]]
    else:
        data_dirs = [Path(plan["data_dir"])]
    for d in data_dirs:
        if not d.is_dir():
            raise SystemExit(f"data dir is not a directory: {d}")
    # Map "<dir-name>/<filename>" -> absolute path, falling back to
    # bare basenames when an old plan was generated against a single dir.
    file_lookup: dict[str, str] = {}
    for d in data_dirs:
        for f in d.glob("*.json"):
            file_lookup.setdefault(f.name, str(f))
            file_lookup[f"{d.name}/{f.name}"] = str(f)

    def resolve_data(rel: str) -> str:
        if rel in file_lookup:
            return file_lookup[rel]
        bn = os.path.basename(rel)
        if bn in file_lookup:
            return file_lookup[bn]
        raise FileNotFoundError(
            f"plan references data file '{rel}' not found in any of: "
            + ", ".join(str(d) for d in data_dirs)
        )

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Ensure parent process also has chrome path set so any in-proc work works.
    os.environ.setdefault("PUPPETEER_EXECUTABLE_PATH", args.chrome_path)

    tracker = ProgressTracker(
        plan=plan,
        output_dir=output_dir,
        match_csv_path=Path(args.match_csv) if args.match_csv else None,
        min_data_for_plan=args.min_data_for_skipped,
    )
    tracker.open_csv(resume=args.resume)

    # Build job list, skipping already-done ones when --resume.
    jobs = []
    skipped_resume = 0
    for tpl in plan["templates"]:
        chart_name = tpl["chart_name"]
        template_out_dir = output_dir / chart_name
        template_out_dir.mkdir(parents=True, exist_ok=True)
        for data_basename in tpl["picked_data_files"]:
            if tracker.already_done(chart_name, data_basename):
                skipped_resume += 1
                continue
            input_path = resolve_data(data_basename)
            jobs.append((chart_name, input_path, data_basename,
                         str(template_out_dir), "", ""))

    if args.limit:
        jobs = jobs[: args.limit]

    print(
        f"Plan: templates={len(plan['templates'])}  "
        f"total-tasks={tracker.total_tasks}  resume-skipped={skipped_resume}  "
        f"to-run={len(jobs)}  threads={args.threads}"
    )

    # Persist an initial PROGRESS.md so the file exists right away.
    tracker.flush()

    t_start = time.time()
    with ProcessPoolExecutor(
        max_workers=args.threads,
        initializer=_worker_init,
        initargs=(args.chrome_path,),
    ) as ex:
        futures = [ex.submit(_run_one, j) for j in jobs]
        done = 0
        for fut in as_completed(futures):
            res = fut.result()
            tracker.add(res)
            done += 1
            elapsed = time.time() - t_start
            rate = done / max(elapsed, 1e-6)
            remaining = (len(jobs) - done) / max(rate, 1e-6)
            ok_str = "T" if res["ok"] else "F"
            print(
                f"[{done:4d}/{len(jobs):4d}] "
                f"ok={ok_str} "
                f"{res['chart_name']:38s} "
                f"data={res['input'][:32]:32s} "
                f"t={res['elapsed_s']:5.1f}s "
                f"size={res['final_svg_bytes']/1024:5.0f}KB "
                f"shapes={res['n_shapes']:4d} "
                f"| eta {remaining/60:.1f}min",
                flush=True,
            )

    tracker.close()
    total_t = time.time() - t_start
    print()
    print(f"Wrote per-task report:      {tracker.tasks_csv}")
    print(f"Wrote per-template summary: {tracker.summary_csv}")
    print(f"Wrote progress doc:         {tracker.progress_md}")
    print(f"Total time this session: {total_t/60:.1f} min")

    # Always (re)generate the HTML preview at the end of a session so the
    # user has a single browsable artifact. Failure to build the preview
    # shouldn't fail the whole run.
    preview_path = output_dir / "preview.html"
    try:
        from scripts.build_quality_preview import build_preview
        build_preview(output_dir, preview_path)
        print(f"Wrote HTML preview:        {preview_path}")
    except Exception as e:
        print(f"WARN: failed to build HTML preview: {e}")


if __name__ == "__main__":
    main()