#!/usr/bin/env python3
"""
Setup script for Production RAG Agent
Run this script to install dependencies and set up the environment
"""

import subprocess
import sys
import os
from pathlib import Path

def run_command(command):
    """Run a shell command and return the result"""
    try:
        result = subprocess.run(command, shell=True, check=True, capture_output=True, text=True)
        return True, result.stdout
    except subprocess.CalledProcessError as e:
        return False, e.stderr

def main():
    print("🚀 Setting up Production RAG Agent...")
    print("=" * 50)
    
    # Check Python version
    print(f"✓ Python version: {sys.version}")
    
    # Install dependencies
    print("\n📦 Installing dependencies...")
    success, output = run_command("pip install -r requirements.txt")
    
    if success:
        print("✓ Dependencies installed successfully")
    else:
        print(f"❌ Error installing dependencies: {output}")
        return
    
    # Check environment variables
    print("\n🔑 Checking environment variables...")
    
    required_vars = ["HF_TOKEN"]
    optional_vars = ["OPENWEATHER_API_KEY", "SERPER_API_KEY"]
    
    for var in required_vars:
        if os.getenv(var):
            print(f"✓ {var} is set")
        else:
            print(f"⚠️ {var} is not set (required for full functionality)")
    
    for var in optional_vars:
        if os.getenv(var):
            print(f"✓ {var} is set")
        else:
            print(f"ℹ️ {var} is not set (optional)")
    
    # Check directories
    print("\n📁 Checking directories...")
    base_dir = Path(__file__).parent
    
    directories = [
        base_dir / "knowledge_base",
        base_dir / "vector_store", 
        base_dir / "logs"
    ]
    
    for directory in directories:
        if directory.exists():
            print(f"✓ {directory.name} directory exists")
        else:
            directory.mkdir(parents=True, exist_ok=True)
            print(f"✓ Created {directory.name} directory")
    
    # Test imports
    print("\n🧪 Testing imports...")
    
    test_imports = [
        "gradio",
        "transformers", 
        "torch",
        "sentence_transformers",
        "langchain"
    ]
    
    for module in test_imports:
        try:
            __import__(module)
            print(f"✓ {module}")
        except ImportError:
            print(f"❌ {module} (install failed)")
    
    print("\n" + "=" * 50)
    print("🎉 Setup complete!")
    print("\nNext steps:")
    print("1. Set your HF_TOKEN environment variable")
    print("2. Add documents to the knowledge_base folder")
    print("3. Run: python app.py")
    print("\nFor Hugging Face Spaces:")
    print("1. Upload all files to your space")
    print("2. Set HF_TOKEN in space settings")
    print("3. Your space will automatically deploy!")

if __name__ == "__main__":
    main()