#Requires -RunAsAdministrator
<#
.SYNOPSIS
    Downloads a Windows ESD from Microsoft (via MCT catalog proxy) and integrates
    the latest cumulative update into a bootable ISO.
.DESCRIPTION
    Steps:
      1. Fetch ESD URL + SHA1 from https://vesperos.chillcraft.me/windows/esd.php
      2. Download and verify ESD (SHA1)
      3. Build ISO staging structure (media skeleton, boot.wim, install.wim)
      4. Download latest cumulative update from Microsoft Update Catalog
      5. Apply CU to install.wim with DISM (mount → add-package → cleanup → commit)
      6. Create final bootable ISO with oscdimg
.PARAMETER Mode
    Win11x64, Win11ARM64, Win10x64, Win10ARM64
.PARAMETER OutputDir
    Directory for output ISO (default: current directory).
.PARAMETER Language
    Full name ("English (United States)") or short code ("en-us").
.PARAMETER Build
    Optional substring to pin a specific build (e.g. "26100", "19045").
.PARAMETER Edition
    Windows edition to extract (default: Professional).
.PARAMETER ImageIndex
    Override which ESD index to use as the install image.
.PARAMETER SkipUpdate
    Skip cumulative update download and integration.
.PARAMETER KeepWork
    Keep the working directory after completion (for debugging).
.EXAMPLE
    .\get-updated-windows-iso.ps1 -Mode Win11x64 -OutputDir C:\ISOs
.EXAMPLE
    .\get-updated-windows-iso.ps1 -Mode Win10x64 -Build "19045" -SkipUpdate
.EXAMPLE
    .\get-updated-windows-iso.ps1 -Mode Win11ARM64 -Language "fr-fr" -OutputDir D:\ISOs
#>
param(
    [Parameter(Mandatory)]
    [ValidateSet("Win11x64", "Win11ARM64", "Win10x64", "Win10ARM64")]
    [string]$Mode,

    [string]$OutputDir  = ".",
    [string]$Language   = "English (United States)",
    [string]$Build,
    [string]$Edition    = "Professional",
    [string]$ImageIndex,
    [switch]$SkipUpdate,
    [switch]$KeepWork
)

Set-StrictMode -Off
$ErrorActionPreference = "Stop"

$OsVer = if ($Mode -like "Win11*") { "11" } else { "10" }
$Arch  = if ($Mode -like "*ARM64") { "arm64" } else { "x64" }

# --- Helpers ------------------------------------------------------------------

function Write-Status { param([string]$Msg) Write-Host "`n==> $Msg" -ForegroundColor Cyan }
function Write-Step   { param([string]$Msg) Write-Host "  - $Msg" }

function Get-OscdimgPath {
    foreach ($hive in @("HKLM:\SOFTWARE\WOW6432Node\Microsoft\Windows Kits\Installed Roots",
                         "HKLM:\SOFTWARE\Microsoft\Windows Kits\Installed Roots")) {
        $root = (Get-ItemProperty $hive -ErrorAction SilentlyContinue)."KitsRoot10"
        if ($root) {
            $p = Join-Path ($root.TrimEnd('\')) "Assessment and Deployment Kit\Deployment Tools\$($Env:PROCESSOR_ARCHITECTURE)\Oscdimg\oscdimg.exe"
            if (Test-Path $p) { return $p }
        }
    }
    $local = Join-Path $PSScriptRoot "..\oscdimg.exe"
    if (Test-Path $local) { return (Resolve-Path $local).Path }
    Write-Step "Downloading oscdimg.exe..."
    Invoke-WebRequest -Uri "https://msdl.microsoft.com/download/symbols/oscdimg.exe/3D44737265000/oscdimg.exe" `
        -OutFile $local -UseBasicParsing -ErrorAction Stop
    return (Resolve-Path $local).Path
}

function Confirm-Hash {
    param([string]$Path, [string]$Expected, [ValidateSet("SHA1","SHA256")][string]$Algo = "SHA1")
    $actual = (Get-FileHash -Path $Path -Algorithm $Algo).Hash
    if ($actual -ne $Expected.ToUpper()) {
        Write-Warning "Hash mismatch for $(Split-Path $Path -Leaf)"
        Write-Warning "  Expected : $Expected"
        Write-Warning "  Actual   : $actual"
        return $false
    }
    return $true
}

$LanguageMap = [ordered]@{
    "Arabic"                   = "ar-sa"; "Bulgarian"                = "bg-bg"
    "Chinese (Simplified)"     = "zh-cn"; "Chinese (Traditional)"    = "zh-tw"
    "Croatian"                 = "hr-hr"; "Czech"                    = "cs-cz"
    "Danish"                   = "da-dk"; "Dutch"                    = "nl-nl"
    "English (United States)"  = "en-us"; "English (United Kingdom)" = "en-gb"
    "Estonian"                 = "et-ee"; "Finnish"                  = "fi-fi"
    "French"                   = "fr-fr"; "French (Canada)"          = "fr-ca"
    "German"                   = "de-de"; "Greek"                    = "el-gr"
    "Hebrew"                   = "he-il"; "Hungarian"                = "hu-hu"
    "Indonesian"               = "id-id"; "Italian"                  = "it-it"
    "Japanese"                 = "ja-jp"; "Korean"                   = "ko-kr"
    "Latvian"                  = "lv-lv"; "Lithuanian"               = "lt-lt"
    "Norwegian"                = "nb-no"; "Polish"                   = "pl-pl"
    "Portuguese (Brazil)"      = "pt-br"; "Portuguese (Portugal)"    = "pt-pt"
    "Romanian"                 = "ro-ro"; "Russian"                  = "ru-ru"
    "Serbian (Latin)"          = "sr-latn-rs"; "Slovak"              = "sk-sk"
    "Slovenian"                = "sl-si"; "Spanish"                  = "es-es"
    "Spanish (Mexico)"         = "es-mx"; "Swedish"                  = "sv-se"
    "Thai"                     = "th-th"; "Turkish"                  = "tr-tr"
    "Ukrainian"                = "uk-ua"
}

function Resolve-LangCode {
    param([string]$Lang)
    if ($Lang -match '^[a-zA-Z]{2,3}(-[a-zA-Z0-9]{2,8}){0,2}$') { return $Lang.ToLower() }
    foreach ($key in $LanguageMap.Keys) {
        if ($key -ieq $Lang) { return $LanguageMap[$key] }
    }
    Write-Error "Unrecognized language '$Lang'. Use full name like 'English (United States)' or code like 'en-us'."
    exit 1
}

function Search-UpdateCatalog {
    param([string]$WinVer, [string]$CatalogArch)
    $query     = "Windows $WinVer $CatalogArch Cumulative Update"
    $searchUrl = "https://www.catalog.update.microsoft.com/Search.aspx?q=$([System.Uri]::EscapeDataString($query))"
    Write-Step "Searching catalog: $query"
    try {
        $html = (Invoke-WebRequest -Uri $searchUrl -UseBasicParsing -ErrorAction Stop).Content
    } catch {
        Write-Warning "Update catalog search failed: $($_.Exception.Message)"
        return @()
    }
    $updates      = @()
    $rowPattern   = '(<tr id="([a-f0-9]{8}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{4}-[a-f0-9]{12})_R\d+"[^>]*>.*?</tr>)'
    $rows         = [regex]::Matches($html, $rowPattern, [System.Text.RegularExpressions.RegexOptions]::Singleline)
    foreach ($row in $rows) {
        $id  = $row.Groups[2].Value
        $htm = $row.Groups[1].Value
        $title = ""; $class = ""
        if ($htm -match 'class="contentTextItemSpacerNoBreakLink">([^<]+)</a>') { $title = $Matches[1].Trim() }
        if ($htm -match '<td[^>]*id="[^"]*_C3_R\d+"[^>]*>(.*?)</td>') { $class = ($Matches[1] -replace '<[^>]+>','').Trim() }
        if ($title) { $updates += [PSCustomObject]@{ UpdateId = $id; Title = $title; Classification = $class } }
    }
    return $updates
}

function Get-UpdateDownloadUrl {
    param([string]$UpdateId)
    $body = @{ updateIDs = "[{`"size`":0,`"languages`":`"`",`"uidInfo`":`"$UpdateId`",`"updateID`":`"$UpdateId`"}]" }
    try {
        $html = (Invoke-WebRequest -Uri "https://www.catalog.update.microsoft.com/DownloadDialog.aspx" `
            -Method Post -Body $body -UseBasicParsing -ContentType "application/x-www-form-urlencoded" -ErrorAction Stop).Content
        if ($html -match "downloadInformation\[\d+\]\.files\[\d+\]\.url\s*=\s*'([^']+)'") { return $Matches[1] }
        if ($html -match "'(https?://[^']+\.msu)'")       { return $Matches[1] }
        if ($html -match '(https?://[^\s"<>]+\.msu)')     { return $Matches[0] }
    } catch {
        Write-Warning "Failed to get update download URL: $($_.Exception.Message)"
    }
    return $null
}

# --- Setup --------------------------------------------------------------------

if (-not (Test-Path $OutputDir)) { New-Item -ItemType Directory -Path $OutputDir | Out-Null }
$OutputDir = (Resolve-Path $OutputDir).Path
$langCode  = Resolve-LangCode $Language

# catalog arch label differs from API arch for search query
$catalogArch = if ($Arch -eq "arm64") { "ARM64" } else { "x64" }

$workDir  = Join-Path $OutputDir "work_win${OsVer}_${Arch}_${langCode}"
$stageDir = Join-Path $workDir "staging"
$mountDir = Join-Path $workDir "mount"
foreach ($d in @($workDir, $stageDir, $mountDir)) {
    if (Test-Path $d) { Remove-Item $d -Recurse -Force }
    New-Item -ItemType Directory -Path $d | Out-Null
}

# -- Step 1: Fetch ESD from MCT catalog proxy ----------------------------------

Write-Status "Step 1: Querying MCT catalog  -  Windows $OsVer $Arch ($langCode)"

$apiUrl = "https://vesperos.chillcraft.me/windows/esd.php?os=$OsVer&arch=$Arch&edition=$Edition&lang=$langCode"
if ($Build) { $apiUrl += "&build=$Build" }
Write-Step "URL: $apiUrl"

try {
    $xmlRaw = (Invoke-WebRequest -Uri $apiUrl -UseBasicParsing -ErrorAction Stop).Content
} catch {
    Write-Error "MCT catalog fetch failed: $($_.Exception.Message)"
    exit 1
}

if (-not ($xmlRaw -match '<FilePath>')) {
    Write-Error "No ESD found for Windows $OsVer $Arch $langCode (edition=$Edition$(if ($Build){", build=$Build"}))."
    exit 1
}

# If multiple File nodes returned (multiple builds), take the last one
$allFileNodes = [regex]::Matches($xmlRaw, '(?s)<File>.*?</File>')
$fileBlock    = if ($allFileNodes.Count -gt 0) { $allFileNodes[$allFileNodes.Count - 1].Value } else { $xmlRaw }

$esdUrl   = [regex]::Match($fileBlock, '<FilePath>\s*(.*?)\s*</FilePath>').Groups[1].Value.Trim()
$sha1     = [regex]::Match($fileBlock, '<Sha1>\s*(.*?)\s*</Sha1>').Groups[1].Value.Trim()
$sizeStr  = [regex]::Match($fileBlock, '<Size>\s*(.*?)\s*</Size>').Groups[1].Value.Trim()
$fileName = [regex]::Match($fileBlock, '<FileName>\s*(.*?)\s*</FileName>').Groups[1].Value.Trim()

if (-not $esdUrl) { Write-Error "Could not parse ESD URL from catalog response." ; exit 1 }

Write-Step "File : $fileName"
if ($sizeStr) { Write-Step "Size : $([math]::Round([long]$sizeStr / 1MB, 1)) MB" }
if ($sha1)    { Write-Step "SHA1 : $sha1" }

$esdPath = Join-Path $workDir "$([System.IO.Path]::GetFileNameWithoutExtension($fileName)).esd"

if (Test-Path $esdPath) {
    Write-Step "Checking cached ESD..."
    if ($sha1 -and (Confirm-Hash -Path $esdPath -Expected $sha1 -Algo SHA1)) {
        Write-Step "Cached ESD valid  -  skipping download."
    } else {
        Write-Step "Cache invalid  -  re-downloading..."
        Remove-Item $esdPath -Force
        Invoke-WebRequest -Uri $esdUrl -OutFile $esdPath -UseBasicParsing -ErrorAction Stop
    }
} else {
    Write-Step "Downloading ESD..."
    Invoke-WebRequest -Uri $esdUrl -OutFile $esdPath -UseBasicParsing -ErrorAction Stop
}

if ($sha1) {
    Write-Step "Verifying ESD hash..."
    if (-not (Confirm-Hash -Path $esdPath -Expected $sha1 -Algo SHA1)) {
        Remove-Item $esdPath -Force
        Write-Error "ESD corrupt after download. Delete working directory and retry."
        exit 1
    }
    Write-Step "Hash OK."
} else {
    Write-Step "No SHA1 in catalog response  -  skipping hash check."
}

# -- Step 2: Extract ESD → ISO staging structure -------------------------------

Write-Status "Step 2: Building ISO staging structure from ESD"

$allIndexes = Get-WindowsImage -ImagePath $esdPath
foreach ($idx in $allIndexes) { Write-Step "  [$($idx.ImageIndex)] $($idx.ImageName)" }

if ($ImageIndex) {
    $installIdx = [int]$ImageIndex
} else {
    $installIdx = ($allIndexes | Where-Object { $_.ImageName -match 'Windows 1[01] Pro$' } |
                   Select-Object -First 1).ImageIndex
    if (-not $installIdx) {
        $installIdx = $allIndexes[-1].ImageIndex
        Write-Step "No Pro edition found  -  using last index: $installIdx"
    } else {
        Write-Step "Auto-selected Pro at index $installIdx"
    }
}

# Index 1: media folder skeleton (boot files, efi, sources stub, etc.)
Write-Step "Extracting media structure (index 1)..."
Expand-WindowsImage -ImagePath $esdPath -Index 1 -ApplyPath $stageDir -ErrorAction Stop

# Indexes 2+3: WinPE + Windows Setup → boot.wim
$bootWim = Join-Path $stageDir "sources\boot.wim"
if (Test-Path $bootWim) { Remove-Item $bootWim -Force }
Write-Step "Building boot.wim (index 2 = WinPE, index 3 = Setup)..."
& dism /English /Export-Image "/SourceImageFile:$esdPath" "/SourceIndex:2" `
    "/DestinationImageFile:$bootWim" /Compress:max /CheckIntegrity | Out-Null
& dism /English /Export-Image "/SourceImageFile:$esdPath" "/SourceIndex:3" `
    "/DestinationImageFile:$bootWim" /Compress:max /CheckIntegrity | Out-Null

# Target OS index → install.wim (WIM format needed for DISM mount in next step)
$installWim = Join-Path $stageDir "sources\install.wim"
Write-Step "Exporting install image (index $installIdx) to WIM..."
& dism /English /Export-Image "/SourceImageFile:$esdPath" "/SourceIndex:$installIdx" `
    "/DestinationImageFile:$installWim" /Compress:max /CheckIntegrity | Out-Null

Remove-Item $esdPath -Force

# -- Step 3: Download + apply latest cumulative update ------------------------

if ($SkipUpdate) {
    Write-Status "Step 3: Skipping cumulative update (-SkipUpdate)"
} else {
    Write-Status "Step 3: Downloading latest cumulative update for Windows $OsVer $catalogArch"

    $allUpdates = Search-UpdateCatalog -WinVer $OsVer -CatalogArch $catalogArch
    $cuList = $allUpdates | Where-Object {
        $_.Title -match "Cumulative Update" -and
        $_.Title -notmatch "Preview|Dynamic|\.NET Framework|Internet Explorer|Server 20" -and
        $_.Title -match "Windows $OsVer" -and
        ($_.Classification -match "Updates|Security Updates")
    } | Sort-Object {
        if ($_.Title -match 'KB(\d+)') { [int]$Matches[1] } else { 0 }
    } -Descending

    $latest = $cuList | Select-Object -First 1

    if (-not $latest) {
        Write-Warning "No cumulative update found  -  skipping update integration."
    } else {
        Write-Step "Found: $($latest.Title)"
        $kbNum = if ($latest.Title -match '(KB\d+)') { $Matches[1] } else { "KBUNKNOWN" }

        $cuUrl = Get-UpdateDownloadUrl -UpdateId $latest.UpdateId
        if (-not $cuUrl) {
            Write-Warning "Could not retrieve download URL for $kbNum  -  skipping."
        } else {
            $cuPath = Join-Path $workDir "${kbNum}.msu"
            if (-not (Test-Path $cuPath)) {
                Write-Step "Downloading $kbNum..."
                Invoke-WebRequest -Uri $cuUrl -OutFile $cuPath -UseBasicParsing -ErrorAction Stop
            } else {
                Write-Step "$kbNum already downloaded."
            }
            Write-Step "SHA256: $((Get-FileHash $cuPath -Algorithm SHA256).Hash)"

            Write-Step "Mounting install.wim..."
            & dism /Mount-Wim "/WimFile:$installWim" /Index:1 "/MountDir:$mountDir" | Out-Null
            if ($LASTEXITCODE -ne 0) {
                Write-Error "Failed to mount install.wim  -  aborting update step."
                exit 1
            }

            Write-Step "Applying $kbNum..."
            & dism "/Image:$mountDir" /Add-Package "/PackagePath:$cuPath" | Out-Null

            Write-Step "Running component cleanup..."
            & dism "/Image:$mountDir" /Cleanup-Image /StartComponentCleanup /ResetBase | Out-Null

            Write-Step "Committing changes..."
            & dism /Unmount-Wim "/MountDir:$mountDir" /Commit | Out-Null
            if ($LASTEXITCODE -ne 0) {
                Write-Warning "Commit failed  -  discarding mount."
                & dism /Unmount-Wim "/MountDir:$mountDir" /Discard | Out-Null
            } else {
                Write-Step "$kbNum applied successfully."
            }
        }
    }
}

# -- Step 4: Create bootable ISO -----------------------------------------------

Write-Status "Step 4: Creating bootable ISO"

$noprompt = Join-Path $stageDir "efi\microsoft\boot\efisys_noprompt.bin"
$efisys   = Join-Path $stageDir "efi\microsoft\boot\efisys.bin"
if (Test-Path $noprompt) {
    Copy-Item $noprompt $efisys -Force
    Write-Step "Using efisys_noprompt.bin (no keypress required at boot)."
}

$OSCDIMG  = Get-OscdimgPath
$isoLabel = "Win${OsVer}_${Arch}_${langCode}".ToUpper() -replace '[^A-Z0-9_]', '_'
$isoOut   = Join-Path $OutputDir "win${OsVer}_${Arch}_${langCode}_updated.iso"

$etfsboot = Join-Path $stageDir "boot\etfsboot.com"
$bootData = if (Test-Path $etfsboot) {
    "2#p0,e,b$etfsboot#pEF,e,b$efisys"
} else {
    Write-Warning "etfsboot.com not found  -  ISO will be EFI-only (no legacy BIOS boot)."
    "1#pEF,e,b$efisys"
}

& "$OSCDIMG" "-l$isoLabel" '-m' '-o' '-u2' '-udfver102' "-bootdata:$bootData" $stageDir $isoOut

if ($LASTEXITCODE -ne 0) {
    Write-Error "oscdimg failed with exit code $LASTEXITCODE."
    exit 1
}

if (-not $KeepWork) {
    Write-Step "Cleaning up working directory..."
    Remove-Item $workDir -Recurse -Force
}

Write-Status "Done."
Write-Host "ISO: $isoOut" -ForegroundColor Green
$isoSize = [math]::Round((Get-Item $isoOut).Length / 1GB, 2)
Write-Host "Size: $isoSize GB" -ForegroundColor Gray