import os
import argparse
from huggingface_hub import HfApi

def upload_to_hf():
    api = HfApi()
    
    # 1. Upload Franka Dataset
    dataset_path = "/storage/ice-shared/ae8803che/hxue/data/dataset/franka"
    dataset_repo = "t1an/world_model_dataset"
    
    if os.path.exists(dataset_path):
        print(f"Uploading Franka dataset to {dataset_repo}...")
        try:
            api.upload_folder(
                folder_path=dataset_path,
                repo_id=dataset_repo,
                repo_type="dataset",
                path_in_repo="franka",
                # ignore_patterns=["*.pt"] # We want metadata.pt though
            )
        except Exception as e:
            print(f"Error uploading dataset: {e}")
    else:
        print(f"Dataset path {dataset_path} not found.")

    # 2. Upload Code and Latest Checkpoints to Model Repo
    model_repo = "t1an/world_model"
    # Create the repo if it doesn't exist
    try:
        api.create_repo(repo_id=model_repo, exist_ok=True)
    except Exception as e:
        print(f"Note on repo creation: {e}")

    # Upload Code (wm folder)
    print(f"Uploading code to {model_repo}...")
    try:
        api.upload_folder(
            folder_path="/storage/ice-shared/ae8803che/hxue/data/world_model/wm",
            repo_id=model_repo,
            path_in_repo="wm"
        )
    except Exception as e:
        print(f"Error uploading code: {e}")

    # Upload only latest.pt from each checkpoint folder
    checkpoint_dir = "/storage/ice-shared/ae8803che/hxue/data/world_model/checkpoints"
    if os.path.exists(checkpoint_dir):
        for run_name in os.listdir(checkpoint_dir):
            run_path = os.path.join(checkpoint_dir, run_name)
            if os.path.isdir(run_path):
                latest_ckpt = os.path.join(run_path, "latest.pt")
                if os.path.exists(latest_ckpt):
                    print(f"Uploading {run_name}/latest.pt to {model_repo}...")
                    try:
                        api.upload_file(
                            path_or_fileobj=latest_ckpt,
                            path_in_repo=f"checkpoints/{run_name}/latest.pt",
                            repo_id=model_repo
                        )
                    except Exception as e:
                        print(f"Error uploading {run_name}/latest.pt: {e}")

if __name__ == "__main__":
    # Check if conversion is finished by checking for metadata.pt
    metadata_path = "/storage/ice-shared/ae8803che/hxue/data/dataset/franka/metadata.pt"
    if not os.path.exists(metadata_path):
        print("Error: Franka dataset conversion not finished (metadata.pt missing).")
        print("Please wait for the conversion to finish before running this script.")
    else:
        upload_to_hf()