import os
import sys
import shutil
from huggingface_hub import snapshot_download, HfApi

# Configuration
REPO_ID = os.environ.get("DATASET_REPO_ID")
HF_TOKEN = os.environ.get("HF_TOKEN")

def download():
    if not REPO_ID:
        print("DATASET_REPO_ID not set, skipping download.")
        return
    
    print(f"Downloading data from {REPO_ID}...")
    try:
        # snapshot_download is more efficient for many files than the CLI
        snapshot_download(
            repo_id=REPO_ID,
            repo_type="dataset",
            local_dir="data_repo",
            token=HF_TOKEN,
            max_workers=8
        )
        print("Download successful.")
    except Exception as e:
        print(f"Download failed: {e}")

def upload():
    if not REPO_ID:
        print("DATASET_REPO_ID not set, skipping upload.")
        return
    if not HF_TOKEN:
        print("HF_TOKEN not set, skipping upload.")
        return

    print(f"Uploading data to {REPO_ID}...")
    try:
        api = HfApi(token=HF_TOKEN)
        api.upload_folder(
            folder_path="data_repo",
            repo_id=REPO_ID,
            repo_type="dataset",
            # This handles large folders by committing in chunks if necessary
        )
        print("Upload successful.")
    except Exception as e:
        print(f"Upload failed: {e}")

def init_local():
    """Ensure data_repo has the necessary structure if download failed or it's new."""
    os.makedirs("data_repo/output", exist_ok=True)
    os.makedirs("data_repo/processed", exist_ok=True)
    os.makedirs("data_repo/uploads", exist_ok=True)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python hf_sync.py [download|upload|init]")
        sys.exit(1)
    
    action = sys.argv[1]
    if action == "download":
        download()
    elif action == "upload":
        upload()
    elif action == "init":
        init_local()
    else:
        print(f"Unknown action: {action}")