#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Script to upload a simple BLIP model to test GPU functionality on Hugging Face Spaces
"""

import os
import sys
from huggingface_hub import HfApi, create_repo, upload_file

# Default repository name
DEFAULT_REPO = "mknolan/simple-gpu-test"
# Token should be provided at runtime
DEFAULT_TOKEN = ""

def main():
    """Main function to upload files to Hugging Face Spaces"""
    # Get Hugging Face token
    token = DEFAULT_TOKEN if DEFAULT_TOKEN else input("Enter your Hugging Face token (with WRITE access): ")
    
    # Get repository name
    repo_name = input("Enter repository name (default: {}): ".format(DEFAULT_REPO)) or DEFAULT_REPO
    
    print("Uploading to Space: {}".format(repo_name))
    
    # Initialize Hugging Face API
    api = HfApi(token=token)
    
    try:
        # Try to get the repository, create if it doesn't exist
        try:
            repo = api.repo_info(repo_id=repo_name, repo_type="space")
            print("Repo {} ready".format(repo_name))
        except Exception:
            print("Creating new Space: {}".format(repo_name))
            create_repo(
                repo_id=repo_name,
                token=token,
                repo_type="space",
                space_sdk="docker",
                private=False
            )
        
        print("Uploading files to Hugging Face Space...")
        
        # Upload Dockerfile
        api.upload_file(
            path_or_fileobj="Dockerfile.simple",
            path_in_repo="Dockerfile",
            repo_id=repo_name,
            repo_type="space",
            token=token,
            commit_message="Add Docker configuration for simple GPU test"
        )
        print("Uploaded Dockerfile")
        
        # Upload the Python script
        api.upload_file(
            path_or_fileobj="simple_gpu_app.py",
            path_in_repo="app.py",
            repo_id=repo_name,
            repo_type="space",
            token=token,
            commit_message="Add simple GPU test application"
        )
        print("Uploaded simple_gpu_app.py as app.py")
        
        # Create a README.md file
        readme_content = """# Simple GPU Test with BLIP

This Space provides a simple test of GPU functionality using the BLIP image captioning model.

## Features
- Uses a lightweight model (~1GB) that runs efficiently even on limited GPU resources
- Provides GPU detection and memory usage information
- Simple interface for testing if GPU is working

## How to Use
1. Upload an image
2. Click "Generate Caption"
3. View the results, including GPU usage information
"""
        
        # Write README to a temporary file
        with open("temp_readme.md", "w") as f:
            f.write(readme_content)
        
        # Upload README
        api.upload_file(
            path_or_fileobj="temp_readme.md",
            path_in_repo="README.md",
            repo_id=repo_name,
            repo_type="space",
            token=token,
            commit_message="Add README"
        )
        print("Uploaded README.md")
        
        # Clean up temp file
        os.remove("temp_readme.md")
        
        print("Upload completed!")
        print("Check your Space at: https://huggingface.co/spaces/{}".format(repo_name))
        
    except Exception as e:
        print("Error: {}".format(e))
        return 1
    
    return 0

if __name__ == "__main__":
    sys.exit(main())