#!/usr/bin/env python3
"""
Automated Setup and Configuration for Unified AI Services
Helps set up the environment, validate configurations, and initialize services
"""

import os
import sys
import json
import asyncio
import subprocess
from pathlib import Path
from typing import Dict, List, Optional
import shutil

def print_header(title: str):
    """Print a formatted header"""
    print("\n" + "=" * 60)
    print(f"  {title}")
    print("=" * 60)

def print_step(step: str):
    """Print a step indicator"""
    print(f"\n🔧 {step}")

def check_python_version():
    """Check Python version compatibility"""
    print_step("Checking Python version...")
    
    version = sys.version_info
    if version.major < 3 or (version.major == 3 and version.minor < 8):
        print("❌ Python 3.8 or higher is required")
        print(f"   Current version: {version.major}.{version.minor}.{version.micro}")
        return False
    
    print(f"✅ Python {version.major}.{version.minor}.{version.micro} is compatible")
    return True

def create_directory_structure():
    """Create necessary directory structure"""
    print_step("Creating directory structure...")
    
    directories = [
        "services",
        "exports",
        "logs",
        "temp",
        "tests",
        "data"
    ]
    
    for directory in directories:
        path = Path(directory)
        if not path.exists():
            path.mkdir(parents=True, exist_ok=True)
            print(f"   ✅ Created directory: {directory}")
        else:
            print(f"   ✓  Directory exists: {directory}")

def check_service_files():
    """Check if service files exist"""
    print_step("Checking service files...")
    
    required_files = {
        "services/ner_service.py": "NER Service",
        "services/ocr_service.py": "OCR Service", 
        "services/rag_service.py": "RAG Service",
        "app.py": "Unified Application",
        "configs.py": "Configuration Management"
    }
    
    missing_files = []
    
    for file_path, description in required_files.items():
        if Path(file_path).exists():
            print(f"   ✅ {description}: {file_path}")
        else:
            print(f"   ❌ {description}: {file_path} (MISSING)")
            missing_files.append(file_path)
    
    if missing_files:
        print(f"\n⚠️  Missing files detected:")
        for file_path in missing_files:
            print(f"   - {file_path}")
        print("\nPlease ensure all service files are in the correct locations.")
        return False
    
    return True

def create_env_file():
    """Create or update .env file with user input"""
    print_step("Setting up environment configuration...")
    
    env_path = Path(".env")
    
    if env_path.exists():
        response = input("   .env file already exists. Overwrite? (y/N): ")
        if response.lower() != 'y':
            print("   Keeping existing .env file")
            return True
    
    print("\n📝 Please provide the following configuration values:")
    print("   (Press Enter to use default values shown in brackets)")
    
    # Collect configuration values
    config_values = {}
    
    # Server Configuration
    print("\n🌐 Server Configuration:")
    config_values['HOST'] = input("   Host [0.0.0.0]: ") or "0.0.0.0"
    config_values['DEBUG'] = input("   Debug mode (true/false) [True]: ") or "True"
    config_values['MAIN_PORT'] = input("   Main app port [8000]: ") or "8000"
    config_values['NER_PORT'] = input("   NER service port [8500]: ") or "8500"
    config_values['OCR_PORT'] = input("   OCR service port [8400]: ") or "8400"
    config_values['RAG_PORT'] = input("   RAG service port [8401]: ") or "8401"
    
    # PostgreSQL Configuration
    print("\n🗄️  PostgreSQL Configuration:")
    config_values['POSTGRES_HOST'] = input("   PostgreSQL host: ")
    config_values['POSTGRES_PORT'] = input("   PostgreSQL port [5432]: ") or "5432"
    config_values['POSTGRES_USER'] = input("   PostgreSQL user: ")
    config_values['POSTGRES_PASSWORD'] = input("   PostgreSQL password: ")
    config_values['POSTGRES_DATABASE'] = input("   PostgreSQL database [postgres]: ") or "postgres"
    
    # Azure OpenAI Configuration
    print("\n🤖 Azure OpenAI Configuration:")
    config_values['AZURE_OPENAI_ENDPOINT'] = input("   Azure OpenAI endpoint: ")
    config_values['AZURE_OPENAI_API_KEY'] = input("   Azure OpenAI API key: ")
    config_values['EMBEDDING_MODEL'] = input("   Embedding model [text-embedding-3-large]: ") or "text-embedding-3-large"
    
    # DeepSeek Configuration
    print("\n🧠 DeepSeek Configuration:")
    config_values['DEEPSEEK_ENDPOINT'] = input("   DeepSeek endpoint: ")
    config_values['DEEPSEEK_API_KEY'] = input("   DeepSeek API key: ")
    config_values['DEEPSEEK_MODEL'] = input("   DeepSeek model [DeepSeek-R1-0528]: ") or "DeepSeek-R1-0528"
    
    # Azure Document Intelligence Configuration
    print("\n📄 Azure Document Intelligence Configuration:")
    config_values['AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT'] = input("   Document Intelligence endpoint: ")
    config_values['AZURE_DOCUMENT_INTELLIGENCE_KEY'] = input("   Document Intelligence API key: ")
    
    # Azure Storage Configuration
    print("\n💾 Azure Storage Configuration:")
    config_values['AZURE_STORAGE_ACCOUNT_URL'] = input("   Storage account URL: ")
    config_values['AZURE_BLOB_SAS_TOKEN'] = input("   Blob SAS token: ")
    config_values['BLOB_CONTAINER'] = input("   Blob container [historylog]: ") or "historylog"
    
    # Write .env file
    try:
        with open(".env", "w") as f:
            f.write("# =================================================================\n")
            f.write("# Unified AI Services - Environment Configuration\n")
            f.write("# Generated by setup.py\n")
            f.write("# =================================================================\n\n")
            
            f.write("# Server Configuration\n")
            f.write(f"HOST={config_values['HOST']}\n")
            f.write(f"DEBUG={config_values['DEBUG']}\n")
            f.write(f"MAIN_PORT={config_values['MAIN_PORT']}\n")
            f.write(f"NER_PORT={config_values['NER_PORT']}\n")
            f.write(f"OCR_PORT={config_values['OCR_PORT']}\n")
            f.write(f"RAG_PORT={config_values['RAG_PORT']}\n\n")
            
            f.write("# PostgreSQL Configuration\n")
            f.write(f"POSTGRES_HOST={config_values['POSTGRES_HOST']}\n")
            f.write(f"POSTGRES_PORT={config_values['POSTGRES_PORT']}\n")
            f.write(f"POSTGRES_USER={config_values['POSTGRES_USER']}\n")
            f.write(f"POSTGRES_PASSWORD={config_values['POSTGRES_PASSWORD']}\n")
            f.write(f"POSTGRES_DATABASE={config_values['POSTGRES_DATABASE']}\n\n")
            
            f.write("# Azure OpenAI Configuration\n")
            f.write(f"AZURE_OPENAI_ENDPOINT={config_values['AZURE_OPENAI_ENDPOINT']}\n")
            f.write(f"AZURE_OPENAI_API_KEY={config_values['AZURE_OPENAI_API_KEY']}\n")
            f.write(f"EMBEDDING_MODEL={config_values['EMBEDDING_MODEL']}\n")
            f.write(f"AZURE_OPENAI_DEPLOYMENT_NAME={config_values['EMBEDDING_MODEL']}\n\n")
            
            f.write("# DeepSeek Configuration\n")
            f.write(f"DEEPSEEK_ENDPOINT={config_values['DEEPSEEK_ENDPOINT']}\n")
            f.write(f"DEEPSEEK_API_KEY={config_values['DEEPSEEK_API_KEY']}\n")
            f.write(f"DEEPSEEK_MODEL={config_values['DEEPSEEK_MODEL']}\n\n")
            
            f.write("# Azure Document Intelligence Configuration\n")
            f.write(f"AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT={config_values['AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT']}\n")
            f.write(f"AZURE_DOCUMENT_INTELLIGENCE_KEY={config_values['AZURE_DOCUMENT_INTELLIGENCE_KEY']}\n\n")
            
            f.write("# Azure Storage Configuration\n")
            f.write(f"AZURE_STORAGE_ACCOUNT_URL={config_values['AZURE_STORAGE_ACCOUNT_URL']}\n")
            f.write(f"AZURE_BLOB_SAS_TOKEN=\"{config_values['AZURE_BLOB_SAS_TOKEN']}\"\n")
            f.write(f"BLOB_CONTAINER={config_values['BLOB_CONTAINER']}\n\n")
            
            f.write("# Processing Configuration\n")
            f.write("MAX_FILE_SIZE=50\n")
            f.write("REQUEST_TIMEOUT=300\n")
            f.write("LOG_LEVEL=INFO\n")
            f.write("ALLOWED_ORIGINS=*\n")
            
            f.write("\n# RAG Specific Configuration\n")
            f.write("PG_HOST=${POSTGRES_HOST}\n")
            f.write("PG_PORT=${POSTGRES_PORT}\n")
            f.write("PG_DATABASE=vectorsearch\n")
            f.write("PG_USER=${POSTGRES_USER}\n")
            f.write("PG_PASSWORD=${POSTGRES_PASSWORD}\n")
            f.write("PG_SSL_MODE=require\n")
            f.write("CHUNK_SIZE=1000\n")
            f.write("CHUNK_OVERLAP=200\n")
            f.write("MIN_CHUNK_SIZE=50\n")
        
        print("   ✅ .env file created successfully")
        return True
        
    except Exception as e:
        print(f"   ❌ Failed to create .env file: {e}")
        return False

def install_dependencies():
    """Install Python dependencies"""
    print_step("Installing Python dependencies...")
    
    # Create requirements.txt if it doesn't exist
    requirements_path = Path("requirements.txt")
    if not requirements_path.exists():
        print("   Creating requirements.txt file...")
        
        requirements = [
            "fastapi>=0.104.1",
            "uvicorn[standard]>=0.24.0",
            "httpx>=0.25.0",
            "asyncpg>=0.29.0",
            "psutil>=5.9.0",
            "pydantic>=2.5.0",
            "python-dotenv>=1.0.0",
            "python-multipart>=0.0.6",
            "azure-ai-inference>=1.0.0",
            "azure-core>=1.29.0",
            "azure-storage-blob>=12.19.0",
            "azure-ai-documentintelligence>=1.0.0",
            "openai>=1.3.0",
            "beautifulsoup4>=4.12.0",
            "requests>=2.31.0",
            "numpy>=1.24.0",
            "Pillow>=10.0.0",
            "python-docx>=1.1.0",
            "lxml>=4.9.0",
            "aiofiles>=23.2.1"
        ]
        
        try:
            with open("requirements.txt", "w") as f:
                for req in requirements:
                    f.write(f"{req}\n")
            print("   ✅ requirements.txt created")
        except Exception as e:
            print(f"   ❌ Failed to create requirements.txt: {e}")
            return False
    
    # Install dependencies
    try:
        print("   Installing dependencies (this may take a few minutes)...")
        result = subprocess.run([
            sys.executable, "-m", "pip", "install", "-r", "requirements.txt"
        ], capture_output=True, text=True)
        
        if result.returncode == 0:
            print("   ✅ Dependencies installed successfully")
            return True
        else:
            print(f"   ❌ Failed to install dependencies:")
            print(f"   {result.stderr}")
            return False
            
    except Exception as e:
        print(f"   ❌ Error installing dependencies: {e}")
        return False

def validate_configuration():
    """Validate the configuration"""
    print_step("Validating configuration...")
    
    try:
        from configs import get_config, validate_environment
        
        config = get_config()
        config.print_configuration_summary()
        
        if validate_environment():
            print("\n   ✅ Configuration validation passed!")
            return True
        else:
            print("\n   ❌ Configuration validation failed!")
            print("   Please check your .env file and update missing values.")
            return False
            
    except ImportError as e:
        print(f"   ❌ Failed to import configuration module: {e}")
        return False
    except Exception as e:
        print(f"   ❌ Configuration validation error: {e}")
        return False

async def test_database_connection():
    """Test database connection"""
    print_step("Testing database connection...")
    
    try:
        from configs import get_config
        import asyncpg
        
        config = get_config()
        
        # Test connection
        conn = await asyncpg.connect(
            host=config.ner.POSTGRES_HOST,
            port=config.ner.POSTGRES_PORT,
            database=config.ner.POSTGRES_DATABASE,
            user=config.ner.POSTGRES_USER,
            password=config.ner.POSTGRES_PASSWORD,
            ssl='require',
            timeout=10
        )
        
        # Test basic query
        version = await conn.fetchval("SELECT version()")
        await conn.close()
        
        print("   ✅ Database connection successful")
        print(f"   Database version: {version[:50]}...")
        return True
        
    except Exception as e:
        print(f"   ❌ Database connection failed: {e}")
        print("\n   💡 Troubleshooting tips:")
        print("   1. Check your PostgreSQL server is running")
        print("   2. Verify host, port, username, and password")
        print("   3. Ensure your IP is allowlisted in firewall rules")
        print("   4. Check SSL configuration")
        return False

def create_startup_script():
    """Create startup script for easy service management"""
    print_step("Creating startup script...")
    
    # Create startup script for Windows
    if sys.platform == "win32":
        script_content = """@echo off
echo Starting Unified AI Services...
echo.

echo Starting in 3 seconds...
timeout /t 3 /nobreak >nul

echo Starting unified application...
python app.py

pause
"""
        with open("start_services.bat", "w") as f:
            f.write(script_content)
        print("   ✅ Created start_services.bat")
    
    # Create startup script for Unix/Linux/Mac
    else:
        script_content = """#!/bin/bash

echo "Starting Unified AI Services..."
echo

echo "Starting in 3 seconds..."
sleep 3

echo "Starting unified application..."
python app.py
"""
        with open("start_services.sh", "w") as f:
            f.write(script_content)
        
        # Make executable
        os.chmod("start_services.sh", 0o755)
        print("   ✅ Created start_services.sh")
    
    return True

def create_test_script():
    """Create test script for easy testing"""
    print_step("Creating test script...")
    
    # Create test script for Windows
    if sys.platform == "win32":
        script_content = """@echo off
echo Running Unified System Tests...
echo.

echo Make sure the unified application is running first!
echo Press any key to continue or Ctrl+C to cancel...
pause >nul

echo Running comprehensive tests...
python test_unified.py

pause
"""
        with open("run_tests.bat", "w") as f:
            f.write(script_content)
        print("   ✅ Created run_tests.bat")
    
    # Create test script for Unix/Linux/Mac
    else:
        script_content = """#!/bin/bash

echo "Running Unified System Tests..."
echo

echo "Make sure the unified application is running first!"
read -p "Press Enter to continue or Ctrl+C to cancel..."

echo "Running comprehensive tests..."
python test_unified.py
"""
        with open("run_tests.sh", "w") as f:
            f.write(script_content)
        
        # Make executable
        os.chmod("run_tests.sh", 0o755)
        print("   ✅ Created run_tests.sh")
    
    return True

def main():
    """Main setup function"""
    print_header("Unified AI Services - Automated Setup")
    
    print("This script will help you set up the Unified AI Services application.")
    print("It will:")
    print("  • Check your Python environment")
    print("  • Create necessary directories")
    print("  • Check for required service files")
    print("  • Set up configuration (.env file)")
    print("  • Install Python dependencies")
    print("  • Validate configuration")
    print("  • Test database connection")
    print("  • Create startup and test scripts")
    
    response = input("\nProceed with setup? (Y/n): ")
    if response.lower() == 'n':
        print("Setup cancelled.")
        return
    
    setup_steps = [
        ("Python Version Check", check_python_version),
        ("Directory Structure", create_directory_structure),
        ("Service Files Check", check_service_files),
        ("Environment Configuration", create_env_file),
        ("Dependencies Installation", install_dependencies),
        ("Configuration Validation", validate_configuration),
        ("Startup Scripts", create_startup_script),
        ("Test Scripts", create_test_script),
    ]
    
    failed_steps = []
    
    for step_name, step_func in setup_steps:
        try:
            if not step_func():
                failed_steps.append(step_name)
        except Exception as e:
            print(f"   ❌ {step_name} failed with exception: {e}")
            failed_steps.append(step_name)
    
    # Database connection test (optional)
    print_step("Testing database connection (optional)...")
    try:
        asyncio.run(test_database_connection())
    except Exception as e:
        print(f"   ⚠️  Database test skipped: {e}")
    
    # Final summary
    print_header("Setup Summary")
    
    if not failed_steps:
        print("🎉 Setup completed successfully!")
        print("\nNext steps:")
        print("1. Review the .env file and update any missing values")
        print("2. Start the unified application:")
        if sys.platform == "win32":
            print("   • Double-click start_services.bat")
            print("   • Or run: python app.py")
        else:
            print("   • Run: ./start_services.sh")
            print("   • Or run: python app.py")
        print("3. Test the system:")
        if sys.platform == "win32":
            print("   • Double-click run_tests.bat")
            print("   • Or run: python test_unified.py")
        else:
            print("   • Run: ./run_tests.sh")
            print("   • Or run: python test_unified.py")
        print("4. Access the API documentation at: http://localhost:8000/docs")
        
    else:
        print("⚠️  Setup completed with some issues:")
        for step in failed_steps:
            print(f"   ❌ {step}")
        
        print("\nPlease resolve the failed steps before proceeding.")
        print("You may need to:")
        print("• Check your internet connection for dependency installation")
        print("• Verify your Azure service credentials")
        print("• Ensure PostgreSQL is accessible")
        print("• Check file permissions")

if __name__ == "__main__":
    main()