#!/usr/bin/env python3
"""
Setup and run script for the llama.cpp Hugging Face Space
"""

import subprocess
import sys
import os

def install_dependencies():
    """Install required dependencies"""
    print("📦 Installing dependencies...")
    
    try:
        # Upgrade pip first
        subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", "pip"], check=True)
        
        # Install requirements
        subprocess.run([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"], check=True)
        
        print("✅ Dependencies installed successfully!")
        return True
        
    except subprocess.CalledProcessError as e:
        print(f"❌ Error installing dependencies: {e}")
        return False

def test_installation():
    """Test if llama.cpp is properly installed"""
    print("🧪 Testing llama.cpp installation...")
    
    try:
        # Test import
        subprocess.run([sys.executable, "-c", "from llama_cpp import Llama; print('✅ llama-cpp-python imported successfully')"], check=True)
        
        # Test other dependencies
        test_imports = [
            "import gradio; print('✅ Gradio imported')",
            "import huggingface_hub; print('✅ Hugging Face Hub imported')",
            "from config import get_recommended_model; print('✅ Config imported')"
        ]
        
        for test_import in test_imports:
            subprocess.run([sys.executable, "-c", test_import], check=True)
        
        print("✅ All tests passed!")
        return True
        
    except subprocess.CalledProcessError as e:
        print(f"❌ Installation test failed: {e}")
        return False

def run_app():
    """Run the Gradio app"""
    print("🚀 Starting the Gradio app...")
    print("📝 Note: The Osmosis model will be downloaded on first use")
    print("🌐 The app will be available at http://localhost:7860")
    print("⏹️  Press Ctrl+C to stop the app")
    
    try:
        subprocess.run([sys.executable, "app.py"], check=True)
    except KeyboardInterrupt:
        print("\n👋 App stopped by user")
    except subprocess.CalledProcessError as e:
        print(f"❌ Error running app: {e}")

def main():
    """Main setup function"""
    print("🔧 llama.cpp Hugging Face Space Setup")
    print("=" * 50)
    
    # Check Python version
    if sys.version_info < (3, 8):
        print("❌ Python 3.8 or higher is required")
        sys.exit(1)
    
    print(f"✅ Python version: {sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}")
    
    # Install dependencies
    if not install_dependencies():
        print("❌ Failed to install dependencies")
        sys.exit(1)
    
    # Test installation
    if not test_installation():
        print("❌ Installation test failed")
        sys.exit(1)
    
    print("\n🎉 Setup completed successfully!")
    print("\n📋 What's installed:")
    print("   - llama-cpp-python for efficient CPU inference")
    print("   - Gradio for the web interface")
    print("   - Hugging Face Hub for model downloading")
    print("   - Osmosis Structure 0.6B model (will download on first use)")
    
    # Ask if user wants to run the app
    print("\n❓ Would you like to run the app now? (y/n): ", end="")
    try:
        response = input().lower().strip()
        if response in ['y', 'yes']:
            run_app()
        else:
            print("👍 Setup complete! Run 'python app.py' when ready.")
    except (EOFError, KeyboardInterrupt):
        print("\n👍 Setup complete! Run 'python app.py' when ready.")

if __name__ == "__main__":
    main()