#!/usr/bin/env python
"""
Helper script to deploy to Hugging Face Spaces.
This script will help you set environment variables and deploy your app.
"""
import os
import sys
import subprocess
from getpass import getpass
from huggingface_hub import HfApi, SpaceHardware, SpaceStage

def setup_deployment():
    """Set up deployment environment variables and authenticate."""
    print("="*50)
    print("Hugging Face Spaces Deployment Setup")
    print("="*50)
    
    # Get user credentials
    username = input("Enter your Hugging Face username: ")
    token = getpass("Enter your Hugging Face token (from https://huggingface.co/settings/tokens): ")
    space_name = input("Enter your Space name (default: personal-rag-assistant): ") or "personal-rag-assistant"
    
    # Set environment variables
    os.environ["HF_USERNAME"] = username
    os.environ["HF_TOKEN"] = token
    os.environ["SPACE_NAME"] = space_name
    
    # Write credentials to .env file
    with open(".env", "w") as f:
        f.write(f"HF_API_KEY={token}\n")
        f.write(f"HF_USERNAME={username}\n")
        f.write(f"SPACE_NAME={space_name}\n")
        f.write("LLM_MODEL=google/flan-t5-large\n")
        f.write("EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2\n")
        f.write("VECTOR_DB_PATH=./data/vector_db\n")
        f.write("COLLECTION_NAME=personal_assistant\n")
        f.write("DEFAULT_TEMPERATURE=0.7\n")
        f.write("CHUNK_SIZE=1000\n")
        f.write("CHUNK_OVERLAP=200\n")
        f.write("MAX_TOKENS=512\n")
    
    # Set up git credential helper for Hugging Face
    try:
        # Configure git to use credential store
        subprocess.run(["git", "config", "--global", "credential.helper", "store"], check=True)
        
        # Create .git-credentials file with token
        home_dir = os.path.expanduser("~")
        credentials_path = os.path.join(home_dir, ".git-credentials")
        
        # Check if credentials already exist
        credentials_exist = os.path.exists(credentials_path)
        
        with open(credentials_path, "a" if credentials_exist else "w") as f:
            f.write(f"https://{username}:{token}@huggingface.co\n")
        
        # Make sure credentials file has correct permissions
        if sys.platform != "win32":  # Skip on Windows
            os.chmod(credentials_path, 0o600)
        
        print("Git credentials configured for Hugging Face.")
    except Exception as e:
        print(f"Warning: Could not set up git credentials: {e}")
        print("You may need to enter your credentials manually during push.")
    
    print(f"\nEnvironment variables set for {username}/{space_name}")
    return username, token, space_name

def create_space(username, token, space_name):
    """Create a Hugging Face Space directly using the HfApi."""
    print("\nCreating Hugging Face Space...")
    
    try:
        # Initialize the API
        api = HfApi(token=token)
        
        # Check if space exists
        try:
            spaces = api.list_spaces(author=username)
            exists = any(space.id == f"{username}/{space_name}" for space in spaces)
            if exists:
                print(f"Space {username}/{space_name} exists.")
            else:
                print(f"Space {username}/{space_name} does not exist. Creating...")
                # Create the space
                api.create_space(
                    name=space_name,
                    organization=None,  # Use None for personal account
                    private=False,
                    sdk="docker",
                    hardware=SpaceHardware.CPU_BASIC,
                    storage=1,
                    sleep_time=3600,  # 1 hour of inactivity before sleep
                    status=SpaceStage.RUNNING,
                )
                print(f"Space created successfully.")
        except Exception as e:
            print(f"Error checking/creating space: {e}")
            print("You may need to create the Space manually in the Hugging Face UI.")
            print(f"Visit: https://huggingface.co/spaces")
            return False
        
        print(f"Space URL: https://huggingface.co/spaces/{username}/{space_name}")
        return True
    except Exception as e:
        print(f"Error creating space: {e}")
        print("\nTrying to proceed anyway, as the space might already exist.")
        return True

def prepare_git_push(username, space_name):
    """Prepare git for pushing to Hugging Face Space."""
    print("\nPreparing to push code to Hugging Face Space...")
    
    try:
        # Initialize git if not already done
        if not os.path.exists(".git"):
            subprocess.run(["git", "init"], check=True)
            print("Git repository initialized.")
        
        # Configure git remote
        remote_url = f"https://huggingface.co/spaces/{username}/{space_name}"
        
        # Check all remotes
        result = subprocess.run(["git", "remote"], capture_output=True, text=True)
        remotes = result.stdout.strip().split('\n') if result.stdout else []
        
        print(f"Existing remotes: {remotes}")
        
        # Check if 'hf' remote exists
        if "hf" not in remotes:
            # Add the remote
            print("Adding 'hf' remote...")
            try:
                add_result = subprocess.run(["git", "remote", "add", "hf", remote_url], capture_output=True, text=True)
                if add_result.returncode != 0:
                    print(f"Error adding remote: {add_result.stderr}")
                    return False
                print("Successfully added 'hf' remote")
            except Exception as e:
                print(f"Error adding remote: {e}")
                return False
        else:
            # Update existing remote
            print("Updating 'hf' remote...")
            try:
                update_result = subprocess.run(["git", "remote", "set-url", "hf", remote_url], capture_output=True, text=True)
                if update_result.returncode != 0:
                    print(f"Error updating remote: {update_result.stderr}")
                    return False
                print("Successfully updated 'hf' remote")
            except Exception as e:
                print(f"Error updating remote: {e}")
                return False
        
        # Verify remote was added/updated
        verify_remote = subprocess.run(["git", "remote", "-v"], capture_output=True, text=True)
        print(f"Remote verification: {verify_remote.stdout}")
        
        # Make sure we have user details
        try:
            user_name = subprocess.run(["git", "config", "user.name"], capture_output=True, text=True).stdout.strip()
            user_email = subprocess.run(["git", "config", "user.email"], capture_output=True, text=True).stdout.strip()
            
            if not user_name or not user_email:
                # Set default values if not configured
                name = input("Enter your name for git config: ") or username
                email = input("Enter your email for git config: ") or f"{username}@example.com"
                
                subprocess.run(["git", "config", "--global", "user.name", name], check=True)
                subprocess.run(["git", "config", "--global", "user.email", email], check=True)
                print("Git user configuration updated.")
        except subprocess.CalledProcessError:
            print("Warning: Could not check git user configuration.")
        
        # Add and commit files
        subprocess.run(["git", "add", "."], check=True)
        try:
            subprocess.run(["git", "commit", "-m", "Initial commit for Hugging Face Space deployment"], check=True)
        except subprocess.CalledProcessError:
            # Check if there are changes to commit
            status = subprocess.run(["git", "status", "--porcelain"], capture_output=True, text=True, check=True).stdout.strip()
            if not status:
                print("No changes to commit.")
            else:
                print("Error making commit. Check git configuration.")
                return False
        
        print("\nGit repository prepared for pushing")
    except Exception as e:
        print(f"Error preparing git: {e}")
        return False
    
    return True

def push_to_space(username, token):
    """Push code to Hugging Face Space."""
    print("\nPushing code to Hugging Face Space...")
    print("This may take a few minutes...")
    
    try:
        # Set git credentials environment variables for this push
        env = os.environ.copy()
        env["GIT_USERNAME"] = username
        env["GIT_PASSWORD"] = token
        
        # Determine current branch
        current_branch = subprocess.run(
            ["git", "branch", "--show-current"], 
            capture_output=True, text=True
        ).stdout.strip()
            
        if not current_branch:
            current_branch = "master"  # Default to master if no branch is returned
        
        # Push code - force push to override any existing content
        print(f"Pushing from branch {current_branch} to main...")
        cmd = ["git", "push", "-f", "hf", f"{current_branch}:main"]
        
        print("\nRunning git push command...")
        print(f"Pushing to Space as user: {username}")
        
        # Try to push
        try:
            subprocess.run(cmd, check=True, env=env)
        except subprocess.CalledProcessError as e:
            print(f"Error during push: {e}")
            
            # Try direct URL push as alternative
            print("\nTrying alternative direct URL push...")
            direct_url = f"https://{username}:{token}@huggingface.co/spaces/{username}/{os.environ.get('SPACE_NAME')}"
            alt_cmd = ["git", "push", "-f", direct_url, f"{current_branch}:main"]
            
            try:
                subprocess.run(alt_cmd, check=True, env=env)
            except subprocess.CalledProcessError as e:
                print(f"Direct URL push also failed: {e}")
                raise
        
        print("\nCode pushed to Hugging Face Space successfully!")
    except subprocess.CalledProcessError as e:
        print(f"Error pushing code: {e}")
        print("\nTroubleshooting git push issues:")
        print("1. Ensure your Hugging Face token has write access")
        print("2. Try manually setting up git credentials:")
        print(f"   git config --global credential.helper store")
        print(f"   echo 'https://{username}:{token}@huggingface.co' > ~/.git-credentials")
        print("3. Try pushing directly with:")
        print(f"   git push -f https://{username}:{token}@huggingface.co/spaces/{username}/{os.environ.get('SPACE_NAME')} main")
        return False
    
    return True

def main():
    """Main function to run the deployment process."""
    username, token, space_name = setup_deployment()
    
    # Create the Space
    if not create_space(username, token, space_name):
        print("Failed to create Space. Attempting to continue anyway.")
    
    # Prepare git for pushing
    if not prepare_git_push(username, space_name):
        print("Failed to prepare git. Exiting.")
        return
    
    # Push code to Space
    if not push_to_space(username, token):
        print("Failed to push code. Exiting.")
        return
    
    print("\n" + "="*50)
    print(f"Deployment completed! Your app should be available at:")
    print(f"https://huggingface.co/spaces/{username}/{space_name}")
    print("="*50)
    print("\nNote: It may take a few minutes for the Space to build and deploy your app.")
    print("You can monitor the build progress on the Space page.")

if __name__ == "__main__":
    main()