#
# Upload-Image.ps1 — Copy MCL cached image to the PXE server
#
# Reads user_selections.json to upload only the selected OS, matching
# packages, and config files. Drivers are EXCLUDED by default.
#
# 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 10.9.100.1           (custom server IP)
#
# After upload, use the PXE webapp (http://10.9.100.1:9009) to import
# the uploaded content into the desired image type.
#

param(
    [string]$CachePath = "C:\ProgramData\GEAerospace\MediaCreator\Cache",
    [string]$Server = "10.9.100.1",
    [string]$User = "pxe-upload",
    [string]$Pass = "pxe",
    [switch]$IncludeDrivers
)

$Share = "\\$Server\image-upload"

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)
    return ($Dir -replace '^\*destinationdir\*\\?', '')
}

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

# --- Validate source paths ---
$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"

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 1
}

# --- 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 1
}

$selections = (Get-Content $SelectionsFile -Raw | ConvertFrom-Json)[0]
$selectedOsId = $selections.OperatingSystemSelection
$selectedModelIds = @($selections.HardwareModelSelection | ForEach-Object { $_.Id } | Select-Object -Unique)

# --- Parse control JSONs ---
$ControlPath = Join-Path $DeployPath "Control"

$osJsonFile     = Join-Path $ControlPath "OperatingSystem.json"
$driverJsonFile = Join-Path $ControlPath "HardwareDriver.json"
$pkgJsonFile    = Join-Path $ControlPath "packages.json"

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

$osJson     = Get-Content $osJsonFile -Raw | ConvertFrom-Json
$driverJson = if (Test-Path $driverJsonFile) { Get-Content $driverJsonFile -Raw | ConvertFrom-Json } else { @() }
$pkgJson    = if (Test-Path $pkgJsonFile)    { Get-Content $pkgJsonFile -Raw | ConvertFrom-Json }    else { @() }

# --- Resolve selections to paths ---

# OS: match OperatingSystemSelection ID to OperatingSystem.json entries
$matchedOs = @($osJson | Where-Object { $_.operatingSystemVersion.id -eq [int]$selectedOsId })
$osDirs = @()
$osTotalSize = [long]0
foreach ($os in $matchedOs) {
    $rel = Resolve-DestDir $os.operatingSystemVersion.wim.DestinationDir
    $osDirs += $rel
    $osTotalSize += $os.operatingSystemVersion.wim.size
}

# Packages: enabled + matching OS ID
$matchedPkgs = @($pkgJson | Where-Object { $_.aOsIds -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 $_.aOsIds -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 "  Server:    $Server"
Write-Host ""
Write-Host "  Upload Plan (from user_selections.json):" -ForegroundColor Cyan
Write-Host "  ------------------------------------------"

if ($matchedOs.Count -gt 0) {
    $osName = $matchedOs[0].operatingSystemVersion.marketingName
    Write-Host "  OS:        $osName ($(Format-Size $osTotalSize))" -ForegroundColor Green
} else {
    Write-Host "  OS:        No match for selection ID $selectedOsId" -ForegroundColor Red
}

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 ""

# --- Connect to SMB share ---
Write-Host "Connecting to $Share ..." -ForegroundColor Gray
net use $Share /delete 2>$null | Out-Null

$netResult = 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 10.9.100.x network" -ForegroundColor Yellow
    Write-Host "  - Samba is running on the PXE server" -ForegroundColor Yellow
    exit 1
}
Write-Host "Connected." -ForegroundColor Green
Write-Host ""

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

# --- Step: Deploy base (Control, Applications, config -- skip big dirs) ---
$stepNum++
Write-Host "========================================" -ForegroundColor Cyan
Write-Host "[$stepNum/$totalSteps] Copying Deploy\ base (Control, Applications, config) ..." -ForegroundColor Cyan
Write-Host "========================================" -ForegroundColor Cyan

robocopy $DeployPath "$Share\Deploy" /E /XD "Operating Systems" "Out-of-box Drivers" "Packages" /R:3 /W:5 /NP /ETA
if ($LASTEXITCODE -ge 8) {
    Write-Host "ERROR: Deploy base copy failed (exit code $LASTEXITCODE)" -ForegroundColor Red
    $failed = $true
}

# --- Step: Operating System ---
if ($matchedOs.Count -gt 0) {
    $stepNum++
    Write-Host ""
    Write-Host "========================================" -ForegroundColor Cyan
    Write-Host "[$stepNum/$totalSteps] Copying Operating System ..." -ForegroundColor Cyan
    Write-Host "========================================" -ForegroundColor Cyan

    foreach ($osDir in $osDirs) {
        $src = Join-Path $CachePath $osDir
        $dst = Join-Path $Share $osDir
        if (Test-Path $src -PathType Container) {
            Write-Host "  $osDir" -ForegroundColor Gray
            robocopy $src $dst /E /R:3 /W:5 /NP /ETA
            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-Host ""
    Write-Host "========================================" -ForegroundColor Cyan
    Write-Host "[$stepNum/$totalSteps] Copying Packages ..." -ForegroundColor Cyan
    Write-Host "========================================" -ForegroundColor Cyan

    # Group packages by destination directory for efficient robocopy
    $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
        $dst = Join-Path $Share $dir
        $files = $pkgGroups[$dir]
        if (Test-Path $src -PathType Container) {
            foreach ($f in $files) { Write-Host "  $f" -ForegroundColor Gray }
            $robocopyArgs = @($src, $dst) + $files + @("/R:3", "/W:5", "/NP", "/ETA")
            & 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-Host ""
    Write-Host "========================================" -ForegroundColor Cyan
    Write-Host "[$stepNum/$totalSteps] Copying Drivers ($($driverDirs.Count) model(s)) ..." -ForegroundColor Cyan
    Write-Host "========================================" -ForegroundColor Cyan

    $drvCopied = 0
    foreach ($drvDir in $driverDirs) {
        $drvCopied++
        $src = Join-Path $CachePath $drvDir
        $dst = Join-Path $Share $drvDir
        if (Test-Path $src -PathType Container) {
            Write-Host "  [$drvCopied/$($driverDirs.Count)] $drvDir" -ForegroundColor Gray
            robocopy $src $dst /E /R:3 /W:5 /NP /ETA
            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-Host ""
    Write-Host "========================================" -ForegroundColor Cyan
    Write-Host "[$stepNum/$totalSteps] Copying Tools\ ..." -ForegroundColor Cyan
    Write-Host "========================================" -ForegroundColor Cyan

    robocopy $ToolsPath "$Share\Tools" /E /R:3 /W:5 /NP /ETA
    if ($LASTEXITCODE -ge 8) {
        Write-Host "ERROR: Tools copy failed (exit code $LASTEXITCODE)" -ForegroundColor Red
        $failed = $true
    }
}

# --- Step: Sources ---
$TempSources = $null
if (Test-Path $SourcesZip) {
    $stepNum++
    Write-Host ""
    Write-Host "========================================" -ForegroundColor Cyan
    Write-Host "[$stepNum/$totalSteps] Extracting and copying Sources\ ..." -ForegroundColor Cyan
    Write-Host "========================================" -ForegroundColor Cyan

    $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"
    }

    robocopy $TempSources "$Share\Sources" /E /R:3 /W:5 /NP /ETA
    if ($LASTEXITCODE -ge 8) {
        Write-Host "ERROR: Sources copy failed (exit code $LASTEXITCODE)" -ForegroundColor Red
        $failed = $true
    }

    Remove-Item -Recurse -Force $TempExtract -ErrorAction SilentlyContinue
}

# --- Verify small files (SMB write-cache workaround) ---
Write-Host ""
Write-Host "========================================" -ForegroundColor Cyan
Write-Host "Verifying small files ..." -ForegroundColor Cyan
Write-Host "========================================" -ForegroundColor Cyan

$fixCount = 0
$verifyPairs = @(
    @{ Local = (Join-Path $DeployPath "Control"); Remote = "$Share\Deploy\Control" }
)
if (Test-Path $ToolsPath -PathType Container) {
    $verifyPairs += @{ Local = $ToolsPath; Remote = "$Share\Tools" }
}

foreach ($pair in $verifyPairs) {
    $localDir  = $pair.Local
    $remoteDir = $pair.Remote
    if (-not (Test-Path $localDir -PathType Container)) { continue }

    Get-ChildItem -Path $localDir -Recurse -File -ErrorAction SilentlyContinue |
        Where-Object { $_.Length -gt 0 -and $_.Length -lt 1MB } |
        ForEach-Object {
            $rel = $_.FullName.Substring($localDir.Length)
            $dstFile = Join-Path $remoteDir $rel
            if (Test-Path $dstFile) {
                $dstSize = (Get-Item $dstFile).Length
                if ($dstSize -ne $_.Length) {
                    Write-Host "  Fixing: $rel ($dstSize -> $($_.Length) bytes)" -ForegroundColor Yellow
                    $bytes = [System.IO.File]::ReadAllBytes($_.FullName)
                    [System.IO.File]::WriteAllBytes($dstFile, $bytes)
                    $fixCount++
                }
            }
        }
}

if ($fixCount -eq 0) {
    Write-Host "  All files verified OK." -ForegroundColor Green
} else {
    Write-Host "  Fixed $fixCount file(s)." -ForegroundColor Yellow
}

# --- Disconnect ---
net use $Share /delete 2>$null | Out-Null

# --- Summary ---
Write-Host ""
if ($failed) {
    Write-Host "========================================" -ForegroundColor Red
    Write-Host "  Upload completed with errors." -ForegroundColor Red
    Write-Host "========================================" -ForegroundColor Red
} else {
    Write-Host "========================================" -ForegroundColor Green
    Write-Host "  Upload complete!" -ForegroundColor Green
    Write-Host "========================================" -ForegroundColor Green
}
Write-Host ""
Write-Host "Next steps:" -ForegroundColor Cyan
Write-Host "  1. Open the PXE webapp: http://$Server`:9009" -ForegroundColor White
Write-Host "  2. Go to Image Import" -ForegroundColor White
Write-Host "  3. Select source 'image-upload' and target image type" -ForegroundColor White
Write-Host "  4. Click Import" -ForegroundColor White
Write-Host ""