#!/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
import time
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)
    
    # Check if running in an environment with saved credentials
    username = os.environ.get("HF_USERNAME")
    token = os.environ.get("HF_TOKEN")
    space_name = os.environ.get("SPACE_NAME")
    
    # If not, ask for credentials
    if not (username and token and space_name):
        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=distilgpt2\n")  # Use smaller model to avoid 403 errors
        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=512\n")  # Smaller chunk size
        f.write("CHUNK_OVERLAP=128\n")  # Smaller overlap
        f.write("MAX_TOKENS=256\n")  # Smaller token limit
    
    # 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.")
                
                # Check if we need to update space configuration
                try:
                    print("Updating Space configuration to use Docker...")
                    api.update_space(
                        repo_id=f"{username}/{space_name}",
                        private=False,
                        sdk="docker",
                        hardware=SpaceHardware.CPU_BASIC
                    )
                    print("Space configuration updated.")
                except Exception as e:
                    print(f"Note: Could not update space configuration: {e}")
            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", "Update 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
        
        # Make sure HUGGINGFACEHUB_API_TOKEN is set in the environment
        env["HUGGINGFACEHUB_API_TOKEN"] = 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
            if not os.path.exists(".git/refs/heads/master"):
                current_branch = "main"  # Try main as another default
        
        # 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 different push methods in sequence until one works
        methods = [
            # Method 1: Standard remote push
            lambda: subprocess.run(cmd, check=True, env=env),
            
            # Method 2: Direct URL push
            lambda: subprocess.run(
                ["git", "push", "-f", f"https://{username}:{token}@huggingface.co/spaces/{username}/{os.environ.get('SPACE_NAME')}", f"{current_branch}:main"],
                check=True, env=env
            ),
            
            # Method 3: Push with credentials explicitly set
            lambda: subprocess.run(
                ["git", "push", "-f", "hf", f"{current_branch}:main"],
                check=True, env={**env, "HUGGINGFACE_TOKEN": token, "HF_TOKEN": token}
            )
        ]
        
        success = False
        for i, method in enumerate(methods, 1):
            try:
                print(f"\nTrying push method {i}...")
                method()
                print(f"Push method {i} succeeded!")
                success = True
                break
            except subprocess.CalledProcessError as e:
                print(f"Push method {i} failed: {e}")
                if i < len(methods):
                    print("Trying next method...")
                    time.sleep(2)  # Give a small delay before trying the next method
        
        if success:
            print("\nCode pushed to Hugging Face Space successfully!")
        else:
            raise Exception("All push methods failed")
            
        # Wait a moment to ensure the Space starts building
        print("\nWaiting for Space to start building...")
        time.sleep(5)
        
        print(f"\nYour Space will be available at: https://huggingface.co/spaces/{username}/{os.environ.get('SPACE_NAME')}")
        print("It may take a few minutes for the Space to build and start.")
        return True
        
    except Exception 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 entry point for the deployment script."""
    print("Hugging Face Space Deployment Script")
    print("="*50)
    print("This script will help you deploy your app to Hugging Face Spaces.")
    
    try:
        # Set up deployment environment
        username, token, space_name = setup_deployment()
        
        # Create the Space
        if not create_space(username, token, space_name):
            print("Error creating Space. Please check your credentials and try again.")
            sys.exit(1)
        
        # Prepare git repository
        if not prepare_git_push(username, space_name):
            print("Error preparing git repository. Please check your git configuration and try again.")
            sys.exit(1)
        
        # Push to Space
        if not push_to_space(username, token):
            print("Error pushing to Space. Please check the logs and try again.")
            sys.exit(1)
        
        print("\nDeployment complete!")
        print(f"Your app is now available at: https://huggingface.co/spaces/{username}/{space_name}")
        print("\nNote: It may take a few minutes for the Space to build and start.")
        print("If your app is not showing up properly, check the Space logs in the Hugging Face UI.")
        print("Common issues:")
        print("1. Permission errors - check that cache directories have proper permissions")
        print("2. Model loading errors - try using a smaller model")
        print("3. Port configuration - ensure app is running on port 7860")
    except KeyboardInterrupt:
        print("\nDeployment interrupted by user.")
        sys.exit(1)
    except Exception as e:
        print(f"\nUnexpected error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()