import os
import logging
from huggingface_hub import HfApi
from dotenv import load_dotenv
import argparse

logging.basicConfig(level=logging.DEBUG)

# Load environment variables from .env file
load_dotenv()

# Access Hugging Face API token
hf_access_token = os.getenv("HUGGINGFACE_API_KEY")

# Set up Hugging Face API
api = HfApi()

# Repository ID to upload to
repo_id = "ErasureResearch/ErasingDiffusionModels"  # Target Hugging Face repository

# Ensure the repository exists
api.create_repo(repo_id, exist_ok=True, token=hf_access_token)

# List existing files in the repository to avoid overwrites
existing_files = api.list_repo_files(repo_id=repo_id, repo_type="model")

def upload_model_files(model_path, subfolder_path):
    for root, _, files in os.walk(subfolder_path):
        for file in files:
            file_path = os.path.join(root, file)
            path_in_repo = os.path.relpath(file_path, model_path)
            
            # Check if the file already exists in the repository
            if path_in_repo not in existing_files:
                logging.info(f"Uploading {file_path} to {repo_id}...")
                api.upload_file(
                    path_or_fileobj=file_path,
                    path_in_repo=path_in_repo,
                    repo_id=repo_id,
                    repo_type="model",
                    token=hf_access_token
                )
                logging.info(f"Successfully uploaded {file_path} to {repo_id}.")
            else:
                logging.info(f"File {path_in_repo} already exists in {repo_id}, skipping upload.")

def main():
    # Argument parser for model path
    parser = argparse.ArgumentParser(description="Upload model files to Hugging Face Hub")
    parser.add_argument("model_path", type=str, help="Path to the models directory")
    args = parser.parse_args()
    
    # Base path to your models directory
    model_path = args.model_path

    # Process each subfolder in the specified model path
    for subfolder in os.listdir(model_path):
        subfolder_path = os.path.join(model_path, subfolder)
        if os.path.isdir(subfolder_path):
            upload_model_files(model_path, subfolder_path)

if __name__ == "__main__":
    main()