#Requires -Version 3.0

<#
.SYNOPSIS
    Downloads the latest Windows Update for Windows 10 or 11
.DESCRIPTION
    Fetches and downloads the latest cumulative update from Microsoft Update Catalog
    for manual injection into Windows images
.PARAMETER WindowsVersion
    Windows version: "10" or "11"
.PARAMETER Architecture
    System architecture: "x64" or "ARM64"
.PARAMETER Edition
    Windows edition (optional): "Home", "Pro", "Enterprise", etc.
.PARAMETER OutputPath
    Download destination path
.PARAMETER Build
    Specific build number (optional, e.g., "22H2", "23H2", "24H2")
.EXAMPLE
    .\Get-LatestWindowsUpdate.ps1 -WindowsVersion 11 -Architecture x64 -OutputPath "C:\Updates"
#>

param(
    [Parameter(Mandatory=$true)]
    [ValidateSet("10", "11")]
    [string]$WindowsVersion,
    
    [Parameter(Mandatory=$true)]
    [ValidateSet("x64", "ARM64")]
    [string]$Architecture,
    
    [Parameter(Mandatory=$false)]
    [string]$Edition,
    
    [Parameter(Mandatory=$true)]
    [string]$OutputPath,
    
    [Parameter(Mandatory=$false)]
    [string]$Build
)

# Create output directory if it doesn't exist
if (-not (Test-Path $OutputPath)) {
    New-Item -ItemType Directory -Path $OutputPath -Force | Out-Null
}

# Microsoft Update Catalog URL
$catalogUrl = "https://www.catalog.update.microsoft.com"

# Build search query based on Windows version
if ($WindowsVersion -eq "11") {
    if ($Build) {
        $searchQuery = "Windows 11 $Build $Architecture Cumulative Update"
    } else {
        $searchQuery = "Windows 11 $Architecture Cumulative Update"
    }
} else {
    if ($Build) {
        $searchQuery = "Windows 10 $Build $Architecture Cumulative Update"
    } else {
        $searchQuery = "Windows 10 $Architecture Cumulative Update"
    }
}

Write-Host "Searching for: $searchQuery" -ForegroundColor Cyan

# Function to parse Microsoft Update Catalog
function Search-UpdateCatalog {
    param([string]$Query)
    
    $searchUrl = "$catalogUrl/Search.aspx?q=$([System.Uri]::EscapeDataString($Query))"
    
    try {
        $response = Invoke-WebRequest -Uri $searchUrl -UseBasicParsing
        $html = $response.Content
        
        # Parse the results table
        $updates = @()
        
        # The catalog structure: each row has an ID like "GUID_R0", "GUID_R1", etc.
        # We need to extract GUIDs from row IDs and then find the corresponding title/product/classification
        
        # Find all table rows with GUID pattern IDs
        $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>)'
        $rowMatches = [regex]::Matches($html, $rowPattern, [System.Text.RegularExpressions.RegexOptions]::Singleline)
        
        foreach ($rowMatch in $rowMatches) {
            $updateId = $rowMatch.Groups[2].Value
            $rowHtml = $rowMatch.Groups[1].Value
            
            # Extract title from the link
            $title = ""
            if ($rowHtml -match 'class="contentTextItemSpacerNoBreakLink">([^<]+)</a>') {
                $title = $Matches[1].Trim()
            }
            
            # Extract all TD contents by finding text between > and <
            # Pattern: extract content from TD cells more carefully
            $tdPattern = '<td[^>]*id="[^"]*_C(\d+)_R\d+"[^>]*>(.*?)</td>'
            $tdMatches = [regex]::Matches($rowHtml, $tdPattern, [System.Text.RegularExpressions.RegexOptions]::Singleline)
            
            $products = ""
            $classification = ""
            
            # C1 = title, C2 = products, C3 = classification
            foreach ($td in $tdMatches) {
                $colNum = $td.Groups[1].Value
                $content = $td.Groups[2].Value
                
                # Strip HTML tags and get text
                $text = $content -replace '<[^>]+>', '' -replace '\s+', ' '
                $text = $text.Trim()
                
                if ($colNum -eq "2") {
                    $products = $text
                }
                elseif ($colNum -eq "3") {
                    $classification = $text
                }
            }
            
            if ($title) {
                $updates += [PSCustomObject]@{
                    UpdateId = $updateId
                    Title = $title
                    Products = $products
                    Classification = $classification
                }
            }
        }
        
        return $updates
    }
    catch {
        Write-Error "Failed to search catalog: $_"
        return @()
    }
}

# Function to get download link for an update
function Get-UpdateDownloadLink {
    param([string]$UpdateId)
    
    $downloadUrl = "$catalogUrl/DownloadDialog.aspx"
    
    try {
        # Format the POST body as the catalog expects
        $body = @{
            updateIDs = "[{`"size`":0,`"languages`":`"`",`"uidInfo`":`"$UpdateId`",`"updateID`":`"$UpdateId`"}]"
        }
        
        $response = Invoke-WebRequest -Uri $downloadUrl -Method Post -Body $body -UseBasicParsing -ContentType "application/x-www-form-urlencoded"
        $html = $response.Content
        
        # Try multiple patterns to extract the download URL
        # Pattern 1: JavaScript variable assignment
        if ($html -match "downloadInformation\[\d+\]\.files\[\d+\]\.url\s*=\s*'([^']+)'") {
            return $Matches[1]
        }
        
        # Pattern 2: Direct URL in quotes
        if ($html -match "'(https?://[^']+\.msu)'") {
            return $Matches[1]
        }
        
        # Pattern 3: URL without quotes
        if ($html -match '(https?://[^\s"<>]+\.msu)') {
            return $Matches[0]
        }
        
        # If we still haven't found it, try to extract any .msu URL
        if ($html -match '(https?://[^\s"<>]+catalog[^\s"<>]+\.msu)') {
            return $Matches[0]
        }
        
        Write-Host "Debug: Could not find download URL in response" -ForegroundColor Red
        return $null
    }
    catch {
        Write-Error "Failed to get download link: $_"
        return $null
    }
}

# Search for updates
Write-Host "Searching Microsoft Update Catalog..." -ForegroundColor Yellow
$updates = Search-UpdateCatalog -Query $searchQuery

if ($updates.Count -eq 0) {
    Write-Error "No updates found matching criteria"
    exit 1
}

Write-Host "Found $($updates.Count) total updates" -ForegroundColor Cyan
Write-Host "`nAll updates found:" -ForegroundColor Yellow
foreach ($update in $updates) {
    Write-Host "  - $($update.Title)" -ForegroundColor Gray
    Write-Host "    Classification: $($update.Classification)" -ForegroundColor DarkGray
}

# Filter for cumulative updates only and exclude preview builds
$cumulativeUpdates = $updates | Where-Object {
    $_.Title -match "Cumulative Update" -and
    $_.Title -notmatch "Preview" -and
    $_.Title -notmatch "Dynamic" -and
    $_.Title -notmatch "\.NET Framework" -and
    $_.Title -notmatch "Internet Explorer" -and
    $_.Title -notmatch "Server 2012" -and
    $_.Title -notmatch "Server 2016" -and
    $_.Title -notmatch "Server 2019" -and
    $_.Title -notmatch "Server 2022" -and
    ($_.Title -match "Windows $WindowsVersion" -or $_.Title -match "Windows $WindowsVersion,") -and
    ($_.Classification -eq "Updates" -or $_.Classification -eq "Security Updates")
}

Write-Host "`nFiltered to $($cumulativeUpdates.Count) cumulative updates" -ForegroundColor Cyan

if ($Edition) {
    $cumulativeUpdates = $cumulativeUpdates | Where-Object {
        $_.Products -match $Edition
    }
}

# Sort by KB number (higher KB = newer) and get the latest
$latestUpdate = $cumulativeUpdates | Sort-Object {
    if ($_.Title -match 'KB(\d+)') {
        [int]$Matches[1]
    } else {
        0
    }
} -Descending | Select-Object -First 1

if (-not $latestUpdate) {
    Write-Error "No suitable cumulative update found"
    exit 1
}

Write-Host "`nLatest Update Found:" -ForegroundColor Green
Write-Host "  Title: $($latestUpdate.Title)" -ForegroundColor White
Write-Host "  Products: $($latestUpdate.Products)" -ForegroundColor White
Write-Host "  Update ID: $($latestUpdate.UpdateId)" -ForegroundColor White

# Extract KB number for filename
$kbNumber = ""
if ($latestUpdate.Title -match '(KB\d+)') {
    $kbNumber = $Matches[1]
}

# Get download link
Write-Host "`nRetrieving download link..." -ForegroundColor Yellow
$downloadLink = Get-UpdateDownloadLink -UpdateId $latestUpdate.UpdateId

if (-not $downloadLink) {
    Write-Error "Failed to retrieve download link"
    exit 1
}

Write-Host "Download URL: $downloadLink" -ForegroundColor Cyan

# Generate filename
$fileName = "Windows$($WindowsVersion)_${Architecture}_${kbNumber}_CumulativeUpdate.msu"
$outputFile = Join-Path $OutputPath $fileName

# Download the update
Write-Host "`nDownloading update to: $outputFile" -ForegroundColor Yellow

try {
    $ProgressPreference = 'SilentlyContinue'
    
    # Use BITS transfer for large files with resume capability
    Import-Module BitsTransfer -ErrorAction SilentlyContinue
    
    if (Get-Command Start-BitsTransfer -ErrorAction SilentlyContinue) {
        Write-Host "Using BITS transfer..." -ForegroundColor Cyan
        Start-BitsTransfer -Source $downloadLink -Destination $outputFile -DisplayName "Windows Update Download" -Description $latestUpdate.Title
    }
    else {
        Write-Host "Using web request..." -ForegroundColor Cyan
        Invoke-WebRequest -Uri $downloadLink -OutFile $outputFile -UseBasicParsing
    }
    
    Write-Host "`nDownload completed successfully!" -ForegroundColor Green
    Write-Host "File location: $outputFile" -ForegroundColor White
    
    # Display file info
    $fileInfo = Get-Item $outputFile
    Write-Host "File size: $([math]::Round($fileInfo.Length / 1MB, 2)) MB" -ForegroundColor White
    Write-Host "SHA256: $((Get-FileHash -Path $outputFile -Algorithm SHA256).Hash)" -ForegroundColor White
    
    # Display injection instructions
    Write-Host "`n=== Injection Instructions ===" -ForegroundColor Cyan
    Write-Host "To inject this update into a WIM file, use DISM:" -ForegroundColor Yellow
    Write-Host "  1. Mount the WIM: " -NoNewline -ForegroundColor White
    Write-Host "dism /Mount-Wim /WimFile:install.wim /Index:1 /MountDir:C:\Mount" -ForegroundColor Gray
    Write-Host "  2. Add update: " -NoNewline -ForegroundColor White
    Write-Host "dism /Image:C:\Mount /Add-Package /PackagePath:`"$outputFile`"" -ForegroundColor Gray
    Write-Host "  3. Commit changes: " -NoNewline -ForegroundColor White
    Write-Host "dism /Unmount-Wim /MountDir:C:\Mount /Commit" -ForegroundColor Gray
    
}
catch {
    Write-Error "Download failed: $_"
    exit 1
}