"""Publish the reproduction (code + docs + results + videos + checkpoint) to the Hub.

Deliberately EXCLUDES:
  * base model weights (models/, ~115GB, all redownloadable from their own repos)
  * our trained checkpoints (checkpoints*/) -- weights are not published
  * source/intermediate video data (data/videos, data/raw, data/cache)
  * third_party/ (the upstream Wan2.2 checkout)
  * the venv and every cache
  * .hf_token and any other secret

Run:  python scripts/push_to_hf.py --repo yqi19/VIPRA-reproduce
"""

from __future__ import annotations

import argparse
import os
import re
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]

ALLOW = [
    "viper/**",
    "scripts/**",
    "*.md",
    "env.sh",
    # NOTE: .gitignore is deliberately NOT published.  The Hub applies the repo's
    # own .gitignore server-side during commits, and ours carries "*.mp4"/"*.pt"
    # (there to keep 100+GB of local data out of git) -- publishing it makes the
    # Hub silently drop every result video from the commit.
    "data/*.jsonl",
    "results/**",
    "logs/*.log",
]

IGNORE = [
    "models/**", "checkpoints/**", "checkpoints_interactive/**",
    "third_party/**", ".venv/**", ".cache/**", ".tmp/**",
    ".uv_python/**", "data/videos/**", "data/raw/**", "data/cache/**",
    "**/__pycache__/**", "*.pyc", ".hf_token", "**/.hf_token",
    "results/smoke_videos/**",           # superseded by results/comparison
]

TOKEN_RE = re.compile(r"hf_[A-Za-z0-9]{30,}")

# Directories never worth walking (huge / irrelevant); keeps the scan fast.
PRUNE = {"models", "third_party", ".venv", ".cache", ".tmp", ".uv_python",
         "__pycache__", ".git", "checkpoints", "checkpoints_interactive"}


def _select(root: Path) -> list[Path]:
    """Files that upload_folder would send, using its fnmatch semantics."""
    from fnmatch import fnmatch

    out = []
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [d for d in dirnames if d not in PRUNE]
        for fn in filenames:
            p = Path(dirpath) / fn
            rel = str(p.relative_to(root))
            if any(fnmatch(rel, pat) for pat in IGNORE):
                continue
            if any(fnmatch(rel, pat) for pat in ALLOW):
                out.append(p)
    return sorted(out)


def secret_scan(root: Path) -> list[str]:
    """Refuse to publish if any tracked text file still contains a token."""
    hits = []
    exts = {".sh", ".py", ".md", ".json", ".jsonl", ".txt", ".yaml", ".yml"}
    skip = {"models", "third_party", ".venv", ".cache", ".tmp", ".uv_python",
            "__pycache__", ".uv-cache-tmp"}
    for p in root.rglob("*"):
        if not p.is_file() or p.suffix not in exts:
            continue
        if any(part in skip for part in p.parts):
            continue
        try:
            txt = p.read_text(errors="ignore")
        except Exception:
            continue
        for m in TOKEN_RE.findall(txt):
            hits.append(f"{p.relative_to(root)}: {m[:8]}...")
    return hits


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", required=True)
    ap.add_argument("--repo_type", default="model")
    ap.add_argument("--token", default=os.environ.get("HF_PUSH_TOKEN"))
    ap.add_argument("--dry_run", action="store_true")
    args = ap.parse_args()

    if not args.token:
        sys.exit("no token: pass --token or set HF_PUSH_TOKEN")

    leaks = secret_scan(ROOT)
    if leaks:
        print("REFUSING TO PUSH -- secrets found in files that would be published:")
        for l in leaks:
            print("   ", l)
        sys.exit(1)
    print("secret scan: clean")

    # Preview using the SAME matching huggingface_hub applies, so the dry run is
    # truthful (its filter is fnmatch over repo-relative paths, where '*' also
    # crosses '/', not pathlib.glob semantics).
    picked = _select(ROOT)
    total = sum(p.stat().st_size for p in picked)
    print(f"{len(picked)} files, {total/1e6:.1f} MB")
    by_top: dict[str, list] = {}
    for p in picked:
        by_top.setdefault(p.relative_to(ROOT).parts[0], []).append(p)
    for top, files in sorted(by_top.items()):
        sz = sum(f.stat().st_size for f in files)
        print(f"    {top:28s} {len(files):4d} files  {sz/1e6:8.1f} MB")

    if args.dry_run:
        print("\n(dry run -- nothing uploaded)")
        return

    from huggingface_hub import HfApi

    api = HfApi(token=args.token)
    api.create_repo(args.repo, repo_type=args.repo_type, exist_ok=True)
    print(f"\nuploading -> https://huggingface.co/{args.repo}")
    api.upload_folder(
        folder_path=str(ROOT),
        repo_id=args.repo,
        repo_type=args.repo_type,
        allow_patterns=ALLOW,
        ignore_patterns=IGNORE,
        commit_message="VIPER reproduction: code, mini VIPER-19K pipeline, "
                       "qualitative comparisons and eval",
    )

    # upload_folder honours the repo's .gitignore, and ours carries "*.mp4"/"*.pt"
    # to keep the 100+GB of local data out of git.  That silently drops every
    # generated video from the pass above, so push the comparison media from its
    # own folder (which contains no .gitignore) as a second commit.
    media = ROOT / "results" / "comparison"
    if media.is_dir():
        n = len(list(media.rglob("*.mp4")))
        print(f"\nuploading {n} videos from results/comparison ...")
        api.upload_folder(
            folder_path=str(media),
            path_in_repo="results/comparison",
            repo_id=args.repo,
            repo_type=args.repo_type,
            commit_message="qualitative comparison videos "
                           "(reference / target GT / Wan2.2 baseline / VIPER)",
        )
    print(f"done: https://huggingface.co/{args.repo}")


if __name__ == "__main__":
    main()