# upload_all.py  — preserves exact paths & case
import time
from pathlib import Path
from typing import Iterable, Set
from huggingface_hub import HfApi

REPO_ID = "yc4ny/SVAD-models"   # change if needed
BASE_DIR = Path(".").resolve()

EXCLUDE_DIRS: Set[str] = {".git", ".hg", ".svn", ".idea", "__pycache__"}
EXCLUDE_FILES: Set[str] = {".DS_Store", "Thumbs.db", "desktop.ini, b.py"}

# If you want to restrict to specific top-level dirs, set like:
# LIMIT_TO_TOPLEVEL = {"avatar", "fitting", "submodules", "weights", "checkpoints", "face_warp"}
LIMIT_TO_TOPLEVEL: Set[str] | None = None

def iter_local_files(base: Path) -> Iterable[Path]:
    for p in base.rglob("*"):
        if p.is_dir():
            if any(part in EXCLUDE_DIRS for part in p.parts):
                continue
            continue
        if p.name in EXCLUDE_FILES:
            continue
        if any(part in EXCLUDE_DIRS for part in p.parts):
            continue
        if LIMIT_TO_TOPLEVEL is not None:
            parts = p.relative_to(base).parts
            if parts:
                top = parts[0]
                if top not in LIMIT_TO_TOPLEVEL:
                    continue
        yield p

def main():
    api = HfApi()
    print(f"📦 Target repo: {REPO_ID}")
    print(f"📂 Local root : {BASE_DIR}")

    print("🔎 Fetching remote file list...")
    remote_files = set(api.list_repo_files(repo_id=REPO_ID, repo_type="model"))
    print(f"   Remote has {len(remote_files)} files.")

    plan: list[tuple[Path, str]] = []
    for local_path in iter_local_files(BASE_DIR):
        rel_posix = local_path.relative_to(BASE_DIR).as_posix()
        dest = rel_posix  # <-- preserve exact path & case
        if dest not in remote_files:
            plan.append((local_path, dest))

    if not plan:
        print("✅ Nothing to upload — already in sync.")
        return

    print(f"📝 Planned uploads: {len(plan)} files")
    for lp, dp in plan[:20]:
        print(f"  + {lp} -> {dp}")
    if len(plan) > 20:
        print("  ... (truncated)")

    print("\n🚀 Uploading…")
    uploaded = 0
    failed = 0
    for idx, (local_path, dest) in enumerate(plan, 1):
        tries = 0
        while True:
            tries += 1
            try:
                api.upload_file(
                    repo_id=REPO_ID,
                    repo_type="model",
                    path_or_fileobj=str(local_path),
                    path_in_repo=dest,   # exact path kept
                )
                uploaded += 1
                if uploaded % 25 == 0 or idx == len(plan):
                    print(f"   … {uploaded}/{len(plan)} uploaded")
                break
            except Exception as e:
                if tries < 5:
                    print(f"⚠️  Retry {tries}/5: {local_path} -> {dest} :: {e}")
                    time.sleep(2 * tries)
                else:
                    print(f"❌ Failed: {local_path} -> {dest} :: {e}")
                    failed += 1
                    break

    print(f"\n✅ Done. Uploaded {uploaded} files. {'❌ '+str(failed)+' failed.' if failed else ''}")
    print("Re-run any time; it only uploads missing files.")

if __name__ == "__main__":
    main()