import os
import argparse
from huggingface_hub import HfApi, create_repo, HfFolder
from dotenv import load_dotenv

def upload_to_huggingface(repo_name = "Zeno0007/zenobot", local_dir='.', token=None):
    """
    Upload files to a Hugging Face repository.
    
    Args:
        repo_name (str): Name of the repository (e.g., 'username/model-name')
        local_dir (str): Local directory containing the files to upload
        token (str, optional): Hugging Face token. If None, will look for it in env vars.
    """
    # Get token from different sources with priority
    if token is None:
        # Try to load from .env file first
        load_dotenv()
        
        # If no variables found, try alternate .env filenames
        if not (os.environ.get("HUGGINGFACE_TOKEN") or os.environ.get("HF_TOKEN")):
            if os.path.exists(".env.sample"):
                load_dotenv(".env.sample")
        
        # Check various environment variables for the token
        token = (os.environ.get("HUGGINGFACE_TOKEN") or 
                os.environ.get("HF_TOKEN") or 
                os.environ.get("HF_API_TOKEN"))
        
        # If still no token, try the Hugging Face CLI cached token
        if token is None:
            token = HfFolder.get_token()
    
    if token is None:
        print("Error: No Hugging Face token found.")
        print("Please provide a token using one of these methods:")
        print("1. Pass it as an argument: --token YOUR_TOKEN")
        print("2. Set it as an environment variable: export HUGGINGFACE_TOKEN=YOUR_TOKEN")
        print("3. Create a .env file with HUGGINGFACE_TOKEN=YOUR_TOKEN")
        print("4. Login using the huggingface-cli: huggingface-cli login")
        return

    # Initialize the Hugging Face API
    api = HfApi(token=token)
    
    # Create or get the repository
    try:
        repo_url = create_repo(repo_name, exist_ok=True, token=token)
        print(f"Repository URL: {repo_url}")
    except Exception as e:
        print(f"Error creating repository: {e}")
        print("\nPlease make sure your token is valid and has write access to this repository.")
        print("You can create or find your token at: https://huggingface.co/settings/tokens")
        return
    
    # Define files to be included/excluded
    ignore_patterns = [
        ".git/", 
        "venv/", 
        "__pycache__/", 
        "*.pyc", 
        ".env",
        ".gitignore",
        "setup_venv.sh"
        
    ]
    
    # Upload the files
    print(f"Uploading files to {repo_name}...")
    try:
        api.upload_folder(
            folder_path=local_dir,
            repo_id=repo_name,
            ignore_patterns=ignore_patterns,
            token=token
        )
        print("Upload complete!")
    except Exception as e:
        print(f"Error during upload: {e}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Setup tokenizers for the model")
    parser.add_argument("repo_name", nargs="?", default="Zeno0007/zenobot", 
                        help="Name of the repository (e.g., 'username/model-name')")
    parser.add_argument("--dir", default=".", help="Local directory containing the files to upload")
    parser.add_argument("--token", help="Hugging Face token (optional if set as environment variable)")
    
    args = parser.parse_args()
    upload_to_huggingface(args.repo_name, args.dir, args.token)