#!/usr/bin/env python3
"""在存放数据集的服务器上运行：批量生成样例并同步到 Hugging Face。

做法与团队 nine-source Space 一致：图片放进一个独立的 HF Dataset 仓库，
Space 只登记固定版本的 URL，网页不依赖 SSH、本机服务或临时隧道。

    export HF_TOKEN=hf_xxx          # 或先 hf auth login
    python tools/server_sync.py manifest.jsonl \
        --media-repo DarthSoluke/EgoNetDatasetResearch-media \
        --space DarthSoluke/EgoNetDatasetResearch

manifest.jsonl 每行一个样例，两种写法任选：

  已经渲染好的叠加图
    {"dataset": "hograspnet", "rgb": "/SSD/.../0042.png", "gt": "/SSD/.../0042_overlay.png",
     "caption": "受试者 01，侧捏", "source": "官方 MANO"}

  只有原图和网格，由本脚本渲染（网格可以是 MANO，也可以是转换后的 SOMA）
    {"dataset": "hograspnet", "rgb": "/SSD/.../0042.png", "mesh": "/SSD/.../0042_soma.npz",
     "caption": "受试者 01，侧捏", "source": "官方 MANO → R1 → SOMA"}
    mesh.npz 的格式见 render_overlay.py（verts、faces、K，可选 joints）

可选字段：name（样例名，默认取 rgb 文件名）。

参数：
  --replace     先清掉这些数据集里以前由本脚本同步的样例，再写入新的
  --local-only  只写到本地目录 --out（默认 ./sync_out），不上传，用于先检查
"""
import argparse
import json
import os
import re
import shutil
import sys
import tempfile
from pathlib import Path

from PIL import Image

sys.path.insert(0, str(Path(__file__).resolve().parent))
from render_overlay import overlay  # noqa: E402

ORIGIN = "server_sync"


def to_jpg(img, dst, max_side):
    img = img.convert("RGB")
    s = max_side / max(img.size)
    if s < 1:
        img = img.resize((round(img.width * s), round(img.height * s)), Image.LANCZOS)
    dst.parent.mkdir(parents=True, exist_ok=True)
    img.save(dst, "JPEG", quality=86, optimize=True)


def build(manifest, stage, max_side):
    import numpy as np
    rows = [json.loads(l) for l in Path(manifest).read_text(encoding="utf-8").splitlines() if l.strip()]
    items = {}
    for i, r in enumerate(rows, 1):
        ds = r["dataset"]
        name = re.sub(r"[^\w\-]", "_", r.get("name") or Path(r["rgb"]).stem.removesuffix("_rgb"))
        rgb = Image.open(r["rgb"])
        if r.get("gt"):
            gt = Image.open(r["gt"])
        elif r.get("mesh"):
            m = np.load(r["mesh"])
            K = np.array(r["K"]) if r.get("K") else m["K"]
            gt = overlay(rgb, m["verts"], m["faces"], K, joints=m["joints"] if "joints" in m.files else None)
        else:
            sys.exit(f"第 {i} 行既没有 gt 也没有 mesh")
        rel_rgb, rel_gt = f"media/{ds}/{name}_rgb.jpg", f"media/{ds}/{name}_gt.jpg"
        to_jpg(rgb, stage / rel_rgb, max_side)
        to_jpg(gt, stage / rel_gt, max_side)
        size = Image.open(stage / rel_rgb).size
        items.setdefault(ds, []).append({"id":name,"rgb":rel_rgb,"gt":rel_gt,"modes":[{"label":r.get("source", ""),"src":rel_gt}],
                                         "width":size[0],"height":size[1],"caption":r.get("caption", ""),"source":r.get("source", ""),
                                         "rgb_label":"原始 RGB","strata":r.get("strata",{}),"source_page":r.get("source_page",""),
                                         "evidence_url":r.get("evidence_url",""),"license":r.get("license",""),"note":"","origin":ORIGIN})
        print(f"[{i}/{len(rows)}] {ds}/{name}")
    return items


def merge(prev, items, prefix, replace):
    for ds, new in items.items():
        old = prev.get(ds, [])
        if replace:
            old = [p for p in old if p.get("origin") != ORIGIN]
        new_paths = {prefix + n["rgb"] for n in new}
        old = [p for p in old if p.get("rgb") not in new_paths]
        prev[ds] = [dict(n, rgb=prefix+n["rgb"], gt=prefix+n["gt"],
                         modes=[{"label":m["label"],"src":prefix+m["src"]} for m in n["modes"]]) for n in new] + old
    return prev


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("manifest")
    ap.add_argument("--media-repo", default="DarthSoluke/EgoNetDatasetResearch-media")
    ap.add_argument("--space", default="DarthSoluke/EgoNetDatasetResearch")
    ap.add_argument("--max-side", type=int, default=960)
    ap.add_argument("--replace", action="store_true")
    ap.add_argument("--local-only", action="store_true")
    ap.add_argument("--out", default="sync_out")
    a = ap.parse_args()

    if a.local_only:
        out = Path(a.out)
        items = build(a.manifest, out, a.max_side)
        pf = out / "data/previews.json"
        prev = json.loads(pf.read_text(encoding="utf-8")) if pf.exists() else {}
        pf.parent.mkdir(parents=True, exist_ok=True)
        pf.write_text(json.dumps(merge(prev, items, "", a.replace), ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
        print(f"已写到 {out}/，检查无误后去掉 --local-only 重新运行")
        return

    from huggingface_hub import HfApi, hf_hub_download
    api = HfApi(token=os.environ.get("HF_TOKEN"))
    stage = Path(tempfile.mkdtemp(prefix="soma_sync_"))
    try:
        items = build(a.manifest, stage, a.max_side)
        api.create_repo(a.media_repo, repo_type="dataset", exist_ok=True)
        info = api.upload_folder(repo_id=a.media_repo, repo_type="dataset", folder_path=str(stage),
                                 commit_message=f"同步 {sum(map(len, items.values()))} 对样例：{', '.join(sorted(items))}")
        rev = info.oid
        prefix = f"https://huggingface.co/datasets/{a.media_repo}/resolve/{rev}/"

        cur = hf_hub_download(a.space, "data/previews.json", repo_type="space", token=api.token)
        prev = json.loads(Path(cur).read_text(encoding="utf-8"))
        merged = merge(prev, items, prefix, a.replace)
        tmp = stage / "previews.json"
        tmp.write_text(json.dumps(merged, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")
        api.upload_file(path_or_fileobj=str(tmp), path_in_repo="data/previews.json", repo_id=a.space, repo_type="space",
                        commit_message=f"登记样例（媒体版本 {rev[:8]}）")
        print(f"完成：媒体仓库版本 {rev}，已更新 {a.space} 的 data/previews.json")
    finally:
        shutil.rmtree(stage, ignore_errors=True)


if __name__ == "__main__":
    main()