scripts/generate_template_samples.py

"""
为每个 chart template 随机生成若干完整 infographic 结果。

默认会跑完整流程:
  preprocess -> datafact_generator -> title_generator -> color_recommender
  -> image_recommender -> infographics_generator

示例:
  PYTHONPATH=. python scripts/generate_template_samples.py --samples-per-template 10 --output-png

输出结构:
  output/chart_template_samples/<timestamp>/
    <chart_template_name>/
      sample_00/
      sample_01/
      ...

为了覆盖所有 chart template，默认忽略 allowed_chart_types.json 白名单。
如果只想跑当前白名单里的 chart_type，加 --respect-allowed-chart-types。
"""

import argparse
import csv
import json
import os
import random
import re
import shutil
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path
from collections import defaultdict


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

from config import api_key as CFG_API_KEY
from config import base_url as CFG_BASE_URL
from config import data_resource_dirs as CFG_DATA_DIRS
from modules.chart_engine.template.template_registry import scan_templates
from modules.infographics_generator.template_utils import (
    analyze_templates,
    check_template_compatibility,
)
from pipeline import run_single_file


DEFAULT_ENGINES = ["d3-js", "echarts-js", "echarts_py"]
PIPELINE_MODULES = ["all", "infographics_generator"]
COMPAT_CACHE_VERSION = 1
SHAPE_TAGS = ("path", "rect", "circle", "line", "polygon", "polyline", "ellipse", "use")
IMAGE_TAGS = ("image",)
FALLBACK_MARKER = "This is a fallback SVG using a PNG screenshot"


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--data",
        nargs="+",
        default=None,
        help="输入 JSON 数据目录；默认使用 config.data_resource_dirs",
    )
    parser.add_argument(
        "--output",
        default=None,
        help="任务输出父目录；脚本会在下面创建 <timestamp>/，默认 output/chart_template_samples",
    )
    parser.add_argument(
        "--samples-per-template",
        type=int,
        default=10,
        help="每个 chart template 生成多少个结果",
    )
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument(
        "--engines",
        nargs="+",
        default=DEFAULT_ENGINES,
        help="要覆盖的 chart engine",
    )
    parser.add_argument(
        "--scan-limit",
        type=int,
        default=0,
        help="候选数据扫描上限；0 表示扫描全部数据",
    )
    parser.add_argument(
        "--limit-templates",
        type=int,
        default=0,
        help="仅跑前 N 个模板，调试用；0 表示不限制",
    )
    parser.add_argument(
        "--output-png",
        action="store_true",
        help="同时生成 PNG，便于人工检查",
    )
    parser.add_argument(
        "--chart-only",
        action="store_true",
        help="只生成 chart SVG/PNG",
    )
    parser.add_argument(
        "--slot-polish-after-chart",
        action="store_true",
        help="先以 chart-only 渲染 template，再直接调用 full_image_polisher 的 slot mode",
    )
    parser.add_argument(
        "--slot-polish-dry-run",
        action="store_true",
        help="只生成 full_image_polisher slot mask/prompt/manifest，不调用图像模型",
    )
    parser.add_argument(
        "--planned-slot-polish",
        action="store_true",
        help="先清理 template 自带 title/image，再规划 editable slots，最后调用 full_image_polisher slot mode",
    )
    parser.add_argument(
        "--planned-slot-dry-run",
        action="store_true",
        help="只生成 planned slot canvas/mask/prompt/manifest，不调用图像模型",
    )
    parser.add_argument(
        "--planned-slot-disallow-chart-overlap",
        action="store_true",
        help="规划 slots 时不允许和 chart bbox 相交；默认允许相交",
    )
    parser.add_argument(
        "--planned-slot-polisher-base-url",
        default=None,
        help=(
            "Override base_url for planned-slot full_image_polisher. "
            "Use 'openai_default' to ignore config.base_url and call the official OpenAI endpoint."
        ),
    )
    parser.add_argument("--slot-polisher-backend", choices=("auto", "openai", "pinco"), default="auto")
    parser.add_argument("--planned-slot-polisher-backend", choices=("auto", "openai", "pinco"), default="auto")
    parser.add_argument(
        "--pinco-command",
        default=None,
        help=(
            "Command template for local Pinco inference. Placeholders: {input}, {mask}, "
            "{foreground}, {prompt_file}, {output}, {model}, {width}, {height}."
        ),
    )
    parser.add_argument("--pinco-url", default=None, help="HTTP endpoint for a Pinco inpainting service.")
    parser.add_argument("--pinco-timeout", type=int, default=600)
    parser.add_argument(
        "--png-longest-side",
        type=int,
        default=1600,
        help="PNG 最长边像素；批量审查默认 1600，线上默认配置目前是 3860",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=1,
        help="并行生成 worker 数；1 表示串行",
    )
    parser.add_argument(
        "--scan-workers",
        type=int,
        default=1,
        help="并行兼容性扫描 worker 数；1 表示串行",
    )
    parser.add_argument(
        "--compat-cache",
        default="output/chart_template_samples/template_compat_cache.json",
        help="template -> compatible data 扫描缓存路径",
    )
    parser.add_argument(
        "--cache-candidates-per-template",
        type=int,
        default=50,
        help="每个 template 最多缓存多少个兼容数据路径",
    )
    parser.add_argument(
        "--rebuild-compat-cache",
        action="store_true",
        help="忽略已有兼容数据缓存，重新扫描",
    )
    parser.add_argument(
        "--respect-allowed-chart-types",
        action="store_true",
        help="尊重 allowed_chart_types.json；默认忽略白名单以覆盖所有模板",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="只生成候选匹配报告，不执行 pipeline",
    )
    parser.add_argument(
        "--clean",
        action="store_true",
        help="如果本次 timestamp 任务目录已存在，先删除再生成",
    )
    parser.add_argument(
        "--resume-root",
        default=None,
        help="继续已有任务目录，跳过 manifest 中已成功的样本",
    )
    return parser.parse_args()


def configure_chart_type_filter(respect_allowed_chart_types: bool):
    if respect_allowed_chart_types:
        return
    os.environ["ALLOWED_CHART_TYPES_FILE"] = str(
        ROOT / "tmp" / "__ignore_allowed_chart_types_for_template_samples__.json"
    )


def safe_slug(value: str, max_len: int = 120) -> str:
    slug = re.sub(r"[^a-zA-Z0-9._-]+", "_", value).strip("_")
    return slug[:max_len] or "template"


def read_json(path: Path):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def write_json(path: Path, data):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)


def append_jsonl(path: Path, record: dict):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")


def cache_meta(data_dirs, engines, respect_allowed_chart_types, scan_limit, template_keys):
    return {
        "version": COMPAT_CACHE_VERSION,
        "data_dirs": [str(Path(p)) for p in data_dirs],
        "engines": list(engines),
        "respect_allowed_chart_types": bool(respect_allowed_chart_types),
        "scan_limit": int(scan_limit),
        "template_keys": list(template_keys),
    }


def load_compatibility_cache(path: Path, expected_meta: dict):
    if not path.is_file():
        return None
    cached = read_json(path)
    cached_meta = cached.get("meta") or {}
    same_scan_scope = all(
        cached_meta.get(key) == expected_meta.get(key)
        for key in ("version", "data_dirs", "respect_allowed_chart_types", "scan_limit")
    )
    cached_template_keys = set(cached_meta.get("template_keys") or [])
    expected_template_keys = set(expected_meta.get("template_keys") or [])
    if not same_scan_scope or not expected_template_keys.issubset(cached_template_keys):
        return None
    candidates = {
        key: [Path(p) for p in paths]
        for key, paths in cached.get("candidates", {}).items()
        if key in expected_template_keys
    }
    return candidates


def save_compatibility_cache(path: Path, meta: dict, candidates_by_template):
    payload = {
        "meta": meta,
        "candidates": {
            key: [str(p) for p in paths]
            for key, paths in candidates_by_template.items()
        },
    }
    write_json(path, payload)


def make_output_root(output_parent, timestamp: str) -> Path:
    parent = Path(output_parent or "output/chart_template_samples")
    return parent / timestamp


def collect_input_files(data_dirs, rng, scan_limit: int):
    files = []
    for data_dir in data_dirs:
        path = Path(data_dir)
        if not path.is_dir():
            raise SystemExit(f"数据目录不存在: {path}")
        files.extend(sorted(path.glob("*.json")))

    rng.shuffle(files)
    if scan_limit and scan_limit < len(files):
        files = files[:scan_limit]
    if not files:
        raise SystemExit("没有找到可用 JSON 数据")
    return files


def enumerate_template_keys(templates, engines):
    keys = []
    for engine, chart_types in templates.items():
        if engine not in engines:
            continue
        for chart_type, chart_names in chart_types.items():
            for chart_name in chart_names:
                if "base" in chart_name:
                    continue
                if engine == "vegalite_py":
                    continue
                keys.append(f"{engine}/{chart_type}/{chart_name}")
    return sorted(keys)


def scan_file_compatibility(path, templates, target_keys, target_chart_names, use_targeted_check):
    data = read_json(path)
    data["name"] = str(path)
    matched_keys = []
    if use_targeted_check:
        for chart_name in target_chart_names:
            compatible = check_template_compatibility(data, templates, chart_name)
            matched_keys.extend(
                template_key
                for template_key, _ordered_fields in compatible
                if template_key in target_keys
            )
    else:
        compatible = check_template_compatibility(data, templates, None)
        matched_keys.extend(
            template_key
            for template_key, _ordered_fields in compatible
            if template_key in target_keys
        )
    return path, matched_keys


def add_compatibility_result(by_template, path, matched_keys, min_candidates: int):
    for template_key in matched_keys:
        if len(by_template[template_key]) < min_candidates:
            by_template[template_key].append(path)


def compatibility_scan_complete(by_template, template_keys, min_candidates: int):
    return all(len(by_template[key]) >= min_candidates for key in template_keys)


def build_compatibility_index_serial(
    input_files,
    templates,
    template_keys,
    min_candidates: int,
    target_keys,
    target_chart_names,
    use_targeted_check: bool,
):
    by_template = defaultdict(list)
    for index, path in enumerate(input_files, 1):
        path, matched_keys = scan_file_compatibility(
            path,
            templates,
            target_keys,
            target_chart_names,
            use_targeted_check,
        )
        add_compatibility_result(by_template, path, matched_keys, min_candidates)
        if index % 500 == 0:
            print(f"  scanned {index}/{len(input_files)} data files", flush=True)
        if compatibility_scan_complete(by_template, template_keys, min_candidates):
            print(f"  found {min_candidates} candidates for every selected template", flush=True)
            break
    return by_template


def build_compatibility_index_parallel(
    input_files,
    templates,
    template_keys,
    min_candidates: int,
    target_keys,
    target_chart_names,
    use_targeted_check: bool,
    scan_workers: int,
):
    by_template = defaultdict(list)
    next_submit_index = 0
    next_result_index = 0
    results = {}
    pending = {}
    max_pending = max(scan_workers * 4, scan_workers)

    with ThreadPoolExecutor(max_workers=scan_workers) as executor:
        while next_submit_index < len(input_files) and len(pending) < max_pending:
            path = input_files[next_submit_index]
            pending[
                executor.submit(
                    scan_file_compatibility,
                    path,
                    templates,
                    target_keys,
                    target_chart_names,
                    use_targeted_check,
                )
            ] = next_submit_index
            next_submit_index += 1

        while pending:
            for future in as_completed(pending):
                result_index = pending.pop(future)
                results[result_index] = future.result()
                break

            while next_result_index in results:
                path, matched_keys = results.pop(next_result_index)
                add_compatibility_result(by_template, path, matched_keys, min_candidates)
                scanned = next_result_index + 1
                if scanned % 500 == 0:
                    print(f"  scanned {scanned}/{len(input_files)} data files", flush=True)
                next_result_index += 1
                if compatibility_scan_complete(by_template, template_keys, min_candidates):
                    print(
                        f"  found {min_candidates} candidates for every selected template",
                        flush=True,
                    )
                    for future in pending:
                        future.cancel()
                    return by_template

            while next_submit_index < len(input_files) and len(pending) < max_pending:
                path = input_files[next_submit_index]
                pending[
                    executor.submit(
                        scan_file_compatibility,
                        path,
                        templates,
                        target_keys,
                        target_chart_names,
                        use_targeted_check,
                    )
                ] = next_submit_index
                next_submit_index += 1

    return by_template


def build_compatibility_index(input_files, templates, template_keys, min_candidates: int, scan_workers: int):
    target_keys = set(template_keys)
    target_chart_names = sorted({key.split("/")[-1] for key in template_keys})
    use_targeted_check = len(template_keys) <= 20
    if scan_workers <= 1:
        return build_compatibility_index_serial(
            input_files,
            templates,
            template_keys,
            min_candidates,
            target_keys,
            target_chart_names,
            use_targeted_check,
        )
    return build_compatibility_index_parallel(
        input_files,
        templates,
        template_keys,
        min_candidates,
        target_keys,
        target_chart_names,
        use_targeted_check,
        scan_workers,
    )


def choose_samples(candidates, count: int, rng):
    if len(candidates) >= count:
        return rng.sample(candidates, count)
    return [rng.choice(candidates) for _ in range(count)]


def write_candidate_report(path: Path, template_keys, candidates_by_template):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        writer.writerow(["template_key", "num_compatible_data"])
        for template_key in template_keys:
            writer.writerow([template_key, len(candidates_by_template.get(template_key, []))])


def newest_root_svg(sample_dir: Path):
    svgs = sorted(
        [p for p in sample_dir.glob("*.svg") if p.is_file()],
        key=lambda p: p.stat().st_mtime,
    )
    return str(svgs[-1]) if svgs else None


def newest_root_png(sample_dir: Path):
    pngs = sorted(
        [p for p in sample_dir.glob("*.png") if p.is_file()],
        key=lambda p: p.stat().st_mtime,
    )
    return str(pngs[-1]) if pngs else None


def nested_chart_svg(sample_dir: Path):
    candidates = sorted(
        [p for p in sample_dir.glob("*/chart.svg") if p.is_file()],
        key=lambda p: p.stat().st_mtime,
    )
    return str(candidates[-1]) if candidates else None


def file_metrics(path):
    if path is not None:
        path = Path(path)
    if path is None or not path.is_file():
        return {
            "exists": False,
            "bytes": 0,
            "fallback": False,
            "n_shapes": 0,
            "n_text": 0,
            "n_images": 0,
            "n_visible": 0,
            "n_elements": 0,
            "empty": True,
        }
    text = path.read_text(encoding="utf-8", errors="ignore")
    n_shapes = sum(text.count(f"<{tag}") for tag in SHAPE_TAGS)
    n_text = text.count("<text")
    n_images = sum(text.count(f"<{tag}") for tag in IMAGE_TAGS)
    n_visible = n_shapes + n_text + n_images
    return {
        "exists": True,
        "bytes": path.stat().st_size,
        "fallback": FALLBACK_MARKER in text,
        "n_shapes": n_shapes,
        "n_text": n_text,
        "n_images": n_images,
        "n_visible": n_visible,
        "n_elements": text.count("<"),
        "empty": n_visible == 0,
    }


def resolve_polisher_base_url(value: str | None):
    if value is None:
        return CFG_BASE_URL
    text = str(value).strip()
    if text.lower() in {"", "none", "null", "openai_default", "official_openai"}:
        return None
    return text


def run_sample(
    template_key: str,
    data_path: Path,
    sample_dir: Path,
    output_png: bool,
    chart_only: bool,
    slot_polish_after_chart: bool = False,
    slot_polish_dry_run: bool = False,
    planned_slot_polish: bool = False,
    planned_slot_dry_run: bool = False,
    planned_slot_allow_chart_overlap: bool = True,
    planned_slot_polisher_base_url: str | None = None,
    slot_polisher_backend: str = "auto",
    planned_slot_polisher_backend: str = "auto",
    pinco_command: str | None = None,
    pinco_url: str | None = None,
    pinco_timeout: int = 600,
):
    chart_name = template_key.split("/")[-1]
    sample_dir.mkdir(parents=True, exist_ok=True)
    output_stub = sample_dir / f"{safe_slug(chart_name)}_{sample_dir.name}"
    sample_start_time = time.time()
    render_chart_only = chart_only or slot_polish_after_chart or planned_slot_polish
    render_output_png = output_png or slot_polish_after_chart or planned_slot_polish
    modules_to_run = ["infographics_generator"] if render_chart_only else PIPELINE_MODULES
    pipeline_result = run_single_file(
        input_path=data_path,
        output_path=output_stub,
        temp_dir=sample_dir / "tmp",
        modules_to_run=modules_to_run,
        chart_name=template_key,
        chart_only=render_chart_only,
        output_png=render_output_png,
        return_timing=True,
    )
    final_svg = newest_root_svg(sample_dir)
    final_png = newest_root_png(sample_dir)
    chart_svg = nested_chart_svg(sample_dir)
    module_seconds = {}
    pipeline_total_seconds = None
    pipeline_success = pipeline_result is not False
    if isinstance(pipeline_result, dict):
        module_seconds = pipeline_result.get("module_seconds") or {}
        pipeline_total_seconds = pipeline_result.get("total_seconds")
        pipeline_success = bool(pipeline_result.get("success"))

    slot_polish_success = None
    slot_polished_png = ""
    slot_polish_manifest = ""
    slot_polish_error = ""
    planned_slot_success = None
    planned_slot_package_dir = ""
    planned_slot_plan = ""
    planned_slot_svg = ""
    planned_slot_png = ""
    planned_slot_reference_map = ""
    planned_slot_sanitized_svg = ""
    planned_slot_sanitized_png = ""
    planned_slot_polished_png = ""
    planned_slot_manifest = ""
    planned_slot_error = ""
    if slot_polish_after_chart and pipeline_success:
        try:
            if not final_png:
                raise FileNotFoundError("slot polish requires chart PNG; no root PNG was produced")
            slot_svg = chart_svg or final_svg
            if not slot_svg:
                raise FileNotFoundError("slot polish requires chart SVG; no SVG was produced")

            from modules.full_image_polisher.full_image_polisher import process as polish_process

            input_png_path = Path(final_png)
            slot_output = input_png_path.with_name(
                f"{input_png_path.stem}.slot_guided_polished.png"
            )
            step_start_time = time.time()
            slot_polish_success = bool(
                polish_process(
                    png_path=str(input_png_path),
                    svg_path=str(slot_svg),
                    output_png=str(slot_output),
                    mode="slot",
                    base_url=CFG_BASE_URL,
                    api_key=CFG_API_KEY,
                    dry_run=slot_polish_dry_run,
                    resize_to_input=True,
                    image_backend=slot_polisher_backend,
                    pinco_command=pinco_command,
                    pinco_url=pinco_url,
                    pinco_timeout=pinco_timeout,
                )
            )
            module_seconds["full_image_polisher.slot"] = time.time() - step_start_time
            manifest_path = slot_output.with_name(f"{slot_output.stem}.manifest.json")
            if slot_output.is_file():
                slot_polished_png = str(slot_output)
            if manifest_path.is_file():
                slot_polish_manifest = str(manifest_path)
        except Exception as exc:
            slot_polish_success = False
            slot_polish_error = f"{type(exc).__name__}: {exc}"

    if planned_slot_polish and pipeline_success:
        try:
            if not final_png:
                raise FileNotFoundError("planned slot polish requires chart PNG; no root PNG was produced")
            slot_svg = chart_svg or final_svg
            if not slot_svg:
                raise FileNotFoundError("planned slot polish requires chart SVG; no SVG was produced")

            from modules.full_image_polisher.full_image_polisher import (
                PLANNED_SLOT_GENERATION_PROMPT,
                process as polish_process,
            )
            from modules.slot_layout_planner.planner import build_planned_slot_package

            input_png_path = Path(final_png)
            package_dir = input_png_path.with_name(f"{input_png_path.stem}.planned_slot_work")
            planned_slot_package_dir = str(package_dir)
            step_start_time = time.time()
            package = build_planned_slot_package(
                chart_svg=Path(slot_svg),
                chart_png=input_png_path,
                data_json=Path(data_path),
                output_dir=package_dir,
                allow_chart_overlap=planned_slot_allow_chart_overlap,
            )
            module_seconds["slot_layout_planner"] = time.time() - step_start_time
            planned_slot_plan = str(package.slot_plan)
            planned_slot_svg = str(package.planned_svg)
            planned_slot_png = str(package.planned_png)
            planned_slot_reference_map = str(package.reference_map)
            planned_slot_sanitized_svg = str(package.sanitized_svg)
            planned_slot_sanitized_png = str(package.sanitized_png)

            planned_output = input_png_path.with_name(
                f"{input_png_path.stem}.planned_slot_polished.png"
            )
            step_start_time = time.time()
            planned_slot_success = bool(
                polish_process(
                    png_path=str(package.planned_png),
                    svg_path=str(package.planned_svg),
                    reference_map=str(package.reference_map),
                    data_json=str(data_path),
                    output_png=str(planned_output),
                    mode="slot",
                    base_url=resolve_polisher_base_url(planned_slot_polisher_base_url),
                    api_key=CFG_API_KEY,
                    dry_run=planned_slot_dry_run,
                    resize_to_input=True,
                    prompt=PLANNED_SLOT_GENERATION_PROMPT,
                    include_title_slots=True,
                    image_backend=planned_slot_polisher_backend,
                    pinco_command=pinco_command,
                    pinco_url=pinco_url,
                    pinco_timeout=pinco_timeout,
                )
            )
            module_seconds["full_image_polisher.planned_slot"] = time.time() - step_start_time
            manifest_path = planned_output.with_name(f"{planned_output.stem}.manifest.json")
            if planned_output.is_file():
                planned_slot_polished_png = str(planned_output)
            if manifest_path.is_file():
                planned_slot_manifest = str(manifest_path)
        except Exception as exc:
            planned_slot_success = False
            planned_slot_error = f"{type(exc).__name__}: {exc}"

    total_seconds = time.time() - sample_start_time
    final_metrics = file_metrics(final_svg)
    chart_metrics = file_metrics(chart_svg or final_svg)
    success = (
        pipeline_success
        and final_metrics["exists"]
        and chart_metrics["exists"]
        and not chart_metrics["fallback"]
        and not chart_metrics["empty"]
    )
    if slot_polish_after_chart:
        success = success and bool(slot_polish_success)
        if not slot_polish_dry_run:
            success = success and bool(slot_polished_png)
    if planned_slot_polish:
        success = success and bool(planned_slot_success)
        if not planned_slot_dry_run:
            success = success and bool(planned_slot_polished_png)
    return {
        "success": success,
        "pipeline_success": pipeline_success,
        "final_svg": final_svg,
        "final_png": final_png,
        "chart_svg": chart_svg,
        "slot_polish_after_chart": slot_polish_after_chart,
        "slot_polish_dry_run": slot_polish_dry_run,
        "slot_polish_success": slot_polish_success,
        "slot_polished_png": slot_polished_png,
        "slot_polish_manifest": slot_polish_manifest,
        "slot_polish_error": slot_polish_error,
        "planned_slot_polish": planned_slot_polish,
        "planned_slot_dry_run": planned_slot_dry_run,
        "planned_slot_allow_chart_overlap": planned_slot_allow_chart_overlap,
        "planned_slot_polisher_base_url": planned_slot_polisher_base_url or "",
        "slot_polisher_backend": slot_polisher_backend,
        "planned_slot_polisher_backend": planned_slot_polisher_backend,
        "pinco_command": pinco_command or "",
        "pinco_url": pinco_url or "",
        "pinco_timeout": pinco_timeout,
        "planned_slot_success": planned_slot_success,
        "planned_slot_package_dir": planned_slot_package_dir,
        "planned_slot_plan": planned_slot_plan,
        "planned_slot_svg": planned_slot_svg,
        "planned_slot_png": planned_slot_png,
        "planned_slot_reference_map": planned_slot_reference_map,
        "planned_slot_sanitized_svg": planned_slot_sanitized_svg,
        "planned_slot_sanitized_png": planned_slot_sanitized_png,
        "planned_slot_polished_png": planned_slot_polished_png,
        "planned_slot_manifest": planned_slot_manifest,
        "planned_slot_error": planned_slot_error,
        "final_svg_bytes": final_metrics["bytes"],
        "final_svg_visible": final_metrics["n_visible"],
        "chart_svg_bytes": chart_metrics["bytes"],
        "chart_shapes": chart_metrics["n_shapes"],
        "chart_text": chart_metrics["n_text"],
        "chart_images": chart_metrics["n_images"],
        "chart_visible": chart_metrics["n_visible"],
        "chart_fallback": chart_metrics["fallback"],
        "chart_empty": chart_metrics["empty"],
        "total_seconds": total_seconds,
        "pipeline_total_seconds": pipeline_total_seconds,
        "module_seconds": module_seconds,
    }


def run_sample_job(job):
    slot_polisher_backend = "auto"
    planned_slot_polisher_backend = "auto"
    pinco_command = None
    pinco_url = None
    pinco_timeout = 600
    if len(job) == 6:
        template_key, sample_index, data_path, sample_dir, output_png, chart_only = job
        slot_polish_after_chart = False
        slot_polish_dry_run = False
        planned_slot_polish = False
        planned_slot_dry_run = False
        planned_slot_allow_chart_overlap = True
        planned_slot_polisher_base_url = None
    elif len(job) == 8:
        (
            template_key,
            sample_index,
            data_path,
            sample_dir,
            output_png,
            chart_only,
            slot_polish_after_chart,
            slot_polish_dry_run,
        ) = job
        planned_slot_polish = False
        planned_slot_dry_run = False
        planned_slot_allow_chart_overlap = True
        planned_slot_polisher_base_url = None
    elif len(job) == 11:
        (
            template_key,
            sample_index,
            data_path,
            sample_dir,
            output_png,
            chart_only,
            slot_polish_after_chart,
            slot_polish_dry_run,
            planned_slot_polish,
            planned_slot_dry_run,
            planned_slot_allow_chart_overlap,
        ) = job
        planned_slot_polisher_base_url = None
    elif len(job) == 12:
        (
            template_key,
            sample_index,
            data_path,
            sample_dir,
            output_png,
            chart_only,
            slot_polish_after_chart,
            slot_polish_dry_run,
            planned_slot_polish,
            planned_slot_dry_run,
            planned_slot_allow_chart_overlap,
            planned_slot_polisher_base_url,
        ) = job
    else:
        (
            template_key,
            sample_index,
            data_path,
            sample_dir,
            output_png,
            chart_only,
            slot_polish_after_chart,
            slot_polish_dry_run,
            planned_slot_polish,
            planned_slot_dry_run,
            planned_slot_allow_chart_overlap,
            planned_slot_polisher_base_url,
            slot_polisher_backend,
            planned_slot_polisher_backend,
            pinco_command,
            pinco_url,
            pinco_timeout,
        ) = job
    result = run_sample(
        template_key,
        data_path,
        sample_dir,
        output_png,
        chart_only,
        slot_polish_after_chart=slot_polish_after_chart,
        slot_polish_dry_run=slot_polish_dry_run,
        planned_slot_polish=planned_slot_polish,
        planned_slot_dry_run=planned_slot_dry_run,
        planned_slot_allow_chart_overlap=planned_slot_allow_chart_overlap,
        planned_slot_polisher_base_url=planned_slot_polisher_base_url,
        slot_polisher_backend=slot_polisher_backend,
        planned_slot_polisher_backend=planned_slot_polisher_backend,
        pinco_command=pinco_command,
        pinco_url=pinco_url,
        pinco_timeout=pinco_timeout,
    )
    return {
        "template_key": template_key,
        "sample_index": sample_index,
        "data_source": str(data_path),
        "sample_dir": str(sample_dir),
        **result,
    }


def main():
    args = parse_args()
    rng = random.Random(args.seed)
    configure_chart_type_filter(args.respect_allowed_chart_types)

    if args.resume_root:
        output_root = Path(args.resume_root)
        timestamp = output_root.name
    else:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_root = make_output_root(args.output, timestamp)
    if args.clean and output_root.exists() and not args.resume_root:
        shutil.rmtree(output_root)
    output_root.mkdir(parents=True, exist_ok=True)

    run_config = {
        "created_at": timestamp,
        "output_parent": args.output or "output/chart_template_samples",
        "output_root": str(output_root),
        "seed": args.seed,
        "samples_per_template": args.samples_per_template,
        "data_dirs": args.data or CFG_DATA_DIRS,
        "engines": args.engines,
        "scan_limit": args.scan_limit,
        "limit_templates": args.limit_templates,
        "output_png": args.output_png,
        "chart_only": args.chart_only,
        "slot_polish_after_chart": args.slot_polish_after_chart,
        "slot_polish_dry_run": args.slot_polish_dry_run,
        "planned_slot_polish": args.planned_slot_polish,
        "planned_slot_dry_run": args.planned_slot_dry_run,
        "planned_slot_allow_chart_overlap": not args.planned_slot_disallow_chart_overlap,
        "planned_slot_polisher_base_url": args.planned_slot_polisher_base_url,
        "slot_polisher_backend": args.slot_polisher_backend,
        "planned_slot_polisher_backend": args.planned_slot_polisher_backend,
        "pinco_command": args.pinco_command or "",
        "pinco_url": args.pinco_url or "",
        "pinco_timeout": args.pinco_timeout,
        "png_longest_side": args.png_longest_side,
        "workers": args.workers,
        "scan_workers": args.scan_workers,
        "compat_cache": args.compat_cache,
        "cache_candidates_per_template": args.cache_candidates_per_template,
        "rebuild_compat_cache": args.rebuild_compat_cache,
        "respect_allowed_chart_types": args.respect_allowed_chart_types,
        "resume_root": args.resume_root,
        "pipeline_modules": PIPELINE_MODULES,
    }
    write_json(output_root / "run_config.json", run_config)
    os.environ["RENDER_LONGEST_SIDE"] = str(args.png_longest_side)

    input_files = collect_input_files(args.data or CFG_DATA_DIRS, rng, args.scan_limit)
    print(f"Collected {len(input_files)} input JSON files", flush=True)

    templates = scan_templates(force=True)
    analyze_templates(templates)
    template_keys = enumerate_template_keys(templates, args.engines)
    if args.limit_templates:
        template_keys = template_keys[:args.limit_templates]
    print(f"Collected {len(template_keys)} chart templates", flush=True)

    compat_cache_meta = cache_meta(
        args.data or CFG_DATA_DIRS,
        args.engines,
        args.respect_allowed_chart_types,
        args.scan_limit,
        template_keys,
    )
    compat_cache_path = Path(args.compat_cache)
    candidates_by_template = None
    if not args.rebuild_compat_cache:
        candidates_by_template = load_compatibility_cache(
            compat_cache_path,
            compat_cache_meta,
        )
        if candidates_by_template is not None:
            print(f"Loaded compatibility cache: {compat_cache_path}", flush=True)
    if candidates_by_template is None:
        cache_target = max(args.samples_per_template, args.cache_candidates_per_template)
        candidates_by_template = build_compatibility_index(
            input_files,
            templates,
            template_keys,
            cache_target,
            args.scan_workers,
        )
        save_compatibility_cache(
            compat_cache_path,
            compat_cache_meta,
            candidates_by_template,
        )
        print(f"Saved compatibility cache: {compat_cache_path}", flush=True)
    write_candidate_report(
        output_root / "template_candidates.csv",
        template_keys,
        candidates_by_template,
    )

    manifest_path = output_root / "manifest.jsonl"
    skipped_path = output_root / "skipped_templates.jsonl"
    completed_samples = set()
    if args.resume_root and manifest_path.exists():
        with open(manifest_path, "r", encoding="utf-8") as f:
            for line in f:
                if not line.strip():
                    continue
                record = json.loads(line)
                if record.get("success"):
                    completed_samples.add((record.get("template_key"), record.get("sample_index")))
    if manifest_path.exists() and not args.resume_root:
        manifest_path.unlink()
    if skipped_path.exists() and not args.resume_root:
        skipped_path.unlink()

    total_requested = len(template_keys) * args.samples_per_template
    total_success = len(completed_samples)
    total_failed = 0
    total_skipped_templates = 0
    jobs = []

    for template_index, template_key in enumerate(template_keys, 1):
        candidates = candidates_by_template.get(template_key, [])
        if not candidates:
            total_skipped_templates += 1
            append_jsonl(
                skipped_path,
                {"template_key": template_key, "reason": "no compatible data"},
            )
            print(
                f"[{template_index}/{len(template_keys)}] SKIP {template_key}: no compatible data",
                flush=True,
            )
            continue

        selected_data = choose_samples(candidates, args.samples_per_template, rng)
        template_slug = safe_slug(template_key)
        print(
            f"[{template_index}/{len(template_keys)}] {template_key}: "
            f"{len(candidates)} candidates, generating {len(selected_data)}",
            flush=True,
        )

        if args.dry_run:
            for sample_index, data_path in enumerate(selected_data):
                append_jsonl(
                    manifest_path,
                    {
                        "template_key": template_key,
                        "sample_index": sample_index,
                        "data_source": str(data_path),
                        "dry_run": True,
                    },
                )
            continue

        for sample_index, data_path in enumerate(selected_data):
            if (template_key, sample_index) in completed_samples:
                continue
            sample_dir = output_root / template_slug / f"sample_{sample_index:02d}"
            jobs.append((
                template_key,
                sample_index,
                data_path,
                sample_dir,
                args.output_png,
                args.chart_only,
                args.slot_polish_after_chart,
                args.slot_polish_dry_run,
                args.planned_slot_polish,
                args.planned_slot_dry_run,
                not args.planned_slot_disallow_chart_overlap,
                args.planned_slot_polisher_base_url,
                args.slot_polisher_backend,
                args.planned_slot_polisher_backend,
                args.pinco_command,
                args.pinco_url,
                args.pinco_timeout,
            ))

    if jobs:
        print(f"Running {len(jobs)} samples with workers={args.workers}", flush=True)
    if jobs and args.workers > 1:
        with ThreadPoolExecutor(max_workers=args.workers) as executor:
            futures = [executor.submit(run_sample_job, job) for job in jobs]
            for done_index, future in enumerate(as_completed(futures), 1):
                record = future.result()
                append_jsonl(manifest_path, record)
                if record["success"]:
                    total_success += 1
                else:
                    total_failed += 1
                print(
                    f"  completed {done_index}/{len(jobs)}: "
                    f"{record['template_key']} sample_{record['sample_index']:02d} "
                    f"success={record['success']}",
                    flush=True,
                )
    else:
        for done_index, job in enumerate(jobs, 1):
            record = run_sample_job(job)
            append_jsonl(manifest_path, record)
            if record["success"]:
                total_success += 1
            else:
                total_failed += 1
            print(
                f"  completed {done_index}/{len(jobs)}: "
                f"{record['template_key']} sample_{record['sample_index']:02d} "
                f"success={record['success']}",
                flush=True,
            )

    summary = {
        "total_templates": len(template_keys),
        "skipped_templates": total_skipped_templates,
        "samples_requested": total_requested,
        "samples_success": total_success,
        "samples_failed": total_failed,
        "output_root": str(output_root),
    }
    write_json(output_root / "summary.json", summary)
    print(json.dumps(summary, indent=2, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()