<#
.SYNOPSIS
    Downloads Windows 10/11 ISOs from the MCT ESD catalog proxy.
.DESCRIPTION
    Supported modes:
      Win11ARM64Latest   -  Latest stable Win11 ARM64 ISO.
      Win11x64Latest     -  Latest stable Win11 x64 ISO.
      Win10ARM64Latest   -  Latest stable Win10 ARM64 ISO.
      Win10x64Latest     -  Latest stable Win10 x64 ISO.
      Win11ARM64RPI      -  Win11 ARM64 build 22631 (last build with broad ARMv8.0 support, for Raspberry Pi).

    All modes use the MCT catalog proxy at vesperos.chillcraft.me.
    Pass -Build to pin to a specific build substring (e.g. "26100", "22631").

.PARAMETER Mode
    Which image to download (see above).
.PARAMETER OutputDir
    Directory to save the output ISO (default: current directory).
.PARAMETER Language
    Full language name e.g. "English (United States)", or short code e.g. "en-us".
.PARAMETER Build
    Optional build substring filter (matched against ESD filename). E.g. "26100" or "22631".
.PARAMETER Edition
    Windows edition to extract (default: Professional).
.PARAMETER Edition
    Windows edition to extract (default: Professional).
.PARAMETER ImageIndex
    Override which ESD index to use for install.esd. Pro is auto-selected if omitted.
.PARAMETER EsdOnly
    Download and verify the ESD only  -  skip ISO assembly. Output is the raw .esd file.
.EXAMPLE
    .\utils\Get-WindowsISO.ps1 -Mode Win11x64Latest -OutputDir C:\ISOs
.EXAMPLE
    .\utils\Get-WindowsISO.ps1 -Mode Win11ARM64RPI -OutputDir C:\ISOs -Language "en-us"
.EXAMPLE
    .\utils\Get-WindowsISO.ps1 -Mode Win10x64Latest -OutputDir C:\ISOs -Build "19045"
.EXAMPLE
    .\utils\Get-WindowsISO.ps1 -Mode Win11x64Latest -OutputDir C:\ISOs -Language "French"
#>

param(
    [Parameter(Mandatory)]
    [ValidateSet("Win11ARM64Latest", "Win11x64Latest", "Win10ARM64Latest", "Win10x64Latest", "Win11ARM64RPI")]
    [string]$Mode,

    [string]$OutputDir = ".",
    [string]$Language  = "English (United States)",
    [string]$Build,
    [string]$Edition   = "Professional",
    [string]$ImageIndex,
    [switch]$EsdOnly    # Download and verify the ESD only  -  skip ISO assembly
)

Set-StrictMode -Off
$ErrorActionPreference = "Stop"

# Win11ARM64RPI pins to build 22631 (last Win11 release with broad ARMv8.0 support)
if ($Mode -eq "Win11ARM64RPI" -and -not $Build) {
    $Build = "22631"
}

# Expand-WindowsImage and dism /Export-Image require elevation; -EsdOnly (download only) does not.
if (-not $EsdOnly) {
    $principal = New-Object Security.Principal.WindowsPrincipal([Security.Principal.WindowsIdentity]::GetCurrent())
    if (-not $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) {
        Write-Host "Restarting as administrator (required for DISM/image operations)..."
        $argStr = "-NoProfile -ExecutionPolicy Bypass -File `"$PSCommandPath`" -Mode `"$Mode`" -OutputDir `"$OutputDir`" -Language `"$Language`""
        if ($Build)      { $argStr += " -Build `"$Build`"" }
        if ($Edition)    { $argStr += " -Edition `"$Edition`"" }
        if ($ImageIndex) { $argStr += " -ImageIndex `"$ImageIndex`"" }
        Start-Process PowerShell -ArgumentList $argStr -Verb RunAs
        exit
    }
}

# -----------------------------------------------------------------
#  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 a full name like 'English (United States)' or short code like 'en-us'."
    Write-Host "Available languages:"
    $LanguageMap.GetEnumerator() | ForEach-Object { Write-Host "  $($_.Value)  ($($_.Key))" }
    exit 1
}

# -----------------------------------------------------------------
#  MCT ESD catalog proxy  (Win10/Win11, latest or pinned build)
#  Proxy API: https://vesperos.chillcraft.me/windows/esd.php
#    ?os=10|11  &arch=x64|arm64|x86  &edition=Professional  &lang=en-us
#    &build=NNNNN   (optional substring match on filename)
# -----------------------------------------------------------------

function Get-WindowsFromMCT {
    param(
        [string]$OutDir,
        [string]$Lang,
        [string]$Arch,      # "ARM64" or "x64"
        [string]$OsVer,     # "10" or "11"
        [string]$ImgIndex,
        [string]$EditionFilter = "Professional",
        [string]$BuildFilter,
        [switch]$EsdOnly
    )

    $langCode = Resolve-LangCode $Lang
    $apiArch  = $Arch.ToLower()   # proxy uses x64/arm64

    Write-Status "Downloading Windows $OsVer $Arch latest ($langCode) via MCT catalog"

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

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

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

    # If multiple File nodes returned, take the last one (newest build tends to be last)
    $allMatches = [regex]::Matches($raw, '(?s)<File>.*?</File>')
    $fileBlock  = if ($allMatches.Count -gt 0) { $allMatches[$allMatches.Count - 1].Value } else { $raw }

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

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

    if (Test-Path $esdPath) {
        Write-Step "Checking existing ESD..."
        if ($sha1 -and (Confirm-Hash -Path $esdPath -Expected $sha1 -Algo SHA1)) {
            Write-Step "Existing ESD valid  -  skipping download."
        } else {
            Write-Step "ESD invalid or unverified  -  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 and retry."
            exit 1
        }
        Write-Step "Hash OK."
    } else {
        Write-Step "No SHA1 in catalog response  -  skipping hash verification."
    }

    if ($EsdOnly) {
        Write-Status "ESD ready: $esdPath"
        return $esdPath
    }

    # Reuse the same ESD→ISO build logic as Get-WindowsFromWOR
    # MCT ESDs have the same structure: idx1=media skeleton, idx2=WinPE, idx3=Setup, idx4+=OS

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

    if ($ImgIndex) {
        $installIdx = [int]$ImgIndex
    } 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"
        }
    }

    $stageDir = Join-Path $OutDir "isoextract_${buildTag}"
    if (Test-Path $stageDir) { Remove-Item $stageDir -Recurse -Force }
    New-Item -ItemType Directory -Path $stageDir | Out-Null

    Write-Status "Extracting setup media structure (index 1)..."
    Expand-WindowsImage -ImagePath $esdPath -Index 1 -ApplyPath $stageDir -ErrorAction Stop
    Write-Step "Setup structure extracted."

    $bootWim = Join-Path $stageDir "sources\boot.wim"
    if (Test-Path $bootWim) { Remove-Item $bootWim -Force }

    Write-Status "Building boot.wim (PE index 2 + Setup index 3)..."
    & 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
    Write-Step "boot.wim built."

    $installEsd = Join-Path $stageDir "sources\install.esd"
    Write-Status "Extracting install image (index $installIdx)..."
    & dism /English /Export-Image "/SourceImageFile:$esdPath" "/SourceIndex:$installIdx" `
        "/DestinationImageFile:$installEsd" /Compress:recovery /CheckIntegrity | Out-Null
    Write-Step "install.esd created."

    $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 "Replaced efisys.bin with noprompt variant."
    }

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

    $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 BIOS boot)."
        "1#pEF,e,b$efisys"
    }

    Write-Status "Creating ISO: $isoOut"
    & "$OSCDIMG" "-l$label" '-m' '-o' '-u2' '-udfver102' "-bootdata:$bootData" $stageDir $isoOut

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

    Write-Step "Cleaning up..."
    Remove-Item $stageDir -Recurse -Force
    Remove-Item $esdPath  -Force

    Write-Status "Done. ISO saved to: $isoOut"
    return $isoOut
}

# -----------------------------------------------------------------
#  Entry point
# -----------------------------------------------------------------

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

$mctArgs = @{
    OutDir        = $OutputDir
    Lang          = $Language
    ImgIndex      = $ImageIndex
    EditionFilter = $Edition
    BuildFilter   = $Build
    EsdOnly       = $EsdOnly
}

switch ($Mode) {
    "Win11ARM64Latest" { Get-WindowsFromMCT @mctArgs -Arch "ARM64" -OsVer "11" }
    "Win11x64Latest"   { Get-WindowsFromMCT @mctArgs -Arch "x64"   -OsVer "11" }
    "Win10ARM64Latest" { Get-WindowsFromMCT @mctArgs -Arch "ARM64" -OsVer "10" }
    "Win10x64Latest"   { Get-WindowsFromMCT @mctArgs -Arch "x64"   -OsVer "10" }
    "Win11ARM64RPI"    { Get-WindowsFromMCT @mctArgs -Arch "ARM64" -OsVer "11" }
}