# Setup Python 3.11 Environment for AudioForge ML Dependencies
# AudioCraft requires Python 3.11 (doesn't support 3.13)

$Colors = @{
    Red = "Red"
    Green = "Green"
    Yellow = "Yellow"
    Blue = "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 "║    Python 3.11 Setup for ML Dependencies                ║" -ForegroundColor Cyan
Write-Host "╚═══════════════════════════════════════════════════════════╝`n" -ForegroundColor Cyan

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

# Check for Python 3.11
Write-Info "Checking for Python 3.11..."

$python311 = $null
$python311Path = $null

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

# Try direct python3.11
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 {}
}

# Try python -V to check current version
if (-not $python311) {
    try {
        $currentVersion = python --version 2>&1
        if ($currentVersion -match "3\.11") {
            $python311 = "python"
            Write-Success "Current Python is 3.11: $currentVersion"
        } else {
            Write-Warning "Current Python is not 3.11: $currentVersion"
        }
    } catch {}
}

if (-not $python311) {
    Write-Error "Python 3.11 not found!"
    Write-Host ""
    Write-Host "📋 To install Python 3.11:" -ForegroundColor Cyan
    Write-Host "1. Download from: https://www.python.org/downloads/release/python-3110/" -ForegroundColor White
    Write-Host "2. Or use Windows Store: Search for 'Python 3.11'" -ForegroundColor White
    Write-Host "3. Or use py launcher: py -3.11 -m pip install --upgrade pip" -ForegroundColor White
    Write-Host ""
    Write-Host "After installing Python 3.11, run this script again." -ForegroundColor Yellow
    exit 1
}

# Create virtual environment
Write-Info "Creating Python 3.11 virtual environment..."
Set-Location $BackendPath

if (Test-Path $VenvPath) {
    Write-Warning "Virtual environment .venv311 already exists"
    $overwrite = Read-Host "Overwrite? (y/N)"
    if ($overwrite -ne "y") {
        Write-Info "Skipping virtual environment creation"
    } else {
        Remove-Item -Recurse -Force $VenvPath
        & $python311 -m venv .venv311
        Write-Success "Virtual environment created"
    }
} else {
    & $python311 -m venv .venv311
    Write-Success "Virtual environment created"
}

# Activate and install dependencies
Write-Info "Installing ML dependencies..."
$activateScript = Join-Path $VenvPath "Scripts\Activate.ps1"

if (-not (Test-Path $activateScript)) {
    Write-Error "Activation script not found: $activateScript"
    exit 1
}

# Install uv first
Write-Info "Installing uv package manager..."
& "$VenvPath\Scripts\python.exe" -m pip install --upgrade pip
& "$VenvPath\Scripts\python.exe" -m pip install uv

# Install ML dependencies
Write-Info "Installing ML dependencies (this may take 5-10 minutes)..."
Write-Warning "This will download ~2GB of files (PyTorch, AudioCraft models)"

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

if ($LASTEXITCODE -eq 0) {
    Write-Success "ML dependencies installed successfully!"
    Write-Host ""
    Write-Host "📝 To use this environment:" -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 ""
} else {
    Write-Error "Failed to install ML dependencies"
    Write-Host "Check the error messages above for details." -ForegroundColor Yellow
    exit 1
}