# --- Config (edit if you want)
$UpdateCabUrl   = "https://catalog.sf.dl.delivery.mp.microsoft.com/filestreamingservice/files/ac20e3c8-6556-4785-8f10-2f631b4e5f5b/public/windows11.0-kb5062688-x64_87a5b5dc864ca204be7ab0ea6082ca542b275665.cab"
$MountDir       = "$env:TEMP\WinRE_Mount"
$CabPath        = "$env:TEMP\KB5062688_WinRE.cab"
$ExpectedSha1   = ""      # leave empty to infer from URL (_<40hex>.cab)
$ForceRedownload = $false # set $true to force re-download

# --- Admin check
$IsAdmin = ([Security.Principal.WindowsPrincipal] [Security.Principal.WindowsIdentity]::GetCurrent()
).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
if (-not $IsAdmin) { Write-Error "Run this script as Administrator."; exit 1 }

# TLS tweak for PS 5.1 if we fall back to Invoke-WebRequest
if ($PSVersionTable.PSVersion.Major -lt 6) {
  try { [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12 } catch {}
}

function Invoke-Exe {
  param([Parameter(Mandatory)] [string]$FilePath,
        [Parameter(Mandatory)] [string]$Arguments)
  $psi = New-Object System.Diagnostics.ProcessStartInfo
  $psi.FileName = $FilePath
  $psi.Arguments = $Arguments
  $psi.RedirectStandardOutput = $true
  $psi.RedirectStandardError  = $true
  $psi.UseShellExecute = $false
  $p = [System.Diagnostics.Process]::Start($psi)
  $stdout = $p.StandardOutput.ReadToEnd()
  $stderr = $p.StandardError.ReadToEnd()
  $p.WaitForExit()
  [pscustomobject]@{ ExitCode=$p.ExitCode; StdOut=$stdout; StdErr=$stderr; Cmd="$FilePath $Arguments" }
}

function Get-ExpectedSha1FromUrl([string]$Url) {
  try {
    $name = [System.IO.Path]::GetFileName($Url)
    if ($name -match "_([0-9a-fA-F]{40})\.cab$") { return $Matches[1].ToLowerInvariant() }
  } catch {}
  ""
}

function Get-FileSha1([string]$Path) {
  (Get-FileHash -Path $Path -Algorithm SHA1).Hash.ToLowerInvariant()
}

function Download-VerifiedCab {
  param([Parameter(Mandatory)] [string]$Url,
        [Parameter(Mandatory)] [string]$Dst,
        [string]$Sha1,
        [bool]$Force)

  # Ensure destination folder exists
  $dstDir = Split-Path -Path $Dst -Parent
  if (-not (Test-Path -LiteralPath $dstDir)) {
    New-Item -ItemType Directory -Path $dstDir -Force | Out-Null
  }

  $dstExists = Test-Path -LiteralPath $Dst
  if ($dstExists -and (-not $Force)) {
    Write-Host "CAB already present: $Dst"
  } else {
    if ($dstExists) { Remove-Item -LiteralPath $Dst -Force -ErrorAction SilentlyContinue }

    # Prefer real curl.exe (not the PS alias)
    $curl = Get-Command curl.exe -ErrorAction SilentlyContinue
    if ($curl) {
      Write-Host "Downloading with curl.exe ..."
      $args = @(
        "-L", "`"$Url`"",
        "-o", "`"$Dst`"",
        "--fail",
        "--retry", "5",
        "--retry-delay", "2",
        "--max-time", "900"
      ) -join ' '
      $res = Invoke-Exe -FilePath $curl.Path -Arguments $args
      if ($res.ExitCode -ne 0) { throw "curl.exe failed (exit $($res.ExitCode)). StdErr: $($res.StdErr)" }
    } else {
      # Fallback to BITS / IWR
      try {
        if (Get-Command Start-BitsTransfer -ErrorAction SilentlyContinue) {
          Start-BitsTransfer -Source $Url -Destination $Dst -ErrorAction Stop
        } else {
          $isPS5 = $PSVersionTable.PSVersion.Major -lt 6
          try {
            if ($isPS5) { Invoke-WebRequest -Uri $Url -OutFile $Dst -UseBasicParsing -ErrorAction Stop }
            else         { Invoke-WebRequest -Uri $Url -OutFile $Dst                 -ErrorAction Stop }
          } catch {
            if ($isPS5) { Invoke-WebRequest -Uri $Url -OutFile $Dst                 -ErrorAction Stop }
            else         { Invoke-WebRequest -Uri $Url -OutFile $Dst -UseBasicParsing -ErrorAction Stop }
          }
        }
      } catch {
        throw "Download failed via fallback: $($_.Exception.Message)"
      }
    }
  }

  if (-not (Test-Path -LiteralPath $Dst)) { throw "Download failed: file not found at $Dst" }
  $len = (Get-Item -LiteralPath $Dst).Length
  if ($len -le 0) { Remove-Item -LiteralPath $Dst -Force -ErrorAction SilentlyContinue; throw "Download failed: file size is 0 bytes" }

  if ([string]::IsNullOrWhiteSpace($Sha1)) {
    $Sha1 = Get-ExpectedSha1FromUrl $Url
    if ($Sha1) { Write-Host "Expected SHA-1 inferred from URL: $Sha1" }
  }

  if ($Sha1) {
    Write-Host "Verifying SHA-1..."
    $actual = Get-FileSha1 $Dst
    if ($actual -ne $Sha1) {
      Remove-Item -LiteralPath $Dst -Force -ErrorAction SilentlyContinue
      throw "Hash mismatch. Expected: $Sha1  Actual: $actual"
    }
    Write-Host "SHA-1 OK."
  } else {
    Write-Host "No expected SHA-1 provided or inferred. Skipping hash verification."
  }
}

# ------------------- main flow (no finally) -------------------

# 1) Download + verify
Download-VerifiedCab -Url $UpdateCabUrl -Dst $CabPath -Sha1 $ExpectedSha1 -Force:$ForceRedownload

# 2) Prep mount dir
if (Test-Path -LiteralPath $MountDir) {
  Remove-Item -LiteralPath $MountDir -Recurse -Force -ErrorAction SilentlyContinue
}
New-Item -ItemType Directory -Path $MountDir -Force | Out-Null

$mounted  = $false
$serviced = $false

try {
  Write-Host "Mounting WinRE to $MountDir ..."
  $m1 = Invoke-Exe -FilePath "reagentc.exe" -Arguments "/mountre /path `"$MountDir`""
  if ($m1.ExitCode -ne 0) {
    Write-Warning "Initial mount failed (exit $($m1.ExitCode)). Attempting to enable WinRE and retry..."
    $en = Invoke-Exe -FilePath "reagentc.exe" -Arguments "/enable"
    if ($en.ExitCode -ne 0) {
      if (Test-Path -LiteralPath $MountDir) { Remove-Item -LiteralPath $MountDir -Recurse -Force -ErrorAction SilentlyContinue }
      throw "Failed to enable WinRE (exit $($en.ExitCode)). StdErr: $($en.StdErr)"
    }
    $m2 = Invoke-Exe -FilePath "reagentc.exe" -Arguments "/mountre /path `"$MountDir`""
    if ($m2.ExitCode -ne 0) {
      if (Test-Path -LiteralPath $MountDir) { Remove-Item -LiteralPath $MountDir -Recurse -Force -ErrorAction SilentlyContinue }
      throw "WinRE mount still failing (exit $($m2.ExitCode)). StdOut: $($m2.StdOut)  StdErr: $($m2.StdErr)"
    } else { $mounted = $true }
  } else { $mounted = $true }

  Write-Host "Adding package to mounted WinRE ..."
  $add = Invoke-Exe -FilePath "dism.exe" -Arguments "/English /Image:`"$MountDir`" /Add-Package /PackagePath:`"$CabPath`""
  if ($add.ExitCode -ne 0) {
    if ($mounted) {
      $ud = Invoke-Exe -FilePath "reagentc.exe" -Arguments "/unmountre /path `"$MountDir`" /discard"
      if ($ud.ExitCode -ne 0) { Write-Warning "Discard failed (exit $($ud.ExitCode)). StdErr: $($ud.StdErr)" }
    }
    if (Test-Path -LiteralPath $MountDir) { Remove-Item -LiteralPath $MountDir -Recurse -Force -ErrorAction SilentlyContinue }
    throw "DISM Add-Package failed (exit $($add.ExitCode)). StdOut: $($add.StdOut)  StdErr: $($add.StdErr)"
  }
  $serviced = $true

  Write-Host "Committing changes and unmounting ..."
  $un = Invoke-Exe -FilePath "reagentc.exe" -Arguments "/unmountre /path `"$MountDir`" /commit"
  if ($un.ExitCode -ne 0) {
    Write-Warning "Commit failed (exit $($un.ExitCode)). Trying /discard ..."
    $ud = Invoke-Exe -FilePath "reagentc.exe" -Arguments "/unmountre /path `"$MountDir`" /discard"
    if ($ud.ExitCode -ne 0) { Write-Warning "Discard also failed (exit $($ud.ExitCode)). Manual cleanup may be required." }
    if (Test-Path -LiteralPath $MountDir) { Remove-Item -LiteralPath $MountDir -Recurse -Force -ErrorAction SilentlyContinue }
    throw "Failed to unmount with commit (exit $($un.ExitCode)). StdErr: $($un.StdErr)"
  }
  $mounted = $false

  # Success cleanup
  if (Test-Path -LiteralPath $MountDir) {
    Remove-Item -LiteralPath $MountDir -Recurse -Force -ErrorAction SilentlyContinue
  }
  Write-Host "`nKB5062688 has been applied to WinRE successfully."
}
catch {
  Write-Warning $_
  if ($mounted) {
    $ud = Invoke-Exe -FilePath "reagentc.exe" -Arguments "/unmountre /path `"$MountDir`" /discard"
    if ($ud.ExitCode -ne 0) { Write-Warning "Discard failed (exit $($ud.ExitCode)). StdErr: $($ud.StdErr)" }
  }
  if (Test-Path -LiteralPath $MountDir) {
    Remove-Item -LiteralPath $MountDir -Recurse -Force -ErrorAction SilentlyContinue
  }
  throw
}