# Install ML Dependencies for AudioForge
# Requires Microsoft Visual C++ Build Tools

$Colors = @{
    Red = "Red"
    Green = "Green"
    Yellow = "Yellow"
    Blue = "Cyan"
    Cyan = "Cyan"
}

function Write-Info { Write-Host "[INFO] $args" -ForegroundColor $Colors.Blue }
function Write-Success { Write-Host "[SUCCESS] $args" -ForegroundColor $Colors.Green }
function Write-Warning { Write-Host "[WARNING] $args" -ForegroundColor $Colors.Yellow }
function Write-Error { Write-Host "[ERROR] $args" -ForegroundColor $Colors.Red }

Write-Host "`n╔═══════════════════════════════════════════════════════════╗" -ForegroundColor Cyan
Write-Host "║    ML Dependencies Installation Guide                    ║" -ForegroundColor Cyan
Write-Host "╚═══════════════════════════════════════════════════════════╝`n" -ForegroundColor Cyan

$ProjectRoot = Split-Path -Parent $PSScriptRoot
$BackendPath = Join-Path $ProjectRoot "backend"
$VenvPath = Join-Path $BackendPath ".venv311"

# Check for Visual C++ Build Tools
Write-Info "Checking for Visual C++ Build Tools..."

$vcToolsFound = $false
$vcToolsPaths = @(
    "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools",
    "C:\Program Files\Microsoft Visual Studio\2022\Community",
    "C:\Program Files\Microsoft Visual Studio\2022\Professional",
    "C:\Program Files\Microsoft Visual Studio\2022\Enterprise"
)

foreach ($path in $vcToolsPaths) {
    if (Test-Path $path) {
        Write-Success "Found Visual Studio at: $path"
        $vcToolsFound = $true
        break
    }
}

# Check for cl.exe (C++ compiler)
try {
    $null = Get-Command cl -ErrorAction Stop
    Write-Success "C++ compiler (cl.exe) found in PATH"
    $vcToolsFound = $true
} catch {
    # Not in PATH, but might still be installed
}

if (-not $vcToolsFound) {
    Write-Warning "Visual C++ Build Tools not detected!"
    Write-Host ""
    Write-Host "📥 Please install Visual C++ Build Tools:" -ForegroundColor Cyan
    Write-Host ""
    Write-Host "1. Download from:" -ForegroundColor White
    Write-Host "   https://visualstudio.microsoft.com/visual-cpp-build-tools/" -ForegroundColor Gray
    Write-Host ""
    Write-Host "2. Run the installer and select:" -ForegroundColor White
    Write-Host "   ☑ Desktop development with C++" -ForegroundColor Gray
    Write-Host "   ☑ C++ build tools" -ForegroundColor Gray
    Write-Host ""
    Write-Host "3. After installation, restart your terminal and run this script again." -ForegroundColor Yellow
    Write-Host ""
    Write-Host "💡 Quick download link:" -ForegroundColor Cyan
    Write-Host "   https://aka.ms/vs/17/release/vs_buildtools.exe" -ForegroundColor Gray
    Write-Host ""
    
    $openBrowser = Read-Host "Open download page in browser? (Y/n)"
    if ($openBrowser -ne "n") {
        Start-Process "https://visualstudio.microsoft.com/visual-cpp-build-tools/"
    }
    
    exit 1
}

Write-Success "Visual C++ Build Tools detected!"

# Check Python 3.11
Write-Info "Checking for Python 3.11..."
$python311 = $null

try {
    $version = py -3.11 --version 2>&1
    if ($LASTEXITCODE -eq 0) {
        $python311 = "py -3.11"
        Write-Success "Found Python 3.11: $version"
    }
} catch {}

if (-not $python311) {
    try {
        $version = python3.11 --version 2>&1
        if ($LASTEXITCODE -eq 0) {
            $python311 = "python3.11"
            Write-Success "Found Python 3.11: $version"
        }
    } catch {}
}

if (-not $python311) {
    Write-Error "Python 3.11 not found!"
    Write-Host "Please install Python 3.11 first." -ForegroundColor Yellow
    exit 1
}

# Create/verify virtual environment
Set-Location $BackendPath

if (-not (Test-Path $VenvPath)) {
    Write-Info "Creating Python 3.11 virtual environment..."
    & $python311 -m venv .venv311
    Write-Success "Virtual environment created"
} else {
    Write-Info "Virtual environment already exists"
}

# Install dependencies
Write-Host ""
Write-Info "Installing ML dependencies..."
Write-Warning "This will take 5-10 minutes and download ~2GB of files"
Write-Host ""

# Upgrade pip and install uv
Write-Info "Setting up package manager..."
& "$VenvPath\Scripts\python.exe" -m pip install --upgrade pip --quiet
& "$VenvPath\Scripts\python.exe" -m pip install uv --quiet
Write-Success "Package manager ready"

# Install ML dependencies
Write-Host ""
Write-Info "Installing PyTorch, AudioCraft, and dependencies..."
Write-Host "This may take several minutes...`n" -ForegroundColor Yellow

& "$VenvPath\Scripts\uv.exe" pip install --python "$VenvPath\Scripts\python.exe" -e ".[ml]"

if ($LASTEXITCODE -eq 0) {
    Write-Host ""
    Write-Success "✅ ML dependencies installed successfully!"
    Write-Host ""
    Write-Host "📝 To use the ML-enabled backend:" -ForegroundColor Cyan
    Write-Host "   cd backend" -ForegroundColor White
    Write-Host "   .venv311\Scripts\Activate.ps1" -ForegroundColor White
    Write-Host "   uvicorn app.main:app --reload" -ForegroundColor White
    Write-Host ""
    Write-Host "🎵 Music generation is now ready!`n" -ForegroundColor Green
    exit 0
} else {
    Write-Host ""
    Write-Error "Installation failed!"
    Write-Host "Check the error messages above for details." -ForegroundColor Yellow
    Write-Host ""
    Write-Host "Common issues:" -ForegroundColor Cyan
    Write-Host "  • Visual C++ Build Tools not properly installed" -ForegroundColor White
    Write-Host "  • Need to restart terminal after installing build tools" -ForegroundColor White
    Write-Host "  • Network issues downloading packages" -ForegroundColor White
    Write-Host ""
    exit 1
}