#
# Download-Drivers.ps1 — Download selected hardware drivers from GE CDN
#
# Reads user_selections.json and HardwareDriver.json from the MCL cache
# to download only the driver packs for your selected hardware models.
# Bypasses Media Creator Lite's unreliable download mechanism.
#
# Downloads go into the MCL cache structure so Upload-Image.ps1 can
# upload them with -IncludeDrivers.
#
# Usage:
#   .\Download-Drivers.ps1                           (download all selected models)
#   .\Download-Drivers.ps1 -ListOnly                 (show what would be downloaded)
#   .\Download-Drivers.ps1 -CachePath "D:\MCL\Cache" (custom cache location)
#   .\Download-Drivers.ps1 -Force                    (re-download even if already cached)
#
# Requires internet access. Run on the workstation, not the PXE server.
#

param(
    [string]$CachePath = "C:\ProgramData\GEAerospace\MediaCreator\Cache",
    [switch]$ListOnly,
    [switch]$Force
)

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 Driver Downloader" -ForegroundColor Cyan
Write-Host "========================================" -ForegroundColor Cyan
Write-Host ""

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

if (-not (Test-Path $ControlPath -PathType Container)) {
    Write-Host "ERROR: Deploy\Control not found at $ControlPath" -ForegroundColor Red
    Write-Host "  Run Media Creator Lite first to cache the base content." -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
    exit 1
}

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

# --- Parse HardwareDriver.json ---
$driverJsonFile = Join-Path $ControlPath "HardwareDriver.json"
if (-not (Test-Path $driverJsonFile)) {
    Write-Host "ERROR: HardwareDriver.json not found in $ControlPath" -ForegroundColor Red
    exit 1
}

$driverJson = Get-Content $driverJsonFile -Raw | ConvertFrom-Json

# --- Match drivers to selections ---
$matchedDrivers = @($driverJson | Where-Object {
    $selectedModelIds -contains $_.family -and $_.aOsIds -contains $selectedOsId
})

# Deduplicate by DestinationDir (some models share a driver pack)
$uniqueDrivers = [ordered]@{}
foreach ($drv in $matchedDrivers) {
    $rel = Resolve-DestDir $drv.DestinationDir
    if (-not $uniqueDrivers.Contains($rel)) {
        $uniqueDrivers[$rel] = $drv
    }
}

$totalSize = [long]0
$uniqueDrivers.Values | ForEach-Object { $totalSize += $_.size }

# --- Display plan ---
Write-Host "  Cache:     $CachePath"
Write-Host "  OS ID:     $selectedOsId"
Write-Host "  Models:    $($selectedModelIds.Count) selected"
Write-Host "  Drivers:   $($uniqueDrivers.Count) unique pack(s) ($(Format-Size $totalSize))" -ForegroundColor Cyan
Write-Host ""

if ($uniqueDrivers.Count -eq 0) {
    Write-Host "No drivers match your selections." -ForegroundColor Yellow
    exit 0
}

# Show each driver pack
$idx = 0
foreach ($rel in $uniqueDrivers.Keys) {
    $idx++
    $drv = $uniqueDrivers[$rel]
    $localDir = Join-Path $CachePath $rel
    $cached = Test-Path $localDir -PathType Container
    $status = if ($cached -and -not $Force) { "[CACHED]" } else { "[DOWNLOAD]" }
    $color  = if ($cached -and -not $Force) { "Green" } else { "Yellow" }

    Write-Host ("  {0,3}. {1,-12} {2} ({3})" -f $idx, $status, $drv.modelsfriendlyname, (Format-Size $drv.size)) -ForegroundColor $color
    Write-Host "       $($drv.FileName)" -ForegroundColor Gray
}
Write-Host ""

if ($ListOnly) {
    Write-Host "  (list only — run without -ListOnly to download)" -ForegroundColor Gray
    exit 0
}

# --- Download and extract ---
$downloadDir = Join-Path $env:TEMP "PXE-DriverDownloads"
if (-not (Test-Path $downloadDir)) { New-Item -ItemType Directory -Path $downloadDir -Force | Out-Null }

$completed = 0
$skipped = 0
$errors = 0

foreach ($rel in $uniqueDrivers.Keys) {
    $drv = $uniqueDrivers[$rel]
    $localDir = Join-Path $CachePath $rel
    $zipFile  = Join-Path $downloadDir $drv.FileName

    Write-Host "========================================" -ForegroundColor Cyan
    Write-Host "[$($completed + $skipped + $errors + 1)/$($uniqueDrivers.Count)] $($drv.modelsfriendlyname)" -ForegroundColor Cyan
    Write-Host "========================================" -ForegroundColor Cyan

    # Skip if already cached (unless -Force)
    if ((Test-Path $localDir -PathType Container) -and -not $Force) {
        Write-Host "  Already cached at $rel" -ForegroundColor Green
        $skipped++
        Write-Host ""
        continue
    }

    # Download
    Write-Host "  Downloading $(Format-Size $drv.size) ..." -ForegroundColor Gray
    Write-Host "  URL: $($drv.url)" -ForegroundColor DarkGray

    try {
        # Use curl.exe for progress display on large files
        if (Get-Command curl.exe -ErrorAction SilentlyContinue) {
            & curl.exe -L -o $zipFile $drv.url --progress-bar --fail
            if ($LASTEXITCODE -ne 0) { throw "curl failed with exit code $LASTEXITCODE" }
        } else {
            # Fallback: WebClient (streams to disk, no buffering)
            $wc = New-Object System.Net.WebClient
            $wc.DownloadFile($drv.url, $zipFile)
        }
    } catch {
        Write-Host "  ERROR: Download failed - $_" -ForegroundColor Red
        $errors++
        Write-Host ""
        continue
    }

    # Verify SHA256 hash
    Write-Host "  Verifying SHA256 hash ..." -ForegroundColor Gray
    $actualHash = (Get-FileHash -Path $zipFile -Algorithm SHA256).Hash
    if ($actualHash -ne $drv.hash) {
        Write-Host "  ERROR: Hash mismatch!" -ForegroundColor Red
        Write-Host "    Expected: $($drv.hash)" -ForegroundColor Red
        Write-Host "    Got:      $actualHash" -ForegroundColor Red
        Remove-Item -Path $zipFile -Force -ErrorAction SilentlyContinue
        $errors++
        Write-Host ""
        continue
    }
    Write-Host "  Hash OK." -ForegroundColor Green

    # Extract to cache destination
    Write-Host "  Extracting to $rel ..." -ForegroundColor Gray
    if (Test-Path $localDir) { Remove-Item -Recurse -Force $localDir -ErrorAction SilentlyContinue }
    New-Item -ItemType Directory -Path $localDir -Force | Out-Null

    try {
        Expand-Archive -Path $zipFile -DestinationPath $localDir -Force
    } catch {
        Write-Host "  ERROR: Extraction failed - $_" -ForegroundColor Red
        $errors++
        Write-Host ""
        continue
    }

    # Clean up zip
    Remove-Item -Path $zipFile -Force -ErrorAction SilentlyContinue

    Write-Host "  Done." -ForegroundColor Green
    $completed++
    Write-Host ""
}

# --- Summary ---
Write-Host "========================================" -ForegroundColor Cyan
Write-Host "  Download Summary" -ForegroundColor Cyan
Write-Host "========================================" -ForegroundColor Cyan
Write-Host "  Downloaded: $completed" -ForegroundColor Green
if ($skipped -gt 0) { Write-Host "  Skipped:    $skipped (already cached)" -ForegroundColor Gray }
if ($errors -gt 0)  { Write-Host "  Failed:     $errors" -ForegroundColor Red }
Write-Host ""

if ($completed -gt 0 -or $skipped -gt 0) {
    Write-Host "Driver packs are in the MCL cache at:" -ForegroundColor Cyan
    Write-Host "  $DeployPath\Out-of-box Drivers\" -ForegroundColor White
    Write-Host ""
    Write-Host "To upload to the PXE server:" -ForegroundColor Cyan
    Write-Host "  .\Upload-Image.ps1 -IncludeDrivers" -ForegroundColor White
    Write-Host ""
}

# Clean up temp dir if empty
if ((Get-ChildItem $downloadDir -Force -ErrorAction SilentlyContinue | Measure-Object).Count -eq 0) {
    Remove-Item $downloadDir -Force -ErrorAction SilentlyContinue
}