import os
from huggingface_hub import HfApi, create_repo
from pathlib import Path
import sys


def check_files_exist(model_path):
    """Check if all necessary files exist in the model directory."""
    required_files = [
        "config.json",
        "model.safetensors",
        "special_tokens_map.json",
        "tokenizer.json",
        "tokenizer_config.json",
        "training_args.bin",
        "vocab.txt",
        "README.md"
    ]

    missing_files = []
    for file in required_files:
        if not os.path.exists(os.path.join(model_path, file)):
            missing_files.append(file)

    return missing_files


def upload_model_to_hf(model_path, repo_name, organization=None):
    """
    Upload a model to Hugging Face Hub.

    Args:
        model_path (str): Path to the model directory
        repo_name (str): Name for the repository on Hugging Face
        organization (str, optional): Organization name if uploading to an organization
    """
    try:
        # Initialize the API
        api = HfApi()

        # Check if all required files exist
        missing_files = check_files_exist(model_path)
        if missing_files:
            print(f"Error: Missing required files: {', '.join(missing_files)}")
            return False

        # Create full repository name
        if organization:
            full_repo_name = f"{organization}/{repo_name}"
        else:
            full_repo_name = f"{api.whoami()['name']}/{repo_name}"

        print(f"Creating repository: {full_repo_name}")

        # Create the repository
        try:
            create_repo(
                repo_id=full_repo_name,
                private=False,
                exist_ok=True
            )
        except Exception as e:
            print(f"Error creating repository: {str(e)}")
            return False

        print("Repository created successfully!")

        # Upload the model files
        print(f"Uploading files from {model_path}")
        api.upload_folder(
            folder_path=model_path,
            repo_id=full_repo_name,
            repo_type="model"
        )

        print("Upload completed successfully!")
        print(f"Your model is now available at: https://huggingface.co/{full_repo_name}")
        return True

    except Exception as e:
        print(f"An error occurred: {str(e)}")
        return False


if __name__ == "__main__":
    # Configurazione
    MODEL_PATH = "/Users/erikbranmarino/BERT-PRCT-fine-tuning/ct-bert-finetuned-20250131_120923"  # Il path al tuo modello
    REPO_NAME = "CT-BERT-PRCT"  # Il nome che vuoi dare al repository

    # Verifica che il path esista
    if not os.path.exists(MODEL_PATH):
        print(f"Error: Model path {MODEL_PATH} does not exist!")
        sys.exit(1)

    # Esegui l'upload
    success = upload_model_to_hf(MODEL_PATH, REPO_NAME)

    if not success:
        sys.exit(1)