<#
.SYNOPSIS
Pre-stage VirtIO drivers on a VMware Windows Server VM before migrating to Proxmox (KVM/QEMU).

.DESCRIPTION
- Optionally downloads the latest virtio-win.iso, or uses a provided local ISO path.
- Mounts the ISO.
- Injects key drivers into the online Windows driver store using pnputil:
    * Storage: vioscsi + viostor
    * Network: NetKVM
    * Memory Balloon: Balloon
- (Optional) Adds QEMU guest extras (commented).
- Verifies presence.
- Unmounts the ISO.

.PARAMETER IsoPath
Path to an existing virtio-win.iso on disk. If omitted, use -Download to fetch it.

.PARAMETER Download
Download the latest virtio-win.iso from Fedora mirrors.

.EXAMPLE
.\Prep-VirtIODrivers.ps1 -Download

.EXAMPLE
.\Prep-VirtIODrivers.ps1 -IsoPath "C:\Temp\virtio-win.iso"

.NOTES
Run from an elevated PowerShell (Run as Administrator).
Tested on Windows Server 2025/2022/2019. Uses the 'w11\amd64' folder for 2025.
#>

param(
    [string]$IsoPath,
    [switch]$Download
)

# --- Require admin ---
$principal = New-Object Security.Principal.WindowsPrincipal([Security.Principal.WindowsIdentity]::GetCurrent())
if (-not $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
    Write-Error "Please run this script as Administrator."
    exit 1
}

$ErrorActionPreference = "Stop"

# --- Determine working dir ---
$WorkDir = Join-Path $env:TEMP "virtio-prep"
New-Item -ItemType Directory -Force -Path $WorkDir | Out-Null

# --- Ensure one of IsoPath or Download is set ---
if (-not $IsoPath -and -not $Download) { $Download = $true }

# --- Download logic ---
if ($Download) {
    try {
        $IsoPath = Join-Path $WorkDir "virtio-win.iso"
        Write-Host "Downloading latest virtio-win.iso..." -ForegroundColor Cyan
        $primary = "https://fedorapeople.org/groups/virt/virtio-win/direct-downloads/latest-virtio/virtio-win.iso"
        Invoke-WebRequest -Uri $primary -OutFile $IsoPath -UseBasicParsing -TimeoutSec 300
    } catch {
        Write-Warning "Primary download failed, trying Fedora mirror..."
        $fallback = "https://dl.fedoraproject.org/pub/alt/virtio-win/latest-virtio/virtio-win.iso"
        Invoke-WebRequest -Uri $fallback -OutFile $IsoPath -UseBasicParsing -TimeoutSec 300
    }
    Write-Host "Downloaded to $IsoPath" -ForegroundColor Green
} elseif (-not (Test-Path $IsoPath)) {
    Write-Error "IsoPath not found: $IsoPath"
    exit 1
}

# --- Mount ISO ---
Write-Host "Mounting ISO..." -ForegroundColor Cyan
$img = Mount-DiskImage -ImagePath $IsoPath -PassThru
$vol = $img | Get-Volume
$cd  = ($vol | Select-Object -First 1).DriveLetter + ":"
Write-Host "Mounted at $cd" -ForegroundColor Green

# --- Choose driver root folders
# Windows Server 2025 uses 'w11\amd64'. Provide a small fallback search if missing.
$CandidateSubDirs = @(
    "w11\amd64",  # Server 2025
    "2k22\amd64", # Server 2022
    "w10\amd64"   # Broad fallback
)

function Resolve-DriverPath {
    param(
        [string]$Base,
        [string[]]$Candidates
    )
    foreach ($c in $Candidates) {
        $p = Join-Path $Base $c
        if (Test-Path $p) { return $p }
    }
    return $null
}

$Drivers = @()
$map = @{
    "vioscsi" = Join-Path $cd "vioscsi"
    "viostor" = Join-Path $cd "viostor"
    "NetKVM"  = Join-Path $cd "NetKVM"
    "Balloon" = Join-Path $cd "Balloon"
}

foreach ($k in $map.Keys) {
    $root = $map[$k]
    $path = Resolve-DriverPath -Base $root -Candidates $CandidateSubDirs
    if ($null -eq $path) {
        Write-Warning "Could not find a suitable subfolder for $k under $root"
    } else {
        $Drivers += @{ Name=$k; Path=$path }
    }
}

# --- Install drivers using pnputil (online install into current OS) ---
Write-Host "`nInstalling VirtIO drivers..." -ForegroundColor Cyan
foreach ($d in $Drivers) {
    $infGlob = Join-Path $d.Path "*.inf"
    if (Test-Path $d.Path) {
        Write-Host ("  -> {0} from {1}" -f $d.Name, $d.Path) -ForegroundColor DarkCyan
        pnputil /add-driver $infGlob /install | Out-Null
    } else {
        Write-Warning "Driver path missing: $($d.Path)"
    }
}

# --- Optional extras (uncomment to include) ---
# pnputil /add-driver "$cd\qxldod\w11\amd64\*.inf" /install | Out-Null   # QXL Display
# pnputil /add-driver "$cd\vioserial\w11\amd64\*.inf" /install | Out-Null # VirtIO Serial
# pnputil /add-driver "$cd\pvpanic\w11\amd64\*.inf" /install | Out-Null   # pvpanic

# --- Verify presence in driver store ---
Write-Host "`nVerifying driver presence..." -ForegroundColor Cyan
$enum = (pnputil /enum-drivers) -join "`n"
foreach ($n in @("vioscsi","viostor","NetKVM","Balloon")) {
    if ($enum -match [Regex]::Escape($n)) {
        Write-Host ("  [+] {0} present" -f $n) -ForegroundColor Green
    } else {
        Write-Warning ("  [!] {0} NOT found in driver store" -f $n)
    }
}

# --- Unmount ISO ---
Write-Host "`nUnmounting ISO..." -ForegroundColor Cyan
Dismount-DiskImage -ImagePath $IsoPath

Write-Host "`nAll done." -ForegroundColor Green
Write-Host @"
Next steps:
1) Shut down this VM cleanly and take your Veeam backup/export OR run your migration.
2) In Proxmox, attach the disks as VirtIO-SCSI and NIC as VirtIO (paravirtualized).
3) First boot should succeed; if needed, keep virtio-win.iso mounted to install guest tools.
"@