$ErrorActionPreference = "Stop"
Set-StrictMode -Version Latest

$InstallerVersion = "0.1.2"
$CudaIndex = "https://download.pytorch.org/whl/cu128"
$CpuIndex = "https://download.pytorch.org/whl/cpu"
$TorchVersion = "2.8.0"
$TorchVisionVersion = "0.23.0"
$TorchAudioVersion = "2.8.0"

Write-Host "WordBeat Captioner Windows installer v$InstallerVersion" -ForegroundColor Cyan

function Refresh-ProcessPath {
    $machinePath = [Environment]::GetEnvironmentVariable("Path", "Machine")
    $userPath = [Environment]::GetEnvironmentVariable("Path", "User")
    $env:Path = "$machinePath;$userPath"
}

function Invoke-NativeChecked {
    param(
        [Parameter(Mandatory = $true)][string]$FilePath,
        [Parameter(ValueFromRemainingArguments = $true)][string[]]$Arguments
    )

    & $FilePath @Arguments
    if ($LASTEXITCODE -ne 0) {
        throw "Command failed with exit code $LASTEXITCODE: $FilePath $($Arguments -join ' ')"
    }
}

function Get-WingetPath {
    $command = Get-Command winget -ErrorAction SilentlyContinue
    if (-not $command) {
        throw "WinGet is required. Install or update 'App Installer' from Microsoft Store, then run this installer again."
    }
    return $command.Source
}

function Find-Python311 {
    $candidates = @()

    $py = Get-Command py -ErrorAction SilentlyContinue
    if ($py) {
        try {
            $candidate = (& $py.Source -3.11 -c "import sys; print(sys.executable)").Trim()
            if ($LASTEXITCODE -eq 0 -and (Test-Path $candidate)) { $candidates += $candidate }
        } catch { }
    }

    $knownPaths = @(
        "$env:LOCALAPPDATA\Programs\Python\Python311\python.exe",
        "$env:ProgramFiles\Python311\python.exe",
        "${env:ProgramFiles(x86)}\Python311\python.exe"
    )
    foreach ($path in $knownPaths) {
        if ($path -and (Test-Path $path)) { $candidates += $path }
    }

    return $candidates | Select-Object -Unique | Select-Object -First 1
}

function Ensure-Python311 {
    $pythonExe = Find-Python311
    if ($pythonExe) {
        Write-Host "Python found: $pythonExe" -ForegroundColor Green
        return $pythonExe
    }

    Write-Host "Python 3.11 was not found. Installing it with WinGet..." -ForegroundColor Yellow
    $winget = Get-WingetPath
    Invoke-NativeChecked $winget install --exact --id Python.Python.3.11 --scope user --accept-package-agreements --accept-source-agreements --silent
    Refresh-ProcessPath

    $pythonExe = Find-Python311
    if (-not $pythonExe) {
        throw "Python installed, but python.exe was not located. Close PowerShell, reopen it, and rerun this installer."
    }

    Write-Host "Python installed: $pythonExe" -ForegroundColor Green
    return $pythonExe
}

function Find-FFmpeg {
    $command = Get-Command ffmpeg -ErrorAction SilentlyContinue
    if ($command) { return $command.Source }

    $packageRoot = "$env:LOCALAPPDATA\Microsoft\WinGet\Packages"
    if (Test-Path $packageRoot) {
        $found = Get-ChildItem -Path $packageRoot -Filter ffmpeg.exe -File -Recurse -ErrorAction SilentlyContinue |
            Where-Object { $_.FullName -like "*Gyan.FFmpeg*" } |
            Select-Object -First 1
        if ($found) {
            $env:Path = "$($found.DirectoryName);$env:Path"
            return $found.FullName
        }
    }
    return $null
}

function Ensure-FFmpeg {
    $ffmpegExe = Find-FFmpeg
    if ($ffmpegExe) {
        Write-Host "FFmpeg found: $ffmpegExe" -ForegroundColor Green
        return $ffmpegExe
    }

    Write-Host "FFmpeg was not found. Installing it with WinGet..." -ForegroundColor Yellow
    $winget = Get-WingetPath
    Invoke-NativeChecked $winget install --exact --id Gyan.FFmpeg --accept-package-agreements --accept-source-agreements --silent
    Refresh-ProcessPath

    $ffmpegExe = Find-FFmpeg
    if (-not $ffmpegExe) {
        throw "FFmpeg installed, but ffmpeg.exe was not located. Close PowerShell, reopen it, and rerun this installer."
    }

    Write-Host "FFmpeg installed: $ffmpegExe" -ForegroundColor Green
    return $ffmpegExe
}

function Assert-FreeSpace {
    param([double]$MinimumGB = 20)

    $projectRoot = (Get-Location).Path
    $root = [System.IO.Path]::GetPathRoot($projectRoot)
    $driveName = $root.Substring(0, 1)
    $drive = Get-PSDrive -Name $driveName
    $freeGB = [math]::Round($drive.Free / 1GB, 1)

    Write-Host "Free space on $($drive.Name): $freeGB GB"
    if ($freeGB -lt $MinimumGB) {
        throw "WordBeat requires at least $MinimumGB GB free. Only $freeGB GB is available on $($drive.Name):."
    }
}

function Test-NvidiaGpu {
    $nvidia = Get-Command nvidia-smi -ErrorAction SilentlyContinue
    if (-not $nvidia) { return $false }
    & $nvidia.Source --query-gpu=name --format=csv,noheader 2>$null | Out-Null
    return ($LASTEXITCODE -eq 0)
}

function Install-PyTorchBuild {
    param(
        [Parameter(Mandatory = $true)][string]$Python,
        [Parameter(Mandatory = $true)][bool]$UseCuda
    )

    # WhisperX declares PyTorch dependencies, so install WordBeat first. Then replace
    # whatever PyPI selected with the matching official CUDA or CPU wheels LAST.
    Write-Host "Removing any previously selected PyTorch build..." -ForegroundColor Cyan
    & $Python -m pip uninstall -y torch torchvision torchaudio 2>$null | Out-Host

    if ($UseCuda) {
        Write-Host "Installing official CUDA 12.8 PyTorch wheels..." -ForegroundColor Cyan
        Invoke-NativeChecked $Python -m pip install --no-cache-dir --force-reinstall `
            "torch==$TorchVersion" "torchvision==$TorchVisionVersion" "torchaudio==$TorchAudioVersion" `
            --index-url $CudaIndex
    } else {
        Write-Host "Installing official CPU PyTorch wheels..." -ForegroundColor Yellow
        Invoke-NativeChecked $Python -m pip install --no-cache-dir --force-reinstall `
            "torch==$TorchVersion" "torchvision==$TorchVisionVersion" "torchaudio==$TorchAudioVersion" `
            --index-url $CpuIndex
    }
}

function Verify-Installation {
    param(
        [Parameter(Mandatory = $true)][string]$Python,
        [Parameter(Mandatory = $true)][bool]$ExpectCuda
    )

    Write-Host "Verifying Python packages..." -ForegroundColor Cyan
    Invoke-NativeChecked $Python -m pip check
    Invoke-NativeChecked $Python -c "import whisperx, torch, torchaudio, torchvision; print('WhisperX import: OK'); print('Torch:', torch.__version__); print('Torch CUDA build:', torch.version.cuda); print('CUDA available:', torch.cuda.is_available())"

    if ($ExpectCuda) {
        Invoke-NativeChecked $Python -c "import torch, sys; ok=torch.cuda.is_available(); print('GPU:', torch.cuda.get_device_name(0) if ok else 'NONE'); sys.exit(0 if ok else 2)"
    }
}

try {
    Assert-FreeSpace -MinimumGB 20
    $pythonExe = Ensure-Python311
    $ffmpegExe = Ensure-FFmpeg

    $version = (& $pythonExe -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')").Trim()
    if ($version -ne "3.11") {
        throw "Python 3.11 is required. Found Python $version at $pythonExe."
    }

    if (-not (Test-Path ".\.venv\Scripts\python.exe")) {
        Write-Host "Creating Python virtual environment..." -ForegroundColor Cyan
        Invoke-NativeChecked $pythonExe -m venv .venv
    } else {
        Write-Host "Using existing .venv virtual environment." -ForegroundColor Cyan
    }

    $venvPython = (Resolve-Path ".\.venv\Scripts\python.exe").Path

    Write-Host "Updating pip and build tools..." -ForegroundColor Cyan
    Invoke-NativeChecked $venvPython -m pip install --upgrade pip setuptools wheel

    Write-Host "Installing WordBeat Captioner and WhisperX dependencies..." -ForegroundColor Cyan
    Invoke-NativeChecked $venvPython -m pip install -e .

    $hasNvidia = Test-NvidiaGpu
    if ($hasNvidia) {
        Write-Host "NVIDIA GPU detected." -ForegroundColor Green
    } else {
        Write-Host "No working NVIDIA GPU detected; CPU mode will be installed." -ForegroundColor Yellow
    }

    Install-PyTorchBuild -Python $venvPython -UseCuda $hasNvidia
    Verify-Installation -Python $venvPython -ExpectCuda $hasNvidia

    Write-Host ""
    Write-Host "INSTALLATION COMPLETE" -ForegroundColor Green
    Write-Host "Python:  $pythonExe"
    Write-Host "FFmpeg:  $ffmpegExe"
    Write-Host "Environment: $venvPython"
    Write-Host "Activate with: .\.venv\Scripts\Activate.ps1"
    Write-Host "Then run: wordbeat --help"
}
catch {
    Write-Host ""
    Write-Host "INSTALLATION FAILED" -ForegroundColor Red
    Write-Host $_.Exception.Message -ForegroundColor Red
    exit 1
}