#
# Upload-Image.ps1 - Copy the Media Creator Lite cache to the PXE server,
# then PROVE every file arrived intact.
#
# 1. Checks the MCL cache before sending anything:
#    - exactly one <OS>-Production.tag at the cache root (MCL writes it only
#      when a run finishes; none = an unfinished or interrupted cache)
#    - no Control file newer than that tag (MCL still writing / crashed)
#    - the OS in OperatingSystem.json is the one in user_selections.json
#    - every Control / Tools JSON is non-empty and parses
#    - every WIM, update and driver pack hashes to the SHA-256 MCL's own
#      Control JSONs publish (MCL resumes downloads onto old files and then
#      trusts any file of the right size, so a cache can hold junk)
# 2. Hashes every file it will send into upload-manifest.json.
# 3. Copies with robocopy as before (SMB share \\<server>\image-upload).
# 4. POSTs the manifest to the PXE webapp, which hashes the files on the
#    server's own disk (no Windows SMB cache in the way) and lists any file
#    that is missing or differs.
# 5. Re-sends those files over HTTP (small files; the server writes a temp
#    file, checks the hash, then renames it) or robocopy /J /IS (large
#    files), and verifies again, up to -MaxRounds times.
# The webapp's Image Import only accepts an upload that passed step 4.
#
# Usage:
#   .\Upload-Image.ps1                              (selected OS + packages, no drivers)
#   .\Upload-Image.ps1 -IncludeDrivers              (also upload selected hardware drivers)
#   .\Upload-Image.ps1 -CachePath "D:\MCL\Cache"    (custom cache location)
#   .\Upload-Image.ps1 -Server 172.16.9.1           (custom server IP)
#
# Exit codes: 0 verified, 1 failed (details printed), 2 MCL cache not usable.
#

param(
    [string]$CachePath = "C:\ProgramData\GEAerospace\MediaCreator\Cache",
    [string]$Server = "172.16.9.1",
    [int]$WebPort = 9009,
    [string]$User = "pxe-upload",
    [string]$Pass = "pxe",
    [int]$MaxRounds = 3,
    [switch]$IncludeDrivers
)

$ErrorActionPreference = "Stop"
$Share = "\\$Server\image-upload"
$ApiBase = "http://${Server}:$WebPort"
$ExcludedDeployDirs = @("Operating Systems", "Out-of-box Drivers", "Packages")

function Format-Size {
    param([long]$Bytes)
    if ($Bytes -ge 1GB) { return "{0:N1} GB" -f ($Bytes / 1GB) }
    if ($Bytes -ge 1MB) { return "{0:N1} MB" -f ($Bytes / 1MB) }
    return "{0:N0} KB" -f ($Bytes / 1KB)
}

function Resolve-DestDir {
    param([string]$Dir)
    $d = $Dir -replace '^\*destinationdir\*\\?', ''
    return (($d -replace '^[A-Za-z]:\\', '').Trim('\'))
}

function Write-Step {
    param([string]$Text)
    Write-Host ""
    Write-Host "========================================" -ForegroundColor Cyan
    Write-Host $Text -ForegroundColor Cyan
    Write-Host "========================================" -ForegroundColor Cyan
}

function Stop-Cache {
    # The cache itself is unusable: nothing was sent.
    param([string[]]$Lines)
    Write-Host ""
    foreach ($l in $Lines) { Write-Host "ERROR: $l" -ForegroundColor Red }
    Write-Host "Nothing was uploaded." -ForegroundColor Red
    exit 2
}

function Read-JsonFile {
    param([string]$Path)
    return (Get-Content -LiteralPath $Path -Raw -Encoding UTF8 | ConvertFrom-Json)
}

# HTTP to the webapp with HttpWebRequest and NO proxy: the work PC's system
# proxy (Zscaler) must not see traffic to the isolated imaging LAN.
# Returns @{ Status = <int>; Body = <string> }; never throws on HTTP errors.
function Invoke-Pxe {
    param([string]$Method, [string]$Path, [byte[]]$Body, [string]$InFile,
          [string]$ContentType = "application/json", [int]$TimeoutSec = 120)
    $req = [System.Net.HttpWebRequest]::Create("$ApiBase$Path")
    $req.Method = $Method
    $req.Proxy = $null
    $req.Timeout = $TimeoutSec * 1000
    $req.ReadWriteTimeout = $TimeoutSec * 1000
    $req.KeepAlive = $false
    if ($Body -or $InFile) {
        $req.ContentType = $ContentType
        if ($InFile) { $req.ContentLength = (Get-Item -LiteralPath $InFile).Length }
        else { $req.ContentLength = $Body.Length }
        try { $rs = $req.GetRequestStream() } catch { return @{ Status = 0; Body = $_.Exception.Message } }
        try {
            if ($InFile) {
                $fs = [System.IO.File]::OpenRead($InFile)
                try { $fs.CopyTo($rs) } finally { $fs.Dispose() }
            } else {
                $rs.Write($Body, 0, $Body.Length)
            }
        } finally { $rs.Dispose() }
    }
    try {
        $resp = $req.GetResponse()
    } catch {
        # A .NET method call wraps the WebException in a
        # MethodInvocationException: look at the inner one.
        $ex = $_.Exception
        while ($ex -and -not ($ex -is [System.Net.WebException])) { $ex = $ex.InnerException }
        if (-not $ex -or -not $ex.Response) { return @{ Status = 0; Body = $_.Exception.Message } }
        $resp = $ex.Response
    }
    try {
        $sr = New-Object System.IO.StreamReader($resp.GetResponseStream())
        $text = $sr.ReadToEnd()
        $sr.Dispose()
        return @{ Status = [int]$resp.StatusCode; Body = $text }
    } finally { $resp.Close() }
}

function Get-ErrorText {
    param($Resp)
    try { return ($Resp.Body | ConvertFrom-Json).error } catch { return $Resp.Body }
}

Write-Host ""
Write-Host "========================================" -ForegroundColor Cyan
Write-Host "  PXE Server Image Uploader" -ForegroundColor Cyan
Write-Host "========================================" -ForegroundColor Cyan
Write-Host ""

# --- Validate source paths ---
$CachePath = $CachePath.TrimEnd('\')
$DeployPath = Join-Path $CachePath "Deploy"
$ToolsPath  = Join-Path (Split-Path $CachePath -Parent) "Tools"
if (-not (Test-Path $ToolsPath -PathType Container)) {
    $ToolsPath = "C:\ProgramData\GEAerospace\MediaCreator\Tools"
}
$SourcesZip = Join-Path $CachePath "Boot\Sources.zip"
$ControlPath = Join-Path $DeployPath "Control"

if (-not (Test-Path $DeployPath -PathType Container)) {
    Write-Host "ERROR: Deploy directory not found at $DeployPath" -ForegroundColor Red
    Write-Host "  .\Upload-Image.ps1 -CachePath ""D:\Path\To\Cache""" -ForegroundColor Yellow
    exit 2
}

# --- MCL finished? One tag at the cache root, written at the very end ---
$tags = @(Get-ChildItem -LiteralPath $CachePath -Filter "*.tag" -File -ErrorAction SilentlyContinue)
if ($tags.Count -eq 0) {
    Stop-Cache @("No MCL completion tag (*.tag) in $CachePath.",
                 "MCL writes it only when a run finishes: the cache is incomplete or the run was interrupted.",
                 "Run Media Creator Lite again and let it finish.")
}
if ($tags.Count -gt 1) {
    Stop-Cache (@("More than one MCL tag in ${CachePath}:") + ($tags | ForEach-Object { "  $($_.Name)" }) +
                @("MCL leaves exactly one. Run Media Creator Lite again."))
}
$tag = $tags[0]
$tagBytes = [System.IO.File]::ReadAllBytes($tag.FullName)
if ($tagBytes.Length -lt 8) { Stop-Cache @("MCL tag $($tag.Name) is damaged ($($tagBytes.Length) bytes). Run MCL again.") }
# MCL writes DateTime.Now.Ticks (local time) into the tag (TagFactory).
try { $tagTime = New-Object DateTime ([BitConverter]::ToInt64($tagBytes, 0)) }
catch { Stop-Cache @("MCL tag $($tag.Name) holds no valid time. Run MCL again.") }
$tagOs = $tag.BaseName -replace '-(Production|QA)$', ''

$late = @(Get-ChildItem -LiteralPath $ControlPath -File -ErrorAction SilentlyContinue |
          Where-Object { $_.LastWriteTime -gt $tagTime.AddSeconds(5) })
if ($late.Count -gt 0) {
    Stop-Cache (@("These Control files are newer than the MCL finish tag ($tagTime):") +
                ($late | ForEach-Object { "  $($_.Name)  $($_.LastWriteTime)" }) +
                @("MCL was still writing or crashed. Run Media Creator Lite again and let it finish."))
}

# --- Every Control / Tools JSON must be non-empty and parse ---
$badJson = @()
foreach ($dir in @($ControlPath, $ToolsPath)) {
    if (-not (Test-Path $dir -PathType Container)) { continue }
    foreach ($j in Get-ChildItem -LiteralPath $dir -Filter "*.json" -File -Recurse) {
        if ($j.Length -eq 0) { $badJson += "$($j.FullName) (0 bytes)"; continue }
        try { $null = Read-JsonFile $j.FullName } catch { $badJson += "$($j.FullName) (not valid JSON)" }
    }
}
if ($badJson.Count -gt 0) {
    Stop-Cache (@("Broken JSON in the MCL cache:") + ($badJson | ForEach-Object { "  $_" }) +
                @("Delete the cache folder ($CachePath) and run Media Creator Lite again."))
}

# --- Parse user_selections.json ---
$SelectionsFile = Join-Path $ToolsPath "user_selections.json"
if (-not (Test-Path $SelectionsFile)) {
    Write-Host "ERROR: user_selections.json not found at $SelectionsFile" -ForegroundColor Red
    Write-Host "  Run Media Creator Lite first to create a configuration." -ForegroundColor Yellow
    exit 2
}

$selections = (Read-JsonFile $SelectionsFile)[0]
$selectedOsId = $selections.OperatingSystemSelection
$selectedModelIds = @($selections.HardwareModelSelection | ForEach-Object { $_.Id } | Select-Object -Unique)

# --- Parse control JSONs ---
$osJsonFile     = Join-Path $ControlPath "OperatingSystem.json"
$driverJsonFile = Join-Path $ControlPath "HardwareDriver.json"
$pkgJsonFile    = Join-Path $ControlPath "packages.json"
$langJsonFile   = Join-Path $ControlPath "LanguagePacks.json"

if (-not (Test-Path $osJsonFile)) {
    Write-Host "ERROR: OperatingSystem.json not found in $ControlPath" -ForegroundColor Red
    exit 2
}

$osJson     = @(Read-JsonFile $osJsonFile)
$driverJson = if (Test-Path $driverJsonFile) { @(Read-JsonFile $driverJsonFile) } else { @() }
$pkgJson    = if (Test-Path $pkgJsonFile)    { @(Read-JsonFile $pkgJsonFile) }    else { @() }
$langJson   = if (Test-Path $langJsonFile)   { @(Read-JsonFile $langJsonFile) }   else { @() }

# --- The cached OS must be the selected one, and the one the tag names ---
$matchedOs = @($osJson | Where-Object { $_.operatingSystemVersion.id -eq [int]$selectedOsId })
if ($matchedOs.Count -eq 0) {
    Stop-Cache @("OperatingSystem.json does not hold the OS selected in user_selections.json (ID $selectedOsId).",
                 "Run Media Creator Lite again.")
}
$osv = $matchedOs[0].operatingSystemVersion
if (-not $tagOs.StartsWith($osv.productName, [StringComparison]::OrdinalIgnoreCase)) {
    Stop-Cache @("The MCL tag says '$tagOs' but OperatingSystem.json holds '$($osv.productName)'.",
                 "The cache is a mix of two runs. Run Media Creator Lite again.")
}
$isLtsc = $osv.productName -match 'LTSC'
$suggested = if ($isLtsc) { "gea-shopfloor" } else { "standard-or-engineer" }

# --- Resolve selections to paths ---
$osDirs = @()
$osTotalSize = [long]0
foreach ($os in $matchedOs) {
    $osDirs += Resolve-DestDir $os.operatingSystemVersion.wim.DestinationDir
    $osTotalSize += $os.operatingSystemVersion.wim.size
}

# Packages: enabled + matching OS ID (osId is a comma string in the new schema)
function Get-OsIds { param($e) if ($e.aOsIds) { return @($e.aOsIds) } return @("$($e.osId)" -split ',' | ForEach-Object { $_.Trim() }) }
$matchedPkgs = @($pkgJson | Where-Object { (Get-OsIds $_) -contains "$selectedOsId" -and $_.enabled -eq 1 })
$pkgTotalSize = [long]0
foreach ($pkg in $matchedPkgs) { $pkgTotalSize += $pkg.size }

# Drivers: match selected model IDs (family) + OS ID, deduplicate by path
$allMatchingDrivers = @($driverJson | Where-Object {
    $selectedModelIds -contains $_.family -and (Get-OsIds $_) -contains "$selectedOsId"
})
$allDriverDirSet = [ordered]@{}
foreach ($drv in $allMatchingDrivers) {
    $rel = Resolve-DestDir $drv.DestinationDir
    if (-not $allDriverDirSet.Contains($rel)) { $allDriverDirSet[$rel] = $drv.size }
}
$allDriverCount = $allDriverDirSet.Count
$allDriverTotalSize = [long]0
$allDriverDirSet.Values | ForEach-Object { $allDriverTotalSize += $_ }

$driverDirs = @()
$driverTotalSize = [long]0
if ($IncludeDrivers) {
    $driverDirs = @($allDriverDirSet.Keys)
    $driverTotalSize = $allDriverTotalSize
}

# --- Display upload plan ---
Write-Host "  Cache:     $CachePath"
Write-Host "  MCL run:   $tagOs, finished $tagTime" -ForegroundColor Green
Write-Host "  Server:    $Server"
Write-Host "  Target:    $(if ($isLtsc) { 'Shop Floor (LTSC)' } else { 'Standard or Engineer (pick it in the webapp)' })"
Write-Host ""
Write-Host "  Upload Plan (from user_selections.json):" -ForegroundColor Cyan
Write-Host "  ------------------------------------------"
Write-Host "  OS:        $($osv.marketingName) $($osv.productName) ($(Format-Size $osTotalSize))" -ForegroundColor Green
Write-Host "  Packages:  $($matchedPkgs.Count) update(s) ($(Format-Size $pkgTotalSize))" -ForegroundColor Green
if ($IncludeDrivers) {
    Write-Host "  Drivers:   $($driverDirs.Count) model(s) ($(Format-Size $driverTotalSize))" -ForegroundColor Green
} else {
    Write-Host "  Drivers:   SKIPPED -- $allDriverCount available, use -IncludeDrivers" -ForegroundColor Yellow
}
Write-Host "  Control:   Always included" -ForegroundColor Gray
Write-Host "  Tools:     $(if (Test-Path $ToolsPath) { 'Yes' } else { 'Not found' })" -ForegroundColor $(if (Test-Path $ToolsPath) { "Gray" } else { "Yellow" })
Write-Host "  Sources:   $(if (Test-Path $SourcesZip) { 'Yes (from Boot\Sources.zip)' } else { 'Not found' })" -ForegroundColor $(if (Test-Path $SourcesZip) { "Gray" } else { "Yellow" })
Write-Host ""

# --- Sources: extract now so it can be hashed with everything else ---
$TempExtract = $null
$TempSources = $null
if (Test-Path $SourcesZip) {
    $TempExtract = Join-Path $env:TEMP "SourcesExtract"
    Remove-Item -Recurse -Force $TempExtract -ErrorAction SilentlyContinue
    Expand-Archive $SourcesZip -DestinationPath $TempExtract -Force
    # Handle nested Sources folder (zip may contain Sources/ at root)
    $TempSources = $TempExtract
    if ((Test-Path (Join-Path $TempExtract "Sources")) -and -not (Test-Path (Join-Path $TempExtract "Diskpart"))) {
        $TempSources = Join-Path $TempExtract "Sources"
    }
}

# --- Everything that will be sent: local file + path on the share ---
$sendList = New-Object System.Collections.Generic.List[object]
function Add-Tree {
    param([string]$LocalDir, [string]$RelPrefix, [string[]]$SkipDirs = @())
    if (-not (Test-Path -LiteralPath $LocalDir -PathType Container)) { return }
    $base = (Get-Item -LiteralPath $LocalDir).FullName.TrimEnd('\')
    foreach ($f in Get-ChildItem -LiteralPath $base -Recurse -File -Force) {
        $rel = $f.FullName.Substring($base.Length + 1)
        if ($SkipDirs.Count -gt 0) {
            # robocopy /XD <name> skips that folder name at any depth
            $segs = $rel.Split('\')
            $skip = $false
            for ($i = 0; $i -lt $segs.Count - 1; $i++) { if ($SkipDirs -contains $segs[$i]) { $skip = $true; break } }
            if ($skip) { continue }
        }
        $sendList.Add([pscustomobject]@{ Local = $f.FullName; Rel = (Join-Path $RelPrefix $rel); Size = $f.Length })
    }
}
Add-Tree $DeployPath "Deploy" $ExcludedDeployDirs
foreach ($d in $osDirs) { Add-Tree (Join-Path $CachePath $d) $d }
foreach ($pkg in $matchedPkgs) {
    $pdir = Resolve-DestDir $pkg.destinationDir
    $rel = if ($pdir) { Join-Path $pdir $pkg.fileName } else { $pkg.fileName }
    $local = Join-Path $CachePath $rel
    if (Test-Path -LiteralPath $local -PathType Leaf) {
        # File name as robocopy will write it (disk case), not as the JSON spells it
        $item = Get-Item -LiteralPath $local
        $rel = Join-Path (Split-Path $rel -Parent) $item.Name
        $sendList.Add([pscustomobject]@{ Local = $item.FullName; Rel = $rel; Size = $item.Length })
    }
}
foreach ($d in $driverDirs) { Add-Tree (Join-Path $CachePath $d) $d }
Add-Tree $ToolsPath "Tools"
if ($TempSources) { Add-Tree $TempSources "Sources" }

# The same file can be reached twice (an OS folder outside the excluded
# Deploy dirs, nested driver folders, two packages naming one file); a
# duplicate path would make the server refuse the whole manifest.
$seenRel = @{}
$dedup = New-Object System.Collections.Generic.List[object]
foreach ($f in $sendList) {
    $k = ($f.Rel -replace '\\', '/').ToLower()
    if (-not $seenRel.ContainsKey($k)) { $seenRel[$k] = $true; $dedup.Add($f) }
}
$sendList = $dedup

# --- Hash everything; check payload against MCL's published hashes ---
Write-Step "Hashing $($sendList.Count) file(s) ($(Format-Size (($sendList | Measure-Object Size -Sum).Sum)))"
$mclHash = @{}
function Add-Mcl { param($Dest, $Name, $Hash)
    # Only a real SHA-256 (as the server checks): anything else is not ours to judge
    if (-not ($Dest -and $Name -and ("$Hash" -match '^[0-9a-fA-F]{64}$'))) { return }
    $dd = Resolve-DestDir $Dest
    $key = if ($dd) { Join-Path $dd $Name } else { "$Name" }
    $mclHash[$key.ToLower()] = "$Hash".ToLower() }
foreach ($e in $osJson) {
    $w = $e.operatingSystemVersion.wim
    $n = if ($w.fileName) { $w.fileName } else { "install.wim" }
    Add-Mcl $w.destinationDir $n $w.hash
}
foreach ($e in @($driverJson) + @($pkgJson)) { if ($e) { Add-Mcl $e.destinationDir $e.fileName $e.hash } }
foreach ($e in $langJson) { foreach ($lf in @($e.files)) { if ($lf) { Add-Mcl $lf.destFolder $lf.fileName $lf.hash } } }

$totalBytes = [long](($sendList | Measure-Object Size -Sum).Sum)
$doneBytes = [long]0
$badCache = @()
$i = 0
foreach ($f in $sendList) {
    $i++
    $pct = if ($totalBytes -gt 0) { [int](100 * $doneBytes / $totalBytes) } else { 100 }
    Write-Progress -Activity "Hashing" -Status "$i/$($sendList.Count) $($f.Rel)" -PercentComplete $pct
    $f | Add-Member -NotePropertyName Sha256 -NotePropertyValue (Get-FileHash -LiteralPath $f.Local -Algorithm SHA256).Hash.ToLower()
    $doneBytes += $f.Size
    $want = $mclHash[$f.Rel.ToLower()]
    if ($want -and $want -ne $f.Sha256) { $badCache += $f }
}
Write-Progress -Activity "Hashing" -Completed
if ($badCache.Count -gt 0) {
    Write-Host ""
    Write-Host "These cache files do not match the hash MCL publishes for them (damaged download):" -ForegroundColor Red
    foreach ($b in $badCache) { Write-Host "  $($b.Rel)" -ForegroundColor Red }
    Write-Host ""
    Write-Host "MCL re-downloads a MISSING file correctly, but trusts a damaged one of the right size." -ForegroundColor Yellow
    $ans = Read-Host "Delete these $($badCache.Count) file(s) from the MCL cache so MCL downloads them again? (Y/N)"
    if ($ans -match '^[Yy]') {
        foreach ($b in $badCache) { Remove-Item -LiteralPath $b.Local -Force; Write-Host "  deleted $($b.Local)" -ForegroundColor Gray }
        Write-Host "Now run Media Creator Lite again, then this script." -ForegroundColor Yellow
    }
    if ($TempExtract) { Remove-Item -Recurse -Force $TempExtract -ErrorAction SilentlyContinue }
    Write-Host "Nothing was uploaded." -ForegroundColor Red
    exit 2
}
Write-Host "  All $($sendList.Count) file(s) hashed; payload matches MCL's published hashes." -ForegroundColor Green

$pcName = ($env:COMPUTERNAME -replace '[^A-Za-z0-9._-]', '')
$uploadId = "{0}-{1}" -f (Get-Date -Format "yyyyMMdd-HHmmss"), $pcName
if ($uploadId.Length -gt 80) { $uploadId = $uploadId.Substring(0, 80) }
$manifest = [ordered]@{
    manifest_version = 1
    upload_id = $uploadId
    created = (Get-Date -Format "yyyy-MM-ddTHH:mm:sszzz")
    source_pc = $pcName
    mcl = [ordered]@{
        tag = $tag.Name
        finished = $tagTime.ToString("yyyy-MM-ddTHH:mm:ss")
        os_id = [int]$osv.id
        os_name = "$($osv.productName)"
        os_version = "$($osv.versionNumber)"
        build = "$($osv.buildNumber)"
        suggested_target = $suggested
    }
    files = @($sendList | ForEach-Object { [ordered]@{ path = ($_.Rel -replace '\\', '/'); size = [long]$_.Size; sha256 = $_.Sha256 } })
}
$manifestBytes = [System.Text.Encoding]::UTF8.GetBytes(($manifest | ConvertTo-Json -Depth 6 -Compress))

# --- Is the webapp reachable? (before spending time on the copy) ---
$ping = Invoke-Pxe -Method GET -Path "/api/upload/verify" -TimeoutSec 15
if ($ping.Status -ne 200) {
    Write-Host "ERROR: the PXE webapp at $ApiBase does not answer ($($ping.Status) $($ping.Body))." -ForegroundColor Red
    Write-Host "  The upload cannot be verified without it. Check the server, then run this again." -ForegroundColor Yellow
    if ($TempExtract) { Remove-Item -Recurse -Force $TempExtract -ErrorAction SilentlyContinue }
    exit 1
}

# --- Connect to SMB share ---
Write-Host ""
Write-Host "Connecting to $Share ..." -ForegroundColor Gray
# net use goes through cmd: in Windows PowerShell 5.1 a native command's
# redirected stderr becomes an error record, which ErrorAction Stop turns
# into a crash ("The network connection could not be found").
cmd /c "net use ""$Share"" /delete /y >nul 2>&1" | Out-Null

$netResult = cmd /c "net use ""$Share"" /user:$User ""$Pass"" 2>&1"
if ($LASTEXITCODE -ne 0) {
    Write-Host "ERROR: Could not connect to $Share" -ForegroundColor Red
    Write-Host $netResult -ForegroundColor Red
    Write-Host ""
    Write-Host "Make sure:" -ForegroundColor Yellow
    Write-Host "  - The PXE server is running at $Server" -ForegroundColor Yellow
    Write-Host "  - This PC is on the 172.16.9.x imaging network" -ForegroundColor Yellow
    Write-Host "  - Samba is running on the PXE server" -ForegroundColor Yellow
    if ($TempExtract) { Remove-Item -Recurse -Force $TempExtract -ErrorAction SilentlyContinue }
    exit 1
}
Write-Host "Connected." -ForegroundColor Green

$failed = $false
$stepNum = 0
$totalSteps = 2  # Deploy base + OS always
if ($matchedPkgs.Count -gt 0) { $totalSteps++ }
if ($IncludeDrivers -and $driverDirs.Count -gt 0) { $totalSteps++ }
if (Test-Path $ToolsPath -PathType Container) { $totalSteps++ }
if ($TempSources) { $totalSteps++ }

# robocopy exit codes 8 and up are failures
$rcArgs = @("/R:3", "/W:5", "/NP", "/ETA")

# --- Step: Deploy base (Control, Applications, config -- skip big dirs) ---
$stepNum++
Write-Step "[$stepNum/$totalSteps] Copying Deploy\ base (Control, Applications, config) ..."
robocopy $DeployPath "$Share\Deploy" /E /XD @ExcludedDeployDirs @rcArgs
if ($LASTEXITCODE -ge 8) { Write-Host "ERROR: Deploy base copy failed (exit code $LASTEXITCODE)" -ForegroundColor Red; $failed = $true }

# --- Step: Operating System ---
$stepNum++
Write-Step "[$stepNum/$totalSteps] Copying Operating System ..."
foreach ($osDir in $osDirs) {
    $src = Join-Path $CachePath $osDir
    if (Test-Path -LiteralPath $src -PathType Container) {
        Write-Host "  $osDir" -ForegroundColor Gray
        robocopy $src (Join-Path $Share $osDir) /E @rcArgs
        if ($LASTEXITCODE -ge 8) { Write-Host "ERROR: OS copy failed (exit code $LASTEXITCODE)" -ForegroundColor Red; $failed = $true }
    } else {
        Write-Host "  SKIPPED (not cached): $osDir" -ForegroundColor Yellow
    }
}

# --- Step: Packages ---
if ($matchedPkgs.Count -gt 0) {
    $stepNum++
    Write-Step "[$stepNum/$totalSteps] Copying Packages ..."
    $pkgGroups = [ordered]@{}
    foreach ($pkg in $matchedPkgs) {
        $rel = Resolve-DestDir $pkg.destinationDir
        if (-not $pkgGroups.Contains($rel)) { $pkgGroups[$rel] = @() }
        $pkgGroups[$rel] += $pkg.fileName
    }
    foreach ($dir in $pkgGroups.Keys) {
        $src = Join-Path $CachePath $dir
        if (Test-Path -LiteralPath $src -PathType Container) {
            foreach ($f in $pkgGroups[$dir]) { Write-Host "  $f" -ForegroundColor Gray }
            $robocopyArgs = @($src, (Join-Path $Share $dir)) + $pkgGroups[$dir] + $rcArgs
            & robocopy @robocopyArgs
            if ($LASTEXITCODE -ge 8) { Write-Host "ERROR: Package copy failed (exit code $LASTEXITCODE)" -ForegroundColor Red; $failed = $true }
        } else {
            Write-Host "  SKIPPED (not cached): $dir" -ForegroundColor Yellow
        }
    }
}

# --- Step: Drivers (only with -IncludeDrivers) ---
if ($IncludeDrivers -and $driverDirs.Count -gt 0) {
    $stepNum++
    Write-Step "[$stepNum/$totalSteps] Copying Drivers ($($driverDirs.Count) model(s)) ..."
    $drvCopied = 0
    foreach ($drvDir in $driverDirs) {
        $drvCopied++
        $src = Join-Path $CachePath $drvDir
        if (Test-Path -LiteralPath $src -PathType Container) {
            Write-Host "  [$drvCopied/$($driverDirs.Count)] $drvDir" -ForegroundColor Gray
            robocopy $src (Join-Path $Share $drvDir) /E @rcArgs
            if ($LASTEXITCODE -ge 8) { Write-Host "ERROR: Driver copy failed (exit code $LASTEXITCODE)" -ForegroundColor Red; $failed = $true }
        } else {
            Write-Host "  SKIPPED (not cached): $drvDir" -ForegroundColor Yellow
        }
    }
}

# --- Step: Tools ---
if (Test-Path $ToolsPath -PathType Container) {
    $stepNum++
    Write-Step "[$stepNum/$totalSteps] Copying Tools\ ..."
    robocopy $ToolsPath "$Share\Tools" /E @rcArgs
    if ($LASTEXITCODE -ge 8) { Write-Host "ERROR: Tools copy failed (exit code $LASTEXITCODE)" -ForegroundColor Red; $failed = $true }
}

# --- Step: Sources ---
if ($TempSources) {
    $stepNum++
    Write-Step "[$stepNum/$totalSteps] Copying Sources\ ..."
    robocopy $TempSources "$Share\Sources" /E @rcArgs
    if ($LASTEXITCODE -ge 8) { Write-Host "ERROR: Sources copy failed (exit code $LASTEXITCODE)" -ForegroundColor Red; $failed = $true }
}

# --- Verify on the server; re-send what is wrong ---
$byRel = @{}
foreach ($f in $sendList) { $byRel[($f.Rel -replace '\\', '/').ToLower()] = $f }
$verified = $false
$result = $null
for ($round = 1; $round -le $MaxRounds + 1; $round++) {
    Write-Step "Verifying on the server (round $round) ..."
    $r = Invoke-Pxe -Method POST -Path "/api/upload/verify" -Body $manifestBytes -TimeoutSec 120
    if ($r.Status -ne 202) {
        Write-Host "ERROR: verify request refused ($($r.Status)): $(Get-ErrorText $r)" -ForegroundColor Red
        break
    }
    $deadline = (Get-Date).AddHours(2)
    $postedAt = Get-Date
    do {
        if ((Get-Date) -gt $deadline) { $result = $null; break }
        if ($result -and $result.upload_id -ne $uploadId -and ((Get-Date) - $postedAt).TotalSeconds -gt 120) {
            Write-Host "ERROR: the server is working on another upload ($($result.upload_id)). Run this again later." -ForegroundColor Red
            $result = $null; break
        }
        Start-Sleep -Seconds 2
        $s = Invoke-Pxe -Method GET -Path "/api/upload/verify" -TimeoutSec 30
        $result = $null
        if ($s.Status -eq 200) { $result = $s.Body | ConvertFrom-Json }
        if ($result -and $result.state -eq "running" -and $result.total_bytes -gt 0) {
            Write-Progress -Activity "Server is hashing the upload" -Status "$($result.done_files)/$($result.total_files) files" `
                -PercentComplete ([Math]::Min(100, [int](100 * $result.done_bytes / $result.total_bytes)))
        }
    } while (-not $result -or $result.state -eq "running" -or $result.upload_id -ne $uploadId)
    Write-Progress -Activity "Server is hashing the upload" -Completed

    if (-not $result) { Write-Host "ERROR: the server did not finish verifying within 2 hours." -ForegroundColor Red; break }
    if ($result.state -eq "error") { Write-Host "ERROR: the server could not verify: $($result.error)" -ForegroundColor Red; break }
    if ($result.ok) { $verified = $true; break }

    foreach ($b in @($result.bad_in_cache)) { if ($b) { Write-Host "  BAD IN MCL CACHE: $($b.path) ($($b.reason))" -ForegroundColor Red } }
    if (@($result.bad_in_cache).Count -gt 0) { break }   # re-sending a bad source cannot help

    $redo = @(@($result.missing) + @($result.mismatch | ForEach-Object { $_.path }) | Where-Object { $_ })
    Write-Host "  $($redo.Count) file(s) missing or different on the server." -ForegroundColor Yellow
    if ($round -gt $MaxRounds) { break }
    foreach ($rel in $redo) {
        $f = $byRel[$rel.ToLower()]
        if (-not $f) { continue }
        if ($f.Size -le [long]$result.small_max) {
            $q = "?upload_id=$([uri]::EscapeDataString($uploadId))&path=$([uri]::EscapeDataString($rel))"
            $p = Invoke-Pxe -Method PUT -Path "/api/upload/file$q" -InFile $f.Local -ContentType "application/octet-stream" -TimeoutSec 300
            if ($p.Status -eq 200) { Write-Host "  re-sent over HTTP: $rel" -ForegroundColor Green }
            else { Write-Host "  HTTP re-send failed ($($p.Status)) for ${rel}: $(Get-ErrorText $p)" -ForegroundColor Red }
        } else {
            $dstDir = Join-Path $Share (Split-Path ($f.Rel) -Parent)
            $dstFile = Join-Path $dstDir (Split-Path $f.Local -Leaf)
            Write-Host "  re-copying (delete, then robocopy /J): $rel" -ForegroundColor Gray
            # robocopy calls a same-size, same-time file "same" or "modified"
            # and skips it, even with /IS /IT. Delete the bad copy so it is a
            # new file; /IS /IT /IM stay as a backstop if the delete fails.
            Remove-Item -LiteralPath $dstFile -Force -ErrorAction SilentlyContinue
            $rc = robocopy (Split-Path $f.Local -Parent) $dstDir (Split-Path $f.Local -Leaf) /J /IS /IT /IM @rcArgs
            if ($LASTEXITCODE -ge 8 -or $LASTEXITCODE -eq 0) {
                # 0 = nothing copied: robocopy skipped it again
                Write-Host "  robocopy did not copy $rel (exit $LASTEXITCODE):" -ForegroundColor Red
                $rc | Where-Object { $_ -match '\S' } | Select-Object -Last 8 | ForEach-Object { Write-Host "    $_" }
            }
        }
    }
}

# --- Disconnect and clean up ---
cmd /c "net use ""$Share"" /delete /y >nul 2>&1" | Out-Null
if ($TempExtract) { Remove-Item -Recurse -Force $TempExtract -ErrorAction SilentlyContinue }

# --- Summary ---
Write-Host ""
if ($failed -and $verified) {
    Write-Host "  (robocopy reported errors on the first pass; the files it missed were re-sent and verified.)" -ForegroundColor Yellow
}
if ($verified) {
    Write-Host "========================================" -ForegroundColor Green
    Write-Host "  Upload complete: all $($sendList.Count) file(s) verified on the server." -ForegroundColor Green
    Write-Host "========================================" -ForegroundColor Green
    if ([int]$result.extra_count -gt 0) {
        Write-Host ""
        Write-Host "  NOTE: the upload folder also holds $($result.extra_count) file(s) from an earlier upload." -ForegroundColor Yellow
        Write-Host "  Import is blocked until you remove them (button on the Image Import page)." -ForegroundColor Yellow
    }
    Write-Host ""
    Write-Host "Next steps:" -ForegroundColor Cyan
    Write-Host "  1. Open the PXE webapp: http://${Server}:$WebPort/images/import" -ForegroundColor White
    Write-Host "  2. Target: $(if ($isLtsc) { 'Shop Floor (the only choice for LTSC)' } else { 'Standard or Engineer' })" -ForegroundColor White
    Write-Host "  3. Click Start import" -ForegroundColor White
    Write-Host ""
    exit 0
}
Write-Host "========================================" -ForegroundColor Red
Write-Host "  Upload NOT verified." -ForegroundColor Red
Write-Host "========================================" -ForegroundColor Red
if ($result) {
    foreach ($m in @($result.missing)) { if ($m) { Write-Host "  missing:  $m" -ForegroundColor Red } }
    foreach ($m in @($result.mismatch)) { if ($m) { Write-Host "  differs:  $($m.path) ($($m.reason))" -ForegroundColor Red } }
}
Write-Host "The webapp will not import this upload. Run the script again." -ForegroundColor Yellow
exit 1