# Safe Dependency Installation Script for Windows
# Handles problematic packages like flash-attn carefully

Write-Host "🚀 OmniAvatar Dependency Installation" -ForegroundColor Green
Write-Host "====================================" -ForegroundColor Green

# Function to run pip command safely
function Install-Package {
    param(
        [string[]]$Command,
        [string]$Description,
        [bool]$Optional = $false
    )
    
    Write-Host "🔄 $Description" -ForegroundColor Yellow
    try {
        $result = & $Command[0] $Command[1..$Command.Length]
        if ($LASTEXITCODE -eq 0) {
            Write-Host "✅ $Description - Success" -ForegroundColor Green
            return $true
        } else {
            throw "Command failed with exit code $LASTEXITCODE"
        }
    } catch {
        if ($Optional) {
            Write-Host "⚠️ $Description - Failed (optional): $($_.Exception.Message)" -ForegroundColor Yellow
            return $false
        } else {
            Write-Host "❌ $Description - Failed: $($_.Exception.Message)" -ForegroundColor Red
            throw
        }
    }
}

try {
    # Step 1: Upgrade pip and essential tools
    Install-Package -Command @("python", "-m", "pip", "install", "--upgrade", "pip", "setuptools", "wheel", "packaging") -Description "Upgrading pip and build tools"
    
    # Step 2: Install PyTorch with CUDA support (if available)
    Write-Host "📦 Installing PyTorch..." -ForegroundColor Cyan
    try {
        Install-Package -Command @("python", "-m", "pip", "install", "torch", "torchvision", "torchaudio", "--index-url", "https://download.pytorch.org/whl/cu124") -Description "Installing PyTorch with CUDA support"
    } catch {
        Write-Host "⚠️ CUDA PyTorch failed, installing CPU version" -ForegroundColor Yellow
        Install-Package -Command @("python", "-m", "pip", "install", "torch", "torchvision", "torchaudio") -Description "Installing PyTorch CPU version"
    }
    
    # Step 3: Install main requirements
    Install-Package -Command @("python", "-m", "pip", "install", "-r", "requirements.txt") -Description "Installing main requirements"
    
    # Step 4: Try optional performance packages
    Write-Host "🎯 Installing optional performance packages..." -ForegroundColor Cyan
    
    # Try xformers
    Install-Package -Command @("python", "-m", "pip", "install", "xformers") -Description "Installing xformers (memory efficient attention)" -Optional $true
    
    # Flash-attn is often problematic, so we'll skip it by default
    Write-Host "ℹ️ Skipping flash-attn installation (often problematic on Windows)" -ForegroundColor Blue
    Write-Host "💡 You can try installing it later with: pip install flash-attn --no-build-isolation" -ForegroundColor Blue
    
    # Step 5: Verify installation
    Write-Host "🔍 Verifying installation..." -ForegroundColor Cyan
    
    python -c @"
import sys
try:
    import torch
    import transformers
    import gradio
    import fastapi
    
    print(f'✅ PyTorch: {torch.__version__}')
    print(f'✅ Transformers: {transformers.__version__}')
    print(f'✅ Gradio: {gradio.__version__}')
    
    if torch.cuda.is_available():
        print(f'✅ CUDA: {torch.version.cuda}')
        print(f'✅ GPU Count: {torch.cuda.device_count()}')
    else:
        print('ℹ️ CUDA not available - will use CPU')
    
    # Check optional packages
    try:
        import xformers
        print(f'✅ xformers: {xformers.__version__}')
    except ImportError:
        print('ℹ️ xformers not available (optional)')
    
    try:
        import flash_attn
        print('✅ flash_attn: Available')
    except ImportError:
        print('ℹ️ flash_attn not available (optional)')
    
    print('🎉 Installation verification successful!')
    
except ImportError as e:
    print(f'❌ Installation verification failed: {e}')
    sys.exit(1)
"@
    
    if ($LASTEXITCODE -eq 0) {
        Write-Host ""
        Write-Host "🎉 Installation completed successfully!" -ForegroundColor Green
        Write-Host ""
        Write-Host "💡 Next steps:" -ForegroundColor Yellow
        Write-Host "1. Download models: .\setup_omniavatar.ps1" -ForegroundColor White
        Write-Host "2. Start the app: python app.py" -ForegroundColor White
        Write-Host ""
    } else {
        throw "Installation verification failed"
    }
    
} catch {
    Write-Host ""
    Write-Host "❌ Installation failed: $($_.Exception.Message)" -ForegroundColor Red
    Write-Host ""
    Write-Host "💡 Troubleshooting tips:" -ForegroundColor Yellow
    Write-Host "1. Make sure Python 3.8+ is installed" -ForegroundColor White
    Write-Host "2. Try running in a virtual environment" -ForegroundColor White
    Write-Host "3. Check your internet connection" -ForegroundColor White
    Write-Host "4. For GPU support, ensure CUDA is properly installed" -ForegroundColor White
    exit 1
}