import os
import sys
from huggingface_hub import HfApi

# Absolute paths — these point at the real training output, not the staging folder
SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
ROOT_DIR = os.path.dirname(SCRIPT_DIR)  # one level up: GaussianAvatars/
OUTPUT_DIR = os.path.join(ROOT_DIR, "output")
MEDIA_DIR = os.path.join(ROOT_DIR, "media")

def upload_heavy_models(token):
    api = HfApi(token=token)
    dataset_repo = "jhduofg/gaussianavatarsstudy"
    print(f"Pushing heavy 3D models to {dataset_repo}...")
    
    # Upload Subject 306 (stored in media/)
    media_306 = os.path.join(MEDIA_DIR, "306")
    if os.path.exists(os.path.join(media_306, "point_cloud.ply")):
        print("  Uploading Subject 306...")
        api.upload_folder(folder_path=media_306, path_in_repo="models/Subject 306", repo_id=dataset_repo, repo_type="dataset")
        
    # Upload all SUBJECT_XXX folders from the real training output directory
    if not os.path.exists(OUTPUT_DIR):
        print(f"ERROR: Cannot find training output at {OUTPUT_DIR}")
        sys.exit(1)
        
    for subj in sorted(os.listdir(OUTPUT_DIR)):
        path = os.path.join(OUTPUT_DIR, subj, "point_cloud", "iteration_600000")
        if os.path.exists(path):
            subj_name = subj.replace("SUBJECT_", "Subject ")
            print(f"  Uploading {subj_name}...")
            api.upload_folder(folder_path=path, path_in_repo=f"models/{subj_name}", repo_id=dataset_repo, repo_type="dataset")
        else:
            print(f"  Skipping {subj} — no iteration_600000 checkpoint found.")
            
    print("Done! All models stored safely in Dataset.")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python upload_models.py <YOUR_HF_TOKEN>")
        sys.exit(1)
    upload_heavy_models(sys.argv[1])