#Requires -RunAsAdministrator
#
# Build-PxeUsb.ps1 - Build the PXE server installer USB on a Windows PC.
#
# The Windows twin of scripts/build-usb.sh: the same checks and the same
# payload, from a checkout of this repo plus its payload folders (for
# example \\<build PC>\pxe-images\pxe-usb, which scripts/stage-usb-share.sh
# fills). No Linux tools are needed.
#
# Layout (UEFI boot only, Secure Boot works): ONE FAT32 partition labelled
# CIDATA holding
#   - every file of the Ubuntu live-server ISO (boot/grub/grub.cfg gets
#     'autoinstall' on the first kernel line and a 5 s timeout), so the
#     firmware boots its shim and grub straight from the stick, and
#   - the CIDATA payload: user-data / meta-data (cloud-init NoCloud finds
#     them by the CIDATA label), packages/, playbook/, webapp/, pip-wheels/,
#     boot files, boot-tools/, enrollment/, bios/, site-content/ ...
# The installer and first-boot.sh find everything by that label, exactly as
# on build-usb.sh's two-partition stick. Windows formats FAT32 only up to
# 32 GB, so the partition is at most 32 GB; the rest of a bigger stick stays
# empty. FAT32 holds files under 4 GiB: boot files, enrollment and drivers
# over 3500 MiB are split into .part.NN files the playbook joins again, as
# build-usb.sh does; site-content and WinPE images are refused instead.
#
# Checks before the stick is touched (refuses to build what the server
# would reject):
#   - the ISO's volume id is the one offline-packages/ was resolved for
#   - the iPXE chain matches the ipxe_*_sha256 pins in the playbook
#   - every wheel in pip-wheels/ is the pinned version with a pinned hash
#   - site-content/ matches its SHA256SUMS list and has no file >= 4 GiB
#   - the site vars set pxe_nic_mac / corp_nic_mac to a MAC or leave them empty
#   - the disk is a USB disk, not the boot or system disk, and big enough
# After copying, every file is read back from the stick and compared with
# its source by SHA-256 (-QuickVerify: by size only).
#
# Usage: double-click Build-PxeUsb.cmd in the build folder (asks for admin
# rights, lists the USB disks and asks which one; uses site-vars.yml from
# the build folder if there is one, else the defaults). Or, elevated:
#   .\Build-PxeUsb.ps1                                    list USB disks
#   .\Build-PxeUsb.ps1 -DiskNumber 3 -NoSiteVars          build on disk 3
#   .\Build-PxeUsb.ps1 -DiskNumber 3 -SiteVars .\site-vars.yml
# Options: -IsoPath FILE, -PayloadRoot DIR, -WinPESource DIR,
#   -SkipEnrollment -SkipBios -SkipDrivers -SkipBootTools -SkipSiteContent,
#   -QuickVerify, -Yes (no typed confirmation).
#
# Exit codes: 0 built and verified, 1 refused or failed (nothing or a
# half-written stick: rebuild), 2 usage.

[CmdletBinding()]
param(
    [int]$DiskNumber = -1,
    [string]$IsoPath = '',
    [string]$PayloadRoot = '',
    [string]$SiteVars = '',
    [switch]$NoSiteVars,
    [string]$WinPESource = '',
    [switch]$SkipEnrollment,
    [switch]$SkipBios,
    [switch]$SkipDrivers,
    [switch]$SkipBootTools,
    [switch]$SkipSiteContent,
    [switch]$QuickVerify,
    [switch]$Yes,
    # Build-PxeUsb.cmd: ask for the disk, take site-vars.yml from the build
    # folder (else the defaults), wait for a key at the end.
    [switch]$Interactive
)

$ErrorActionPreference = 'Stop'
Set-StrictMode -Version 2

$RepoRoot = Split-Path -Parent $PSScriptRoot
if (-not $PayloadRoot) { $PayloadRoot = $RepoRoot }
$PayloadRoot = (Resolve-Path -LiteralPath $PayloadRoot).Path

$FatMax    = [int64]4GB - 1                # largest file FAT32 can hold
$ChunkSize = [int64]3500 * 1MB             # split size (build-usb.sh FAT32_MAX)
$PartMax   = [int64]32GB - 64MB            # Windows formats FAT32 up to 32 GB
$Utf8NoBom = New-Object System.Text.UTF8Encoding($false)

function Say([string]$m)  { Write-Host $m }
function Fail([string]$m) {
    Write-Host "ERROR: $m" -ForegroundColor Red
    if ($Interactive) { Write-Host ''; Read-Host 'Press Enter to close' | Out-Null }
    exit 1
}

function Show-UsbDisks {
    Say 'USB disks:'
    $d = @(Get-Disk | Where-Object { $_.BusType -eq 'USB' })
    if ($d.Count -eq 0) { Say '    (none found)'; return }
    foreach ($x in $d) {
        Say ('    {0,3}  {1,-34} {2,8:N1} GB  {3}' -f $x.Number, $x.FriendlyName, ($x.Size / 1GB), $x.PartitionStyle)
    }
    Say ''
    Say 'Then: .\Build-PxeUsb.ps1 -DiskNumber N -NoSiteVars   (or -SiteVars FILE)'
}

function Get-Sha256([string]$path) {
    (Get-FileHash -LiteralPath $path -Algorithm SHA256).Hash.ToLowerInvariant()
}

function Get-StreamSha256([System.IO.Stream]$s) {
    $sha = [System.Security.Cryptography.SHA256]::Create()
    try { -join ($sha.ComputeHash($s) | ForEach-Object { $_.ToString('x2') }) }
    finally { $sha.Dispose() }
}

# Every regular file under a folder (follows no links), relative paths with \.
function Get-Tree([string]$root, [string[]]$skipDirs = @(), [string[]]$skipExt = @()) {
    $base = (Resolve-Path -LiteralPath $root).Path.TrimEnd('\')
    $out = New-Object System.Collections.Generic.List[object]
    $stack = New-Object System.Collections.Stack
    $stack.Push($base)
    while ($stack.Count -gt 0) {
        $dir = $stack.Pop()
        foreach ($e in (Get-ChildItem -LiteralPath $dir -Force)) {
            if ($e.Attributes -band [IO.FileAttributes]::ReparsePoint) { continue }
            if ($e.PSIsContainer) {
                if ($skipDirs -notcontains $e.Name) { $stack.Push($e.FullName) }
                continue
            }
            if ($skipExt -contains $e.Extension.ToLowerInvariant()) { continue }
            $out.Add([pscustomobject]@{ Full = $e.FullName; Rel = $e.FullName.Substring($base.Length + 1); Size = $e.Length })
        }
    }
    , $out
}

# --- the copy plan -------------------------------------------------------------
# Each entry: Src, Dest (relative to the stick root), Size, Split (may be split
# into .part.NN), Data (bytes written instead of Src: the patched grub.cfg).
$Plan = New-Object System.Collections.Generic.List[object]
function Add-Plan([string]$src, [string]$dest, [int64]$size, [bool]$split = $false, [byte[]]$data = $null) {
    $Plan.Add([pscustomobject]@{ Src = $src; Dest = $dest.TrimStart('\'); Size = $size; Split = $split; Data = $data; Hashes = $null })
}
function Add-File([string]$src, [string]$destDir, [bool]$split = $false) {
    $f = Get-Item -LiteralPath $src
    $d = if ($destDir) { Join-Path $destDir $f.Name } else { $f.Name }
    Add-Plan $f.FullName $d $f.Length $split
}
function Add-Tree([string]$src, [string]$destDir, [bool]$split = $false, [string[]]$skipDirs = @(), [string[]]$skipExt = @(), [string[]]$skipRel = @()) {
    foreach ($t in (Get-Tree $src $skipDirs $skipExt)) {
        if ($skipRel -contains $t.Rel) { continue }
        Add-Plan $t.Full (Join-Path $destDir $t.Rel) $t.Size $split
    }
}

function Pause-IfInteractive {
    if ($Interactive) { Write-Host ''; Read-Host 'Press Enter to close' | Out-Null }
}
trap {
    Write-Host "ERROR: $_" -ForegroundColor Red
    Write-Host ("    at line {0}: {1}" -f $_.InvocationInfo.ScriptLineNumber, $_.InvocationInfo.Line.Trim())
    Pause-IfInteractive
    exit 1
}

# --- no disk given: list and stop (or ask) -------------------------------------
if ($DiskNumber -lt 0 -and $Interactive) {
    Show-UsbDisks
    $ans = Read-Host 'Disk number of the USB stick to build (Enter to cancel)'
    if ($ans -notmatch '^\d+$') { Say 'Cancelled.'; Pause-IfInteractive; exit 2 }
    $DiskNumber = [int]$ans
}
if ($Interactive -and -not $SiteVars -and -not $NoSiteVars) {
    $sv = Join-Path $RepoRoot 'site-vars.yml'
    if (Test-Path -LiteralPath $sv) { $SiteVars = $sv } else { $NoSiteVars = [switch]$true }
}
if ($DiskNumber -lt 0) { Show-UsbDisks; exit 2 }
if ($SiteVars -and $NoSiteVars) { Say 'Give -SiteVars FILE or -NoSiteVars, not both.'; exit 2 }
if (-not $SiteVars -and -not $NoSiteVars) {
    Say 'Give -SiteVars FILE (per-site values, see playbook\site-vars.example.yml)'
    Say 'or -NoSiteVars (the first boot then picks the imaging NIC itself and uses the defaults).'
    exit 2
}

$Autoinstall = Join-Path $RepoRoot 'autoinstall'
$PlaybookDir = Join-Path $RepoRoot 'playbook'
$WebappDir   = Join-Path $RepoRoot 'webapp'
$OfflineDir  = Join-Path $PayloadRoot 'offline-packages'
$WheelsDir   = Join-Path $PayloadRoot 'pip-wheels'
$BootFiles   = Join-Path $PayloadRoot 'boot-files'
$BootTools   = if ($SkipBootTools)   { '' } else { Join-Path $PayloadRoot 'boot-tools' }
$Enrollment  = if ($SkipEnrollment)  { '' } else { Join-Path $PayloadRoot 'enrollment' }
$BiosDir     = if ($SkipBios)        { '' } else { Join-Path $PayloadRoot 'bios-staging' }
$DriversDir  = if ($SkipDrivers)     { '' } else { Join-Path $PayloadRoot 'drivers-staging' }
$SiteContent = if ($SkipSiteContent) { '' } else { Join-Path $PayloadRoot 'site-content' }

Say '============================================'
Say 'PXE Server USB Builder (Windows)'
Say '============================================'
Say "Repo     : $RepoRoot"
Say "Payloads : $PayloadRoot"

# --- ISO ---------------------------------------------------------------------
if (-not $IsoPath) {
    $iso = Get-ChildItem -LiteralPath $PayloadRoot -Filter 'ubuntu-*-live-server-amd64.iso' -File -ErrorAction SilentlyContinue |
        Sort-Object Name | Select-Object -Last 1
    if (-not $iso) { Fail "no -IsoPath and no ubuntu-*-live-server-amd64.iso in $PayloadRoot" }
    $IsoPath = $iso.FullName
}
if (-not (Test-Path -LiteralPath $IsoPath -PathType Leaf)) { Fail "ISO not found: $IsoPath" }
$IsoPath = (Resolve-Path -LiteralPath $IsoPath).Path
Say "ISO      : $IsoPath"

# Volume id from the ISO9660 primary volume descriptor (sector 16, offset 40,
# 32 bytes): the same string xorriso -pvd_info prints on Linux.
$fs = [IO.File]::OpenRead($IsoPath)
try {
    $buf = New-Object byte[] 2048
    [void]$fs.Seek(16 * 2048, 'Begin')
    if ($fs.Read($buf, 0, 2048) -ne 2048 -or $buf[0] -ne 1 -or [Text.Encoding]::ASCII.GetString($buf, 1, 5) -ne 'CD001') {
        Fail "$IsoPath has no ISO9660 primary volume descriptor"
    }
    $IsoVolId = [Text.Encoding]::ASCII.GetString($buf, 40, 32).TrimEnd()
} finally { $fs.Dispose() }

# --- checks (as build-usb.sh) ------------------------------------------------
Say ''
Say '[1/6] Checking the payload...'
foreach ($f in @((Join-Path $Autoinstall 'user-data'), (Join-Path $Autoinstall 'meta-data'),
                 (Join-Path $PlaybookDir 'pxe_server_setup.yml'), (Join-Path $RepoRoot 'scripts\offline-apt-install.sh'),
                 (Join-Path $OfflineDir 'Packages'), (Join-Path $OfflineDir 'explicit.list'))) {
    if (-not (Test-Path -LiteralPath $f -PathType Leaf)) { Fail "$f not found (offline-packages: built by scripts/download-packages.sh)" }
}

$volFile = Join-Path $OfflineDir 'iso-volid'
if (Test-Path -LiteralPath $volFile) {
    $want = (Get-Content -LiteralPath $volFile -TotalCount 1).Trim()
    if ($want -ne $IsoVolId) { Fail "offline-packages\ was built for ISO '$want', this ISO is '$IsoVolId'. Use that ISO or rebuild the set." }
    Say "    offline-packages\ matches ISO $IsoVolId"
} else {
    Say "    WARNING: offline-packages\iso-volid missing; cannot check the set was resolved against $IsoVolId."
}

if ($SiteVars) {
    if (-not (Test-Path -LiteralPath $SiteVars -PathType Leaf)) { Fail "-SiteVars file not found: $SiteVars" }
    $SiteVars = (Resolve-Path -LiteralPath $SiteVars).Path
    $lines = Get-Content -LiteralPath $SiteVars
    foreach ($v in 'pxe_nic_mac', 'corp_nic_mac') {
        $line = @($lines | Where-Object { $_ -match "^${v}:" }) | Select-Object -Last 1
        if (-not $line) { continue }
        if ($line -match "^${v}:\s*([`"'][`"'])?\s*(#.*)?$") { continue }
        if ($line -notmatch "^${v}:\s*[`"']?([0-9A-Fa-f]{2}:){5}[0-9A-Fa-f]{2}[`"']?\s*(#.*)?$") {
            Fail "$SiteVars sets $v to something that is not a MAC (aa:bb:cc:dd:ee:ff) or empty"
        }
    }
}

# iPXE chain: the playbook's ipxe_*_sha256 pins.
$yml = [IO.File]::ReadAllText((Join-Path $PlaybookDir 'pxe_server_setup.yml'))
function Get-Pin([string]$name) {
    $m = [regex]::Match($yml, "(?m)^\s*${name}:\s*`"([0-9a-f]{64})`"")
    if (-not $m.Success) { Fail "no $name pin in playbook\pxe_server_setup.yml" }
    $m.Groups[1].Value
}
$tarball = Join-Path $BootFiles 'ipxeboot.tar.gz'
$shim    = Join-Path $BootFiles 'ipxe-shimx64.efi'
if (-not (Test-Path -LiteralPath $tarball) -or (Get-Sha256 $tarball) -ne (Get-Pin 'ipxe_tarball_sha256')) {
    Fail "$tarball is missing or does not match ipxe_tarball_sha256 (scripts/fetch-ipxe.sh fetches the pinned release)"
}
if (-not (Test-Path -LiteralPath $shim) -or (Get-Sha256 $shim) -ne (Get-Pin 'ipxe_shim_sha256')) {
    Fail "$shim is missing or does not match ipxe_shim_sha256"
}
$tmpTar = Join-Path ([IO.Path]::GetTempPath()) ('pxeusb-' + [guid]::NewGuid().ToString('N'))
New-Item -ItemType Directory -Path $tmpTar | Out-Null
try {
    & tar.exe -xzf $tarball -C $tmpTar 'ipxeboot/x86_64-sb/snponly.efi' 'ipxeboot/x86_64/ipxe.pxe' 2>$null
    foreach ($p in @(@('ipxeboot\x86_64-sb\snponly.efi', 'ipxe_snponly_sha256'), @('ipxeboot\x86_64\ipxe.pxe', 'ipxe_bios_sha256'))) {
        $f = Join-Path $tmpTar $p[0]
        if (-not (Test-Path -LiteralPath $f) -or (Get-Sha256 $f) -ne (Get-Pin $p[1])) { Fail "$($p[0]) in ipxeboot.tar.gz does not match $($p[1])" }
    }
} finally { Remove-Item -LiteralPath $tmpTar -Recurse -Force -ErrorAction SilentlyContinue }
Say '    iPXE chain matches the playbook pins'

# Wheels: each pinned package exactly once, at its pinned version and hash.
$reqs = [IO.File]::ReadAllText((Join-Path $WebappDir 'requirements.txt'))
function Norm([string]$n) { ($n -replace '[-_.]+', '_').ToLowerInvariant() }
$pins = @{}; $hashes = @{}
foreach ($m in [regex]::Matches($reqs, '(?m)^([A-Za-z0-9_.-]+)==([^\s\\]+)((?:\s*\\?\s*--hash=sha256:[0-9a-f]{64})+)')) {
    $n = Norm $m.Groups[1].Value
    $pins[$n] = $m.Groups[2].Value
    $hashes[$n] = @([regex]::Matches($m.Groups[3].Value, 'sha256:([0-9a-f]{64})') | ForEach-Object { $_.Groups[1].Value })
}
if ($pins.Count -eq 0) { Fail 'no hashed pins in webapp\requirements.txt' }
if (-not (Test-Path -LiteralPath $WheelsDir)) { Fail "no $WheelsDir (scripts/download-packages.sh builds it)" }
$bad = New-Object System.Collections.Generic.List[string]
$seen = @{}
foreach ($w in (Get-ChildItem -LiteralPath $WheelsDir -File)) {
    if ($w.Extension -ne '.whl') { $bad.Add("$($w.Name): not a wheel"); continue }
    $parts = $w.Name.Split('-')
    $n = Norm $parts[0]; $ver = $parts[1]
    if (-not $pins.ContainsKey($n)) { $bad.Add("$($w.Name): $n is not pinned") }
    elseif ($ver -ne $pins[$n]) { $bad.Add("$($w.Name): version $ver, pinned $($pins[$n])") }
    elseif ($hashes[$n] -notcontains (Get-Sha256 $w.FullName)) { $bad.Add("$($w.Name): sha256 is not a pinned hash of $n") }
    $seen[$n] = 1 + [int]($seen[$n])
}
foreach ($n in $pins.Keys) { if ([int]($seen[$n]) -ne 1) { $bad.Add("${n}: $([int]($seen[$n])) wheels (need exactly 1)") } }
if ($bad.Count -gt 0) { Fail ("pip-wheels\ does not match webapp\requirements.txt:`n       " + ($bad -join "`n       ")) }
Say "    $($pins.Count) wheels match their pinned hashes"

# site-content: listed in SHA256SUMS, nothing FAT32 cannot hold.
$scSums = @{}
if ($SiteContent -and -not (Test-Path -LiteralPath $SiteContent)) {
    # Optional: a site with no other PXE server has none. The server installs
    # and images without it; that content is added later if the site uses it.
    Say '    no site-content\ in the build folder (optional): building without it'
    $SiteContent = ''
}
if ($SiteContent) {
    $sumFile = Join-Path $SiteContent 'SHA256SUMS'
    if (-not (Test-Path -LiteralPath $sumFile)) { Fail 'site-content\ has no SHA256SUMS list; remove the folder or replace it with a complete copy.' }
    foreach ($l in [IO.File]::ReadAllLines($sumFile)) {
        if ($l -match '^([0-9a-f]{64})\s+\*?\./(.+)$') { $scSums[($Matches[2] -replace '/', '\')] = $Matches[1] }
    }
    $scFiles = @(Get-Tree $SiteContent | Where-Object { $_.Rel -ne 'SHA256SUMS' })
    if ($scFiles.Count -ne $scSums.Count) { Fail "site-content\ holds $($scFiles.Count) files but SHA256SUMS lists $($scSums.Count). Re-run the capture." }
    foreach ($t in $scFiles) {
        if (-not $scSums.ContainsKey($t.Rel)) { Fail "site-content\$($t.Rel) is not in SHA256SUMS" }
        if ($t.Size -gt $FatMax) { Fail "site-content\$($t.Rel) is 4 GiB or more (FAT32 cannot hold it)" }
    }
    Say "    site-content\ lists $($scSums.Count) files (hashed when they are read back)"
}
if ($WinPESource) {
    if (-not (Test-Path -LiteralPath $WinPESource)) { Fail "WinPE source not found: $WinPESource" }
    $big = @(Get-Tree $WinPESource | Where-Object { $_.Size -gt $FatMax })
    if ($big.Count) { Fail "WinPE content has files of 4 GiB or more; import those through the webapp instead: $($big[0].Rel)" }
}

# --- plan --------------------------------------------------------------------
Add-File (Join-Path $Autoinstall 'user-data') ''
Add-File (Join-Path $Autoinstall 'meta-data') ''
Add-File (Join-Path $OfflineDir 'Packages') 'packages'
Add-File (Join-Path $OfflineDir 'explicit.list') 'packages'
$debs = @(Get-ChildItem -LiteralPath $OfflineDir -Filter '*.deb' -File)
foreach ($d in $debs) { Add-Plan $d.FullName (Join-Path 'packages' $d.Name) $d.Length }
Add-File (Join-Path $RepoRoot 'scripts\offline-apt-install.sh') ''
# The playbook without a repo-local site-vars.yml: the stick carries exactly
# -SiteVars (or none).
Add-Tree $PlaybookDir 'playbook' $false @() @() @('site-vars.yml')
if ($SiteVars) { Add-Plan $SiteVars 'playbook\site-vars.yml' (Get-Item -LiteralPath $SiteVars).Length }
Add-File (Join-Path $RepoRoot 'scripts\lint-unattend.py') 'scripts'
Add-Tree $WebappDir 'webapp' $false @('__pycache__') @('.pyc', '.pyo')
Add-Tree $WheelsDir 'pip-wheels'
if (Test-Path -LiteralPath $BootFiles) {
    foreach ($f in (Get-ChildItem -LiteralPath $BootFiles -File)) { Add-Plan $f.FullName $f.Name $f.Length $true }
}
if ($BootTools -and (Test-Path -LiteralPath $BootTools)) { Add-Tree $BootTools 'boot-tools' }
if ($Enrollment -and (Test-Path -LiteralPath $Enrollment)) { Add-Tree $Enrollment 'enrollment' $true }
if ($SiteContent) { Add-Tree $SiteContent 'site-content' }
if ($BiosDir -and (Test-Path -LiteralPath $BiosDir)) { Add-Tree $BiosDir 'bios' }
if ($DriversDir -and (Test-Path -LiteralPath $DriversDir)) { Add-Tree $DriversDir 'drivers' $true }
if ($WinPESource) { Add-Tree $WinPESource 'images' }

foreach ($p in $Plan) {
    if ($p.Size -gt $FatMax -and -not $p.Split) { Fail "$($p.Src) is 4 GiB or more and FAT32 cannot hold it" }
}
$payloadBytes = ($Plan | Measure-Object -Property Size -Sum).Sum
$isoBytes = (Get-Item -LiteralPath $IsoPath).Length
$need = [int64]$payloadBytes + $isoBytes
$need = $need + [int64]($need / 10) + 256MB
Say ('    payload {0:N1} GB + ISO {1:N1} GB; needs about {2:N1} GB' -f ($payloadBytes / 1GB), ($isoBytes / 1GB), ($need / 1GB))
if ($need -gt $PartMax) { Fail ('that is more than the {0:N0} GB FAT32 partition Windows can format; skip a payload (-SkipDrivers, -SkipEnrollment ...)' -f ($PartMax / 1GB)) }

# --- the disk ------------------------------------------------------------------
$disk = Get-Disk -Number $DiskNumber -ErrorAction SilentlyContinue
if (-not $disk) { Show-UsbDisks; Fail "no disk $DiskNumber" }
if ($disk.IsBoot -or $disk.IsSystem) { Fail "disk $DiskNumber is the boot/system disk" }
if ($disk.BusType -ne 'USB') { Fail "disk $DiskNumber is on bus $($disk.BusType), not USB. Refusing." }
$partSize = [Math]::Min([int64]$disk.Size - 4MB, $PartMax)
if ($partSize -lt $need) { Fail ('disk {0} is {1:N1} GB; the build needs about {2:N1} GB' -f $DiskNumber, ($disk.Size / 1GB), ($need / 1GB)) }

Say ''
Say ('Disk {0}: {1}, {2:N1} GB, {3}' -f $disk.Number, $disk.FriendlyName, ($disk.Size / 1GB), $disk.PartitionStyle)
foreach ($v in @(Get-Partition -DiskNumber $DiskNumber -ErrorAction SilentlyContinue | Get-Volume -ErrorAction SilentlyContinue)) {
    Say ('    now holds: {0}: "{1}" {2} {3:N1} GB' -f $v.DriveLetter, $v.FileSystemLabel, $v.FileSystem, ($v.Size / 1GB))
}
if (-not $Yes) {
    $ans = Read-Host "EVERYTHING on disk $DiskNumber will be erased. Type the disk number to continue"
    if ($ans -ne [string]$DiskNumber) { Say 'Aborted.'; Pause-IfInteractive; exit 1 }
}

# --- partition and format ------------------------------------------------------
Say ''
Say "[2/6] Partitioning disk $DiskNumber (MBR, one FAT32 partition, label CIDATA)..."
# diskpart, not Clear-Disk / New-Partition: on removable USB sticks those
# often leave the old table half in place ("Not enough available capacity").
# Windows formats FAT32 up to 32 GB, so the partition is at most that.
$sizeMB = [int64][Math]::Floor($partSize / 1MB)
$dpScript = Join-Path ([IO.Path]::GetTempPath()) ('pxeusb-diskpart-' + [guid]::NewGuid().ToString('N') + '.txt')
@(
    "select disk $DiskNumber"
    'clean'
    'convert mbr'
    "create partition primary size=$sizeMB"
    'active'
    'format fs=fat32 quick label=CIDATA'
    'assign'
) | Set-Content -LiteralPath $dpScript -Encoding ASCII
$dpOut = & diskpart.exe /s $dpScript 2>&1
$dpCode = $LASTEXITCODE
Remove-Item -LiteralPath $dpScript -Force -ErrorAction SilentlyContinue
if ($dpCode -ne 0) {
    $dpOut | Select-Object -Last 15 | ForEach-Object { Say "    $_" }
    Fail "diskpart failed (exit $dpCode) on disk $DiskNumber"
}
Start-Sleep -Seconds 3
Update-Disk -Number $DiskNumber -ErrorAction SilentlyContinue
$part = @(Get-Partition -DiskNumber $DiskNumber | Where-Object { $_.Type -ne 'Reserved' }) | Select-Object -First 1
if (-not $part) { Fail "no partition on disk $DiskNumber after diskpart" }
if (-not $part.DriveLetter) {
    $part | Add-PartitionAccessPath -AssignDriveLetter
    Start-Sleep -Seconds 2
    $part = Get-Partition -DiskNumber $DiskNumber -PartitionNumber $part.PartitionNumber
}
if (-not $part.DriveLetter) { Fail 'the new partition got no drive letter' }
$vol = Get-Volume -DriveLetter $part.DriveLetter
if ($vol.FileSystem -ne 'FAT32' -or $vol.FileSystemLabel -ne 'CIDATA') {
    Fail "disk $DiskNumber is $($vol.FileSystem) '$($vol.FileSystemLabel)' after diskpart, not FAT32 'CIDATA'"
}
$Usb = "$($part.DriveLetter):\"
Say ('    {0} FAT32 "CIDATA" {1:N1} GB' -f $Usb, ($vol.Size / 1GB))

# --- the ISO -------------------------------------------------------------------
Say ''
Say '[3/6] Adding the Ubuntu ISO files...'
# Drive letter of the mounted ISO: the CD-ROM letter that was not there
# before. (Mount-DiskImage -PassThru | Get-Volume fails on some Windows
# PowerShell 5.1 builds with "Argument types do not match".)
function Get-CdLetters { @(Get-CimInstance Win32_LogicalDisk -Filter 'DriveType=5' | ForEach-Object { $_.DeviceID.Substring(0, 1) }) }
$before = Get-CdLetters
Mount-DiskImage -ImagePath $IsoPath | Out-Null
try {
    $isoLetter = $null
    for ($w = 0; $w -lt 30 -and -not $isoLetter; $w++) {
        $isoLetter = @(Get-CdLetters | Where-Object { $before -notcontains $_ }) | Select-Object -First 1
        if (-not $isoLetter) {
            try { $isoLetter = (Get-DiskImage -ImagePath $IsoPath | Get-Volume -ErrorAction Stop).DriveLetter } catch { }
        }
        if (-not $isoLetter) { Start-Sleep -Seconds 1 }
    }
    if (-not $isoLetter) { Fail "the ISO was mounted but got no drive letter: $IsoPath" }
    $isoRoot = "${isoLetter}:\"
    Say "    ISO mounted as $isoRoot"
    $isoTree = Get-Tree $isoRoot
    $grubRel = 'boot\grub\grub.cfg'
    $grubSrc = Join-Path $isoRoot $grubRel
    if (-not (Test-Path -LiteralPath $grubSrc)) { Fail "the ISO has no $grubRel" }
    if (-not (Test-Path -LiteralPath (Join-Path $isoRoot 'EFI\boot\bootx64.efi'))) { Fail 'the ISO has no EFI\boot\bootx64.efi' }

    # Same edit as build-usb.sh: 'autoinstall' on the default kernel line, 5 s timeout.
    $grub = [IO.File]::ReadAllText($grubSrc)
    $patched = $grub.Replace("linux`t/casper/vmlinuz  ---", "linux`t/casper/vmlinuz autoinstall ---")
    $patched = [regex]::Replace($patched, '(?m)^set timeout=30', 'set timeout=5')
    if ($patched -notmatch 'vmlinuz[^\n]*autoinstall') { Fail 'grub.cfg patch failed: no autoinstall on a kernel line (the ISO grub.cfg format changed)' }
    $grubBytes = $Utf8NoBom.GetBytes($patched)

    # The payload must not land on an ISO path.
    $isoNames = @{}
    foreach ($t in $isoTree) { $isoNames[$t.Rel.ToLowerInvariant()] = 1 }
    foreach ($p in $Plan) {
        if ($isoNames.ContainsKey($p.Dest.ToLowerInvariant())) { Fail "payload file $($p.Dest) would overwrite an ISO file" }
    }
    $isoPlan = New-Object System.Collections.Generic.List[object]
    foreach ($t in $isoTree) {
        if ($t.Size -gt $FatMax) { Fail "ISO file $($t.Rel) is 4 GiB or more" }
        if ($t.Rel -eq $grubRel) {
            $isoPlan.Add([pscustomobject]@{ Src = $t.Full; Dest = $t.Rel; Size = [int64]$grubBytes.Length; Split = $false; Data = $grubBytes; Hashes = $null })
        } else {
            $isoPlan.Add([pscustomobject]@{ Src = $t.Full; Dest = $t.Rel; Size = $t.Size; Split = $false; Data = $null; Hashes = $null })
        }
    }

    # --- copying -----------------------------------------------------------
    function Copy-PlanItem($p) {
        $dest = Join-Path $Usb $p.Dest
        $dir = Split-Path -Parent $dest
        if (-not (Test-Path -LiteralPath $dir)) { New-Item -ItemType Directory -Path $dir -Force | Out-Null }
        if ($p.Data) { [IO.File]::WriteAllBytes($dest, $p.Data); return }
        if (-not ($p.Split -and $p.Size -gt $ChunkSize)) { [IO.File]::Copy($p.Src, $dest, $true); return }
        # Split like `split -b 3500M -d -a 2 FILE FILE.part.`; the playbook joins them.
        $p.Hashes = New-Object System.Collections.Generic.List[string]
        $in = [IO.File]::OpenRead($p.Src)
        try {
            $buf = New-Object byte[] (8MB)
            $n = 0
            while ($in.Position -lt $in.Length) {
                $chunk = '{0}.part.{1:D2}' -f $dest, $n
                $out = [IO.File]::Create($chunk)
                $sha = [System.Security.Cryptography.SHA256]::Create()
                try {
                    $left = [Math]::Min($ChunkSize, $in.Length - $in.Position)
                    while ($left -gt 0) {
                        $r = $in.Read($buf, 0, [int][Math]::Min($buf.Length, $left))
                        if ($r -le 0) { throw "short read on $($p.Src)" }
                        $out.Write($buf, 0, $r)
                        [void]$sha.TransformBlock($buf, 0, $r, $null, 0)
                        $left -= $r
                    }
                    [void]$sha.TransformFinalBlock($buf, 0, 0)
                    $p.Hashes.Add((-join ($sha.Hash | ForEach-Object { $_.ToString('x2') })))
                } finally { $out.Dispose(); $sha.Dispose() }
                $n++
            }
        } finally { $in.Dispose() }
        Say "      split $($p.Dest) into $n parts"
    }

    # One list (not @() + @(): on Windows PowerShell 5.1 that adds the two
    # List objects themselves and fails "Argument types do not match").
    $all = New-Object System.Collections.Generic.List[object]
    $all.AddRange($isoPlan)
    $all.AddRange($Plan)
    $total = [int64](($all | Measure-Object -Property Size -Sum).Sum)
    $done = [int64]0
    $i = 0
    foreach ($p in $isoPlan) {
        Copy-PlanItem $p
        $done += $p.Size; $i++
        if ($i % 50 -eq 0) { Write-Progress -Activity 'Copying to the USB' -Status $p.Dest -PercentComplete ([int](100 * $done / $total)) }
    }
    Say "    $($isoPlan.Count) ISO files copied (grub.cfg: autoinstall, 5 s timeout)"

    Say ''
    Say '[4/6] Copying the payload (autoinstall, packages, playbook, webapp, boot files...)...'
    foreach ($p in $Plan) {
        Copy-PlanItem $p
        $done += $p.Size; $i++
        Write-Progress -Activity 'Copying to the USB' -Status $p.Dest -PercentComplete ([int](100 * $done / $total))
    }
    Write-Progress -Activity 'Copying to the USB' -Completed
    Say ("    {0} files, {1:N1} GB ({2} .deb, site vars: {3})" -f $Plan.Count, ($payloadBytes / 1GB), $debs.Count, $(if ($SiteVars) { $SiteVars } else { 'none' }))

    Say ''
    Say '[5/6] Flushing...'
    Write-VolumeCache -DriveLetter $part.DriveLetter

    # --- read back ---------------------------------------------------------
    Say ''
    Say ('[6/6] Reading every file back from the stick ({0})...' -f $(if ($QuickVerify) { 'sizes' } else { 'SHA-256' }))
    $errors = New-Object System.Collections.Generic.List[string]
    $done = [int64]0
    foreach ($p in $all) {
        $dest = Join-Path $Usb $p.Dest
        $done += $p.Size
        Write-Progress -Activity 'Verifying the USB' -Status $p.Dest -PercentComplete ([int](100 * $done / $total))
        if ($p.Hashes) {
            for ($k = 0; $k -lt $p.Hashes.Count; $k++) {
                $chunk = '{0}.part.{1:D2}' -f $dest, $k
                if (-not (Test-Path -LiteralPath $chunk)) { $errors.Add("$($p.Dest).part.$k missing"); continue }
                if (-not $QuickVerify -and (Get-Sha256 $chunk) -ne $p.Hashes[$k]) { $errors.Add("$($p.Dest).part.$k differs") }
            }
            continue
        }
        if (-not (Test-Path -LiteralPath $dest)) { $errors.Add("$($p.Dest) missing"); continue }
        if ((Get-Item -LiteralPath $dest).Length -ne $p.Size) { $errors.Add("$($p.Dest) has the wrong size"); continue }
        if ($QuickVerify) { continue }
        $got = Get-Sha256 $dest
        if ($p.Data) {
            $ms = New-Object IO.MemoryStream (, $p.Data)
            $want = Get-StreamSha256 $ms
        } elseif ($p.Dest.StartsWith('site-content\') -and $scSums.ContainsKey($p.Dest.Substring(13))) {
            $want = $scSums[$p.Dest.Substring(13)]   # the captured hash, not today's source
        } else {
            $want = Get-Sha256 $p.Src
        }
        if ($got -ne $want) { $errors.Add("$($p.Dest) differs from its source") }
    }
    Write-Progress -Activity 'Verifying the USB' -Completed
} finally {
    Dismount-DiskImage -ImagePath $IsoPath | Out-Null
}

if ($errors.Count -gt 0) {
    Say ''
    Say "The stick is NOT good: $($errors.Count) problem(s):"
    $errors | Select-Object -First 20 | ForEach-Object { Say "    $_" }
    Fail 'rebuild it (try another stick or port)'
}

Say ''
Say '============================================'
Say "USB build complete and verified: $Usb"
Say '============================================'
Say 'Next:'
Say '  1. Eject the stick (Safely Remove).'
Say '  2. On the server: plug in ONLY the imaging switch (GE unplugged), boot the'
Say '     stick in UEFI mode (F12 / boot menu). Ubuntu installs itself.'
Say '  3. First boot picks the imaging port, pins it as pxe0 = 172.16.9.1 and runs'
Say '     the playbook. Then plug GE into any other port.'
Say '  4. If the console shows a banner, read /var/log/first-boot.log; for the'
Say '     imaging port, log in and run: sudo pxe-pick-nic'
Pause-IfInteractive
exit 0