#!/usr/bin/env python3
"""给某个数据集添加一对「RGB / GT 叠加」样例。

用法：
    python tools/add_preview.py <数据集 id> --rgb a_rgb.png --gt a_gt.png \
        [--caption "说明"] [--source "GT 来源"] [--name 样例名] [--max-side 960]

图片会缩放到最长边 --max-side 像素，存为 JPEG 放进 media/<id>/，
登记到 data/preview_pool.json，并按现有分类重新抽样。同名样例会被覆盖。
数据集 id 见 data/datasets.json（页面地址 #ds-<id> 里的那段）。
"""
import argparse
import json
import sys
import subprocess
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent


def save(src: Path, dst: Path, max_side: int):
    try:
        from PIL import Image
    except ImportError:
        sys.exit("需要 Pillow：pip install pillow")
    im = Image.open(src).convert("RGB")
    w, h = im.size
    s = max_side / max(w, h)
    if s < 1:
        im = im.resize((round(w * s), round(h * s)), Image.LANCZOS)
    dst.parent.mkdir(parents=True, exist_ok=True)
    im.save(dst, "JPEG", quality=88, optimize=True)
    return im.size


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("dataset")
    ap.add_argument("--rgb", required=True, type=Path)
    ap.add_argument("--gt", required=True, type=Path)
    ap.add_argument("--caption", default="")
    ap.add_argument("--source", default="")
    ap.add_argument("--name", default="", help="样例名，默认取 RGB 文件名")
    ap.add_argument("--max-side", type=int, default=960)
    a = ap.parse_args()

    ids = {d["id"] for d in json.loads((ROOT / "data/datasets.json").read_text(encoding="utf-8"))["datasets"]}
    if a.dataset not in ids:
        sys.exit(f"未知的数据集 id：{a.dataset}\n可用的 id：{', '.join(sorted(ids))}")
    for p in (a.rgb, a.gt):
        if not p.is_file():
            sys.exit(f"找不到文件：{p}")

    name = a.name or a.rgb.stem.removesuffix("_rgb")
    rel_rgb = f"media/{a.dataset}/{name}_rgb.jpg"
    rel_gt = f"media/{a.dataset}/{name}_gt.jpg"
    width, height = save(a.rgb, ROOT / rel_rgb, a.max_side)
    save(a.gt, ROOT / rel_gt, a.max_side)

    pf = ROOT / "data/preview_pool.json"
    prev = json.loads(pf.read_text(encoding="utf-8")) if pf.exists() else {}
    lst = [p for p in prev.get(a.dataset, []) if p.get("rgb") != rel_rgb]
    lst.append({"id":name,"rgb":rel_rgb,"gt":rel_gt,"modes":[{"label":a.source,"src":rel_gt}],
                "width":width,"height":height,"caption":a.caption,"source":a.source,"rgb_label":"原始 RGB",
                "strata":{},"source_page":"","evidence_url":"","license":"","note":""})
    prev[a.dataset] = lst
    pf.write_text(json.dumps(prev, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
    subprocess.run([sys.executable, str(ROOT / "tools/import_team_previews.py"), "--per", "12"], check=True)
    print(f"已登记素材并更新抽样：{a.dataset} / {name}")


if __name__ == "__main__":
    main()