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

$InstallerVersion = "0.1.3"
$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"

# Always run relative to the folder containing this installer.
Set-Location -LiteralPath $PSScriptRoot

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,
        [string[]]$Arguments = @()
    )

    & $FilePath @Arguments
    $code = $LASTEXITCODE
    if ($null -ne $code -and $code -ne 0) {
        throw "Command failed with exit code ${code}: $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 -LiteralPath $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 -LiteralPath $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 -FilePath $winget -Arguments @(
        "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."
    }

    # Make the newly installed interpreter immediately usable in this PowerShell process.
    $pythonDir = Split-Path -Parent $pythonExe
    $scriptsDir = Join-Path $pythonDir "Scripts"
    $env:Path = "$pythonDir;$scriptsDir;$env:Path"

    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 -LiteralPath $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 -FilePath $winget -Arguments @(
        "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 Ensure-FreshVenv {
    param([Parameter(Mandatory = $true)][string]$Python)

    $venvPython = Join-Path $PSScriptRoot ".venv\Scripts\python.exe"

    if (Test-Path -LiteralPath $venvPython) {
        $valid = $false
        try {
            & $venvPython -c "import sys; print(sys.executable)" *> $null
            $valid = ($LASTEXITCODE -eq 0)
        } catch {
            $valid = $false
        }

        if (-not $valid) {
            Write-Host "Existing .venv is stale or belongs to another Python installation. Rebuilding it..." -ForegroundColor Yellow
            Remove-Item -LiteralPath (Join-Path $PSScriptRoot ".venv") -Recurse -Force
        }
    }

    if (-not (Test-Path -LiteralPath $venvPython)) {
        Write-Host "Creating Python virtual environment..." -ForegroundColor Cyan
        Invoke-NativeChecked -FilePath $Python -Arguments @("-m", "venv", ".venv")
    } else {
        Write-Host "Using existing .venv virtual environment." -ForegroundColor Cyan
    }

    return (Resolve-Path -LiteralPath $venvPython).Path
}

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

    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 -FilePath $Python -Arguments @(
            "-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 -FilePath $Python -Arguments @(
            "-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 -FilePath $Python -Arguments @("-m", "pip", "check")
    Invoke-NativeChecked -FilePath $Python -Arguments @(
        "-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 -FilePath $Python -Arguments @(
            "-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."
    }

    $venvPython = Ensure-FreshVenv -Python $pythonExe

    Write-Host "Updating pip and build tools..." -ForegroundColor Cyan
    Invoke-NativeChecked -FilePath $venvPython -Arguments @("-m", "pip", "install", "--upgrade", "pip", "setuptools", "wheel")

    Write-Host "Installing WordBeat Captioner and WhisperX dependencies..." -ForegroundColor Cyan
    # Pass -e as a literal argument array so PowerShell never interprets it as its own parameter.
    Invoke-NativeChecked -FilePath $venvPython -Arguments @("-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
}