
#Requires -Modules Hyper-V, Storage
#Requires -RunAsAdministrator

<#
.SYNOPSIS
A PowerShell script to create a Windows 10/11 FFU file. 

.DESCRIPTION
This script creates a Windows 10/11 FFU and USB drive to help quickly get a Windows device reimaged. FFU can be customized with drivers, apps, and additional settings. 

.PARAMETER AdditionalFFUFiles
Array of full file paths to existing FFU files that should also be copied to the deployment USB when -CopyAdditionalFFUFiles is set to $true.

.PARAMETER AllowExternalHardDiskMedia
When set to $true, will allow the use of media identified as External Hard Disk media via WMI class Win32_DiskDrive. Default is not defined.

.PARAMETER AllowVHDXCaching
When set to $true, will cache the VHDX file to the $FFUDevelopmentPath\VHDXCache folder and create a config json file that will keep track of the Windows build information, the updates installed, and the logical sector byte size information. Default is $false.

.PARAMETER AppListPath
Path to a JSON file containing a list of applications to install using WinGet. Default is $FFUDevelopmentPath\Apps\AppList.json.

.PARAMETER AppsScriptVariables
When passed a hashtable, the script will create an AppsScriptVariables.json file in the OrchestrationPath. This file will be used to pass variables to the Apps script. The hashtable should contain key-value pairs where the key is the variable name and the value is the variable value.

.PARAMETER BitsPriority
BITS transfer priority used for downloads. Accepted values are 'Foreground', 'High', 'Normal', and 'Low'. Default is 'Normal'.

.PARAMETER BuildUSBDrive
When set to $true, will partition and format a USB drive and copy the captured FFU to the drive. 

.PARAMETER Cleanup
Switch to run cleanup-only mode. When specified, the script performs cleanup and exits without starting a new build.

.PARAMETER CleanupAppsISO
When set to $true, will remove the Apps ISO after the FFU has been captured. Default is $true.

.PARAMETER CleanupCurrentRunDownloads
When set to $true, cleanup mode will remove downloads created during the current run and restore backed up run JSON files. Default is $false.

.PARAMETER CleanupDeployISO
When set to $true, will remove the WinPE deployment ISO after the FFU has been captured. Default is $true.

.PARAMETER CleanupDrivers
When set to $true, will remove the drivers folder after the FFU has been captured. Default is $true.

.PARAMETER CompactOS
When set to $true, will compact the OS when building the FFU. Default is $true.

.PARAMETER CompressDownloadedDriversToWim
When set to $true, compresses downloaded drivers into a WIM file. Default is $false.

.PARAMETER ConfigFile
Path to a JSON file containing parameters to use for the script. Default is $null.

.PARAMETER CopyAdditionalFFUFiles
When set to $true, enables copying additional FFU files from $FFUDevelopmentPath\FFU to the deployment USB alongside the current build output.

.PARAMETER CopyAutopilot
When set to $true, will copy the $FFUDevelopmentPath\Autopilot folder to the Deployment partition of the USB drive. Default is $false.

.PARAMETER CopyDrivers
When set to $true, will copy the drivers from the $FFUDevelopmentPath\Drivers folder to the Drivers folder on the deploy partition of the USB drive. Default is $false.

.PARAMETER CopyPEDrivers
When set to $true, enables adding WinPE drivers. By default copies drivers from $FFUDevelopmentPath\PEDrivers to the WinPE deployment media unless -UseDriversAsPEDrivers is also $true.

.PARAMETER CopyPPKG
When set to $true, will copy the provisioning package from the $FFUDevelopmentPath\PPKG folder to the Deployment partition of the USB drive. Default is $false.

.PARAMETER CopyUnattend
When set to $true, stages the selected architecture-specific unattend XML file as Unattend.xml on the deployment partition of the USB drive. Default is $false.

.PARAMETER DeviceNamingMode
Controls how device naming is handled when unattend content is copied to USB media or injected into the FFU. Supported values are Legacy, None, Prompt, Template, Prefixes, and SerialComputerNames.

.PARAMETER DeviceNameTemplate
Sets the device name used when DeviceNamingMode is Template. Supports a static name or the %serial% token when CopyUnattend is used.

.PARAMETER DeviceNamePrefixes
Sets the prefixes used when DeviceNamingMode is Prefixes. Each entry becomes a line in prefixes.txt on the deployment media.

.PARAMETER DeviceNamePrefixesPath
Path to the source prefixes file used for legacy copy or when DeviceNamePrefixes is not supplied. Default is $FFUDevelopmentPath\Unattend\prefixes.txt.

.PARAMETER DeviceNameSerialComputerNames
Sets the CSV content used when DeviceNamingMode is SerialComputerNames. The CSV must include SerialNumber and ComputerName headers.

.PARAMETER DeviceNameSerialComputerNamesPath
Path to the source CSV file used when DeviceNamingMode is SerialComputerNames and DeviceNameSerialComputerNames is not supplied. Default is $FFUDevelopmentPath\Unattend\SerialComputerNames.csv.

.PARAMETER UnattendX64FilePath
Path to the x64 unattend XML source file. Default is $FFUDevelopmentPath\Unattend\unattend_x64.xml.

.PARAMETER UnattendArm64FilePath
Path to the arm64 unattend XML source file. Default is $FFUDevelopmentPath\Unattend\unattend_arm64.xml.

.PARAMETER CreateDeploymentMedia
When set to $true, this will create WinPE deployment media for use when deploying to a physical device.

.PARAMETER CustomFFUNameTemplate
Sets a custom FFU output name with placeholders. Allowed placeholders are: {WindowsRelease}, {WindowsVersion}, {SKU}, {BuildDate}, {yyyy}, {MM}, {dd}, {H}, {hh}, {mm}, {tt}.

.PARAMETER Disksize
Size of the virtual hard disk for the virtual machine. Default is a 50GB dynamic disk.

.PARAMETER OSPartitionSize
Fixed size of the Windows partition in bytes. Leave as 0 to let Windows use the space remaining after Recovery and fixed-size data partitions are reserved.

.PARAMETER RecoveryPartitionSize
Optional fixed size of the Recovery partition in bytes. Leave as 0 to calculate the Recovery partition size from winre.wim plus buffer space.

.PARAMETER CreateRecoveryPartition
When set to $false, skips creating the Windows Recovery partition. Default is $true.

.PARAMETER AdditionalDataPartitions
Optional data partitions to create after the Recovery partition. Each item supports Name, Label, DriveLetter, SizeBytes or SizeGB, FillRemaining, FileSystem, and PersistDriveLetter. The configured letter is always used in the build VM. On physical devices, persisted partitions keep that letter and other data partitions receive the next available letter from D through Z.

.PARAMETER DriversFolder
Path to the drivers folder. Default is $FFUDevelopmentPath\Drivers.

.PARAMETER DriversJsonPath
Path to a JSON file that specifies which drivers to download.

.PARAMETER ExportConfigFile
Path to a JSON file to export the parameters used for the script.

.PARAMETER EnableVMNetworking
When set to $true, connects the build VM to the Hyper-V virtual switch named in -VMSwitchName during provisioning. Default is $false because internet-connected Sysprep is experimental.

.PARAMETER FFUCaptureLocation
Path to the folder where the captured FFU will be stored. Default is $FFUDevelopmentPath\FFU.

.PARAMETER FFUDevelopmentPath
Path to the FFU development folder. Default is $PSScriptRoot.

.PARAMETER FFUPrefix
Prefix for the generated FFU file. Default is _FFU.

.PARAMETER Headers
Headers to use when downloading files. Not recommended to modify.

.PARAMETER InjectUnattend
When set to $true and InstallApps is also $true, stages the selected architecture-specific unattend XML file to $FFUDevelopmentPath\Apps\Unattend\Unattend.xml so sysprep can use it inside the VM. Default is $false.

.PARAMETER InstallApps
When set to $true, the script will create an Apps.iso file from the $FFUDevelopmentPath\Apps folder. It will also create a VM, mount the Apps.iso, install the apps, sysprep, and capture the VM. When set to $false, the FFU is created from a VHDX file, and no VM is created.

.PARAMETER InstallDrivers
Install device drivers from the specified $FFUDevelopmentPath\Drivers folder if set to $true. Download the drivers and put them in the Drivers folder. The script will recurse the drivers folder and add the drivers to the FFU.

.PARAMETER InstallOffice
Install Microsoft Office if set to $true. The script will download the latest ODT and Office files in the $FFUDevelopmentPath\Apps\Office folder and install Office in the FFU via VM.

.PARAMETER ISOPath
Path to the Windows 10/11 ISO file.

.PARAMETER LogicalSectorSizeBytes
UInt32 value of 512 or 4096. Useful for 4Kn drives or devices shipping with UFS drives. Default is 512.

.PARAMETER SystemPartitionDriveLetter
Drive letter used for the System partition while building the FFU VHDX. Default is S.

.PARAMETER WindowsPartitionDriveLetter
Drive letter used for the Windows partition while building the FFU VHDX. Default is W.

.PARAMETER RecoveryPartitionDriveLetter
Drive letter used for the Recovery partition while building the FFU VHDX. Default is R.

.PARAMETER Make
Make of the device to download drivers. Accepted values are: 'Microsoft', 'Dell', 'HP', 'Lenovo'.

.PARAMETER MaxUSBDrives
Maximum number of USB drives to build in parallel. Default is 5. Set to 0 to process all discovered drives (or all selected drives when USBDriveList or selection is used). Actual throttle will never exceed the number of drives discovered.

.PARAMETER MediaType
String value of either 'business' or 'consumer'. This is used to identify which media type to download. Default is 'consumer'.

.PARAMETER Memory
Amount of memory to allocate for the virtual machine. Recommended to use 8GB if possible, especially for Windows 11. Default is 4GB.

.PARAMETER Model
Model of the device to download drivers. This is required if Make is set.

.PARAMETER OfficeConfigXMLFile
Path to a custom Office configuration XML file to use for installation.

.PARAMETER Optimize
When set to $true, will optimize the FFU file. Default is $true.

.PARAMETER OptimizeFFUPartitionNumber
Optional partition number to pass to DISM /Optimize-FFU /PartitionNumber. Leave as 0 to optimize the FillRemaining partition, or Windows when no partition uses FillRemaining.

.PARAMETER OptionalFeatures
Provide a semicolon-separated list of Windows optional features you want to include in the FFU (e.g., netfx3;TFTP).

.PARAMETER OrchestrationPath
Path to the orchestration folder containing scripts that run inside the VM. Default is $FFUDevelopmentPath\Apps\Orchestration.

.PARAMETER PEDriversFolder
Path to the folder containing drivers to be injected into the WinPE deployment media. Default is $FFUDevelopmentPath\PEDrivers.

.PARAMETER Processors
Number of virtual processors for the virtual machine. Recommended to use at least 4.

.PARAMETER ProductKey
Product key for the Windows edition specified in WindowsSKU. This will overwrite whatever SKU is entered for WindowsSKU. Recommended to use if you want to use a MAK or KMS key to activate Enterprise or Education. If using VL media instead of consumer media, you'll want to enter a MAK or KMS key here.

.PARAMETER PromptExternalHardDiskMedia
When set to $true, will prompt the user to confirm the use of media identified as External Hard Disk media via WMI class Win32_DiskDrive. Default is $true.

.PARAMETER RemoveApps
When set to $true, will remove the application content in the Apps folder after the FFU has been captured. Default is $true.

.PARAMETER RemoveFFU
When set to $true, will remove the FFU file from the $FFUDevelopmentPath\FFU folder after it has been copied to the USB drive. Default is $false.

.PARAMETER RemoveUpdates
When set to $true, will remove the downloaded CU, MSRT, Defender, Edge, OneDrive, and .NET files downloaded. Default is $true.

.PARAMETER RemoveDownloadedESD
When set to $true, will remove downloaded Windows ESD files after they have been applied. Default is $true.

.PARAMETER Threads
Controls the throttle applied to parallel tasks inside the script. Default is 5, matching the UI Threads field, and applies to driver downloads invoked through Invoke-ParallelProcessing.

.PARAMETER UpdateADK
When set to $true, the script will check for and install the latest Windows ADK and WinPE add-on if they are not already installed or up-to-date. Default is $true.

.PARAMETER UpdateEdge
When set to $true, will download and install the latest Microsoft Edge. Default is $false.

.PARAMETER UpdateLatestCU
When set to $true, will download and install the latest cumulative update. Default is $false.

.PARAMETER UpdateLatestDefender
When set to $true, will download and install the latest Windows Defender definitions and Defender platform update. Default is $false.

.PARAMETER UpdateLatestMicrocode
When set to $true, will download and install the latest microcode updates for applicable Windows releases (e.g., Windows Server 2016/2019, Windows 10 LTSC 2016/2019) into the FFU. Default is $false.

.PARAMETER UpdateLatestMSRT
When set to $true, will download and install the latest Windows Malicious Software Removal Tool. Default is $false.

.PARAMETER UpdateLatestNet
When set to $true, will download and install the latest .NET Framework. Default is $false.

.PARAMETER UpdateOneDrive
When set to $true, will download and install the latest OneDrive and install it as a per-machine installation instead of per-user. Default is $false.

.PARAMETER UpdatePreviewCU
When set to $true, will download and install the latest Preview cumulative update. Default is $false.

.PARAMETER USBDriveList
A hashtable containing USB drives from win32_diskdrive where:
- Key: USB drive model name (partial match supported)
- Value: USB drive UniqueId string, or an array of UniqueIds (to support selecting multiple drives with the same model)

Examples:
@{ "SanDisk Ultra" = "1234567890" }
@{ "SanDisk Ultra" = @("1234567890", "ABCDEFG"); "Kingston DataTraveler" = "0987654321" }

.PARAMETER UseDriversAsPEDrivers
When set to $true (and -CopyPEDrivers is also $true), bypasses the contents of $FFUDevelopmentPath\PEDrivers and instead builds the WinPE driver set dynamically from the $DriversFolder path, copying only the required WinPE drivers. Has no effect if -CopyPEDrivers is not specified. Default is $false.

.PARAMETER UserAgent
User agent string to use when downloading files.

.PARAMETER UserAppListPath
Path to a JSON file containing a list of user-defined applications to install. Default is $FFUDevelopmentPath\Apps\UserAppList.json.

.PARAMETER VMLocation
Default is $FFUDevelopmentPath\VM. This is the location of the VHDX that gets created where Windows will be installed to.

.PARAMETER VMSwitchName
Name of the Hyper-V virtual switch used when -EnableVMNetworking is set to $true. Provide it only if the VM needs network connectivity during provisioning.

.PARAMETER WindowsArch
String value of 'x86', 'x64', or 'arm64'. This is used to identify which architecture of Windows to download. Default is 'x64'.

.PARAMETER WindowsLang
String value in language-region format (e.g., 'en-us'). This is used to identify which language of media to download. Default is 'en-us'.

.PARAMETER WindowsRelease
Integer value of 10, 11, 2016, 2019, 2021, 2022, 2024, or 2025. This is used to identify which Windows client/LTSC/server release to use. Default is 11.

.PARAMETER WindowsSKU
Edition/SKU to install. Accepted values are: 'Home', 'Home N', 'Home Single Language', 'Education', 'Education N', 'Pro', 'Pro N', 'Pro Education', 'Pro Education N', 'Pro for Workstations', 'Pro N for Workstations', 'Enterprise', 'Enterprise N', 'Enterprise 2016 LTSB', 'Enterprise N 2016 LTSB', 'Enterprise LTSC', 'Enterprise N LTSC', 'IoT Enterprise LTSC', 'IoT Enterprise N LTSC', 'Standard', 'Standard (Desktop Experience)', 'Datacenter', 'Datacenter (Desktop Experience)'.

.PARAMETER WindowsVersion
String value of the Windows version to download. This is used to identify which version of Windows to download. Default is '25h2'.

.EXAMPLE
Command line for most people who want to download the latest Windows 11 Pro x64 media in English (US) with the latest Windows Cumulative Update, .NET Framework, Defender platform and definition updates, Edge, OneDrive, and Office/M365 Apps. It will also copy drivers to the FFU. This can take about 40 minutes to create the FFU due to the time it takes to download and install the updates.
.\BuildFFUVM.ps1 -WindowsSKU 'Pro' -Installapps $true -InstallOffice $true -InstallDrivers $true -CreateDeploymentMedia $true -BuildUSBDrive $true -UpdateLatestCU $true -UpdateLatestNet $true -UpdateLatestDefender $true -UpdateEdge $true -UpdateOneDrive $true -verbose

Command line for most people who want to create an FFU with Office and drivers and have downloaded their own ISO. This assumes you have copied this script and associated files to the C:\FFUDevelopment folder. If you need to use another drive or folder, change the -FFUDevelopment parameter (e.g. -FFUDevelopment 'D:\FFUDevelopment')
.\BuildFFUVM.ps1 -ISOPath 'C:\path_to_iso\Windows.iso' -WindowsSKU 'Pro' -Installapps $true -InstallOffice $true -InstallDrivers $true -CreateDeploymentMedia $true -BuildUSBDrive $true -verbose

Command line for those who just want a FFU with no drivers, apps, or Office and have downloaded their own ISO.
.\BuildFFUVM.ps1 -ISOPath 'C:\path_to_iso\Windows.iso' -WindowsSKU 'Pro' -Installapps $false -InstallOffice $false -InstallDrivers $false -CreateDeploymentMedia $true -BuildUSBDrive $true -verbose

Command line for those who just want a FFU with Apps and drivers, no Office and have downloaded their own ISO.
.\BuildFFUVM.ps1 -ISOPath 'C:\path_to_iso\Windows.iso' -WindowsSKU 'Pro' -Installapps $true -InstallOffice $false -InstallDrivers $true -CreateDeploymentMedia $true -BuildUSBDrive $true -verbose

Command line for those who want to download the latest Windows 11 Pro x64 media in English (US) and install the latest version of Office and drivers.
.\BuildFFUVM.ps1 -WindowsSKU 'Pro' -Installapps $true -InstallOffice $true -InstallDrivers $true -CreateDeploymentMedia $true -BuildUSBDrive $true -verbose

Command line for those who want to download the latest Windows 11 Pro x64 media in French (CA) and install the latest version of Office and drivers.
.\BuildFFUVM.ps1 -WindowsSKU 'Pro' -Installapps $true -InstallOffice $true -InstallDrivers $true -CreateDeploymentMedia $true -BuildUSBDrive $true -WindowsRelease 11 -WindowsArch 'x64' -WindowsLang 'fr-ca' -MediaType 'consumer' -verbose

Command line for those who want to download the latest Windows 11 Pro x64 media in English (US) and install the latest version of Office and drivers.
.\BuildFFUVM.ps1 -WindowsSKU 'Pro' -Installapps $true -InstallOffice $true -InstallDrivers $true -CreateDeploymentMedia $true -BuildUSBDrive $true -verbose

.NOTES
    Additional notes about your script.

.LINK
    https://github.com/rbalsleyMSFT/FFU
#>


[CmdletBinding()]
param(
    [Parameter(Mandatory = $false, Position = 0)]
    [ValidateScript({ Test-Path $_ })]
    [string]$ISOPath,
    [ValidateSet(
        'Home',
        'Home N',
        'Home Single Language',
        'Education',
        'Education N',
        'Pro',
        'Pro N',
        'Pro Education',
        'Pro Education N',
        'Pro for Workstations',
        'Pro N for Workstations',
        'Enterprise',
        'Enterprise N',
        'Enterprise 2016 LTSB',
        'Enterprise N 2016 LTSB',
        'Enterprise LTSC',
        'Enterprise N LTSC',
        'IoT Enterprise LTSC',
        'IoT Enterprise N LTSC',
        'Standard',
        'Standard (Desktop Experience)',
        'Datacenter',
        'Datacenter (Desktop Experience)'
    )]
    [string]$WindowsSKU = 'Pro',
    [ValidateScript({ Test-Path $_ })]
    [string]$FFUDevelopmentPath = $PSScriptRoot,
    [bool]$InstallApps,
    [string]$AppListPath,
    [string]$UserAppListPath,

    [hashtable]$AppsScriptVariables,
    [bool]$InstallOffice,
    [string]$OfficeConfigXMLFile,
    [ValidateSet('Microsoft', 'Dell', 'HP', 'Lenovo')]
    [string]$Make,
    [string]$Model,
    [bool]$InstallDrivers,
    [uint64]$Memory = 4GB,
    [uint64]$Disksize = 50GB,
    [uint64]$OSPartitionSize = 0,
    [uint64]$RecoveryPartitionSize = 0,
    [bool]$CreateRecoveryPartition = $true,
    [object[]]$AdditionalDataPartitions = @(),
    [int]$Processors = 4,
    [bool]$EnableVMNetworking,
    [string]$VMSwitchName,
    [string]$VMLocation,
    [string]$FFUPrefix = '_FFU',
    [string]$FFUCaptureLocation,
    [string]$CustomFFUNameTemplate,
    [bool]$CreateDeploymentMedia,
    [ValidateScript({
            $allowedFeatures = @("Windows-Defender-Default-Definitions", "Printing-PrintToPDFServices-Features", "Printing-XPSServices-Features", "TelnetClient", "TFTP",
                "TIFFIFilter", "LegacyComponents", "DirectPlay", "MSRDC-Infrastructure", "Windows-Identity-Foundation", "MicrosoftWindowsPowerShellV2Root", "MicrosoftWindowsPowerShellV2",
                "SimpleTCP", "NetFx4-AdvSrvs", "NetFx4Extended-ASPNET45", "WCF-Services45", "WCF-HTTP-Activation45", "WCF-TCP-Activation45", "WCF-Pipe-Activation45", "WCF-MSMQ-Activation45",
                "WCF-TCP-PortSharing45", "IIS-WebServerRole", "IIS-WebServer", "IIS-CommonHttpFeatures", "IIS-HttpErrors", "IIS-HttpRedirect", "IIS-ApplicationDevelopment", "IIS-Security",
                "IIS-RequestFiltering", "IIS-NetFxExtensibility", "IIS-NetFxExtensibility45", "IIS-HealthAndDiagnostics", "IIS-HttpLogging", "IIS-LoggingLibraries", "IIS-RequestMonitor",
                "IIS-HttpTracing", "IIS-URLAuthorization", "IIS-IPSecurity", "IIS-Performance", "IIS-HttpCompressionDynamic", "IIS-WebServerManagementTools", "IIS-ManagementScriptingTools",
                "IIS-IIS6ManagementCompatibility", "IIS-Metabase", "WAS-WindowsActivationService", "WAS-ProcessModel", "WAS-NetFxEnvironment", "WAS-ConfigurationAPI", "IIS-HostableWebCore",
                "WCF-HTTP-Activation", "WCF-NonHTTP-Activation", "IIS-StaticContent", "IIS-DefaultDocument", "IIS-DirectoryBrowsing", "IIS-WebDAV", "IIS-WebSockets", "IIS-ApplicationInit",
                "IIS-ISAPIFilter", "IIS-ISAPIExtensions", "IIS-ASPNET", "IIS-ASPNET45", "IIS-ASP", "IIS-CGI", "IIS-ServerSideIncludes", "IIS-CustomLogging", "IIS-BasicAuthentication",
                "IIS-HttpCompressionStatic", "IIS-ManagementConsole", "IIS-ManagementService", "IIS-WMICompatibility", "IIS-LegacyScripts", "IIS-LegacySnapIn", "IIS-FTPServer", "IIS-FTPSvc",
                "IIS-FTPExtensibility", "MSMQ-Container", "MSMQ-DCOMProxy", "MSMQ-Server", "MSMQ-ADIntegration", "MSMQ-HTTP", "MSMQ-Multicast", "MSMQ-Triggers", "IIS-CertProvider",
                "IIS-WindowsAuthentication", "IIS-DigestAuthentication", "IIS-ClientCertificateMappingAuthentication", "IIS-IISCertificateMappingAuthentication", "IIS-ODBCLogging",
                "NetFx3", "SMB1Protocol-Deprecation", "MediaPlayback", "WindowsMediaPlayer", "Client-DeviceLockdown", "Client-EmbeddedShellLauncher", "Client-EmbeddedBootExp",
                "Client-EmbeddedLogon", "Client-KeyboardFilter", "Client-UnifiedWriteFilter", "HostGuardian", "MultiPoint-Connector", "MultiPoint-Connector-Services", "MultiPoint-Tools"
                , "AppServerClient", "SearchEngine-Client-Package", "WorkFolders-Client", "Printing-Foundation-Features", "Printing-Foundation-InternetPrinting-Client",
                "Printing-Foundation-LPDPrintService", "Printing-Foundation-LPRPortMonitor", "HypervisorPlatform", "VirtualMachinePlatform", "Microsoft-Windows-Subsystem-Linux",
                "Client-ProjFS", "Containers-DisposableClientVM", 'Containers-DisposableClientVM', 'Microsoft-Hyper-V-All', 'Microsoft-Hyper-V', 'Microsoft-Hyper-V-Tools-All',
                'Microsoft-Hyper-V-Management-PowerShell', 'Microsoft-Hyper-V-Hypervisor', 'Microsoft-Hyper-V-Services', 'Microsoft-Hyper-V-Management-Clients', 'DataCenterBridging',
                'DirectoryServices-ADAM-Client', 'Windows-Defender-ApplicationGuard', 'ServicesForNFS-ClientOnly', 'ClientForNFS-Infrastructure', 'NFS-Administration', 'Containers', 'Containers-HNS',
                'Containers-SDN', 'SMB1Protocol', 'SMB1Protocol-Client', 'SMB1Protocol-Server', 'SmbDirect')
            $inputFeatures = $_ -split ';'
            foreach ($feature in $inputFeatures) {
                if (-not ($allowedFeatures -contains $feature)) {
                    throw "Invalid optional feature '$feature'. Allowed values: $($allowedFeatures -join ', ')"
                }
            }
            return $true
        })]
    [string]$OptionalFeatures,
    [string]$ProductKey,
    [bool]$BuildUSBDrive,
    [hashtable]$USBDriveList,
    [int]$MaxUSBDrives = 5,
    [ValidateRange(1, 64)]
    [int]$Threads = 5,
    [ValidateSet('Foreground', 'High', 'Normal', 'Low')]
    [string]$BitsPriority = 'Normal',
    [Parameter(Mandatory = $false)]
    [ValidateSet(10, 11, 2016, 2019, 2021, 2022, 2024, 2025)]
    [int]$WindowsRelease = 11,
    [Parameter(Mandatory = $false)]
    [string]$WindowsVersion = '25h2',
    [Parameter(Mandatory = $false)]
    [ValidateSet('x86', 'x64', 'arm64')]
    [string]$WindowsArch = 'x64',
    [ValidateScript({
            $allowedLang = @('ar-sa', 'bg-bg', 'cs-cz', 'da-dk', 'de-de', 'el-gr', 'en-gb', 'en-us', 'es-es', 'es-mx', 'et-ee', 'fi-fi', 'fr-ca', 'fr-fr', 'he-il', 'hr-hr', 'hu-hu',
                'it-it', 'ja-jp', 'ko-kr', 'lt-lt', 'lv-lv', 'nb-no', 'nl-nl', 'pl-pl', 'pt-br', 'pt-pt', 'ro-ro', 'ru-ru', 'sk-sk', 'sl-si', 'sr-latn-rs', 'sv-se', 'th-th', 'tr-tr', 'uk-ua',
                'zh-cn', 'zh-tw')
            if ($allowedLang -contains $_) { $true } else { throw "Invalid WindowsLang value. Allowed values: $($allowedLang -join ', ')" }
            return $true
        })]
    [Parameter(Mandatory = $false)]
    [string]$WindowsLang = 'en-us',
    [Parameter(Mandatory = $false)]
    [ValidateSet('consumer', 'business')]
    [string]$MediaType = 'consumer',
    [ValidateSet(512, 4096)]
    [uint32]$LogicalSectorSizeBytes = 512,
    [string]$SystemPartitionDriveLetter = 'S',
    [string]$WindowsPartitionDriveLetter = 'W',
    [string]$RecoveryPartitionDriveLetter = 'R',
    [bool]$Optimize = $true,
    [int]$OptimizeFFUPartitionNumber = 0,
    [string]$DriversJsonPath,
    [bool]$CompressDownloadedDriversToWim = $false,
    [bool]$CopyDrivers,
    [bool]$CopyPEDrivers,
    [bool]$UseDriversAsPEDrivers,
    [bool]$RemoveFFU,
    [bool]$CopyAdditionalFFUFiles,
    [string[]]$AdditionalFFUFiles,
    [bool]$UpdateLatestCU,
    [bool]$UpdatePreviewCU,
    [bool]$UpdateLatestMicrocode,
    [bool]$UpdateLatestNet,
    [bool]$UpdateLatestDefender,
    [bool]$UpdateLatestMSRT,
    [bool]$UpdateEdge,
    [bool]$UpdateOneDrive,
    [bool]$AllowVHDXCaching,
    [bool]$CopyPPKG,
    [bool]$CopyUnattend,
    [ValidateSet('Legacy', 'None', 'Prompt', 'Template', 'Prefixes', 'SerialComputerNames')]
    [string]$DeviceNamingMode = 'Legacy',
    [string]$DeviceNameTemplate,
    [string[]]$DeviceNamePrefixes,
    [string]$DeviceNamePrefixesPath,
    [string[]]$DeviceNameSerialComputerNames,
    [string]$DeviceNameSerialComputerNamesPath,
    [string]$UnattendX64FilePath,
    [string]$UnattendArm64FilePath,
    [bool]$CopyAutopilot,
    [bool]$CompactOS = $true,
    [bool]$CleanupDeployISO = $true,
    [bool]$CleanupAppsISO = $true,
    [bool]$RemoveUpdates = $true,
    [bool]$RemoveDownloadedESD = $true,
    [bool]$RemoveApps = $true,
    [string]$DriversFolder,
    [string]$PEDriversFolder,
    [bool]$CleanupDrivers = $true,
    [string]$UserAgent = 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/125.0.0.0 Safari/537.36 Edg/125.0.0.0',
    #Microsoft sites will intermittently fail on downloads. These headers are to help with that.
    $Headers = @{
        "Accept"                    = "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7"
        "Accept-Encoding"           = "gzip, deflate, br, zstd"
        "Accept-Language"           = "en-US,en;q=0.9"
        "Priority"                  = "u=0, i"
        "Sec-Ch-Ua"                 = "`"Microsoft Edge`";v=`"125`", `"Chromium`";v=`"125`", `"Not.A/Brand`";v=`"24`""
        "Sec-Ch-Ua-Mobile"          = "?0"
        "Sec-Ch-Ua-Platform"        = "`"Windows`""
        "Sec-Fetch-Dest"            = "document"
        "Sec-Fetch-Mode"            = "navigate"
        "Sec-Fetch-Site"            = "none"
        "Sec-Fetch-User"            = "?1"
        "Upgrade-Insecure-Requests" = "1"
    },
    [bool]$AllowExternalHardDiskMedia,
    [bool]$PromptExternalHardDiskMedia = $true,
    [Parameter(Mandatory = $false)]
    [ValidateScript({ $_ -eq $null -or (Test-Path $_) })]
    [string]$ConfigFile,
    [Parameter(Mandatory = $false)]
    [string]$ExportConfigFile,
    [bool]$InjectUnattend = $false,
    [string]$orchestrationPath,
    [bool]$UpdateADK = $true,
    [bool]$CleanupCurrentRunDownloads = $false,
    [switch]$Cleanup
)
$ProgressPreference = 'SilentlyContinue'
$version = '2608.1'

# Remove any existing modules to avoid conflicts
if (Get-Module -Name 'FFU.Common.Core' -ErrorAction SilentlyContinue) {
    Remove-Module -Name 'FFU.Common.Core' -Force
}
if (Get-Module -Name 'FFU.Common.Winget' -ErrorAction SilentlyContinue) {
    Remove-Module -Name 'FFU.Common.Winget' -Force
}
if (Get-Module -Name 'FFU.Common.Drivers' -ErrorAction SilentlyContinue) {
    Remove-Module -Name 'FFU.Common.Drivers' -Force
}
# Import the required modules
Import-Module "$PSScriptRoot\FFU.Common" -Force

# Set the module's verbose preference to match the script's - allows logging verbose output to console.
$moduleInfo = Get-Module -Name 'FFU.Common'
if ($moduleInfo) {
    & $moduleInfo { $script:VerbosePreference = $args[0] } $VerbosePreference | Out-Null
}

# If a config file is specified and it exists, load it
if ($ConfigFile -and (Test-Path -Path $ConfigFile)) {
    $configData = Get-Content $ConfigFile -Raw | ConvertFrom-Json
    $keys = $configData.psobject.Properties.Name
    
    # Iterate through the keys in the config data
    foreach ($key in $keys) {
        $value = $configdata.$key

        $valueIsEmptyString = ($value -is [string]) -and [string]::IsNullOrEmpty($value)
        $valueIsEmptyArray = ($value -is [System.Array]) -and ($value.Count -eq 0)
        $valueIsEmptyHashtable = ($value -is [System.Collections.Hashtable]) -and ($value.Count -eq 0)
        $valueIsZero = (($value -is [System.UInt32]) -or ($value -is [System.UInt64]) -or ($value -is [System.Int32])) -and ($value -eq 0)
        
        # If $value is empty, skip
        if ($null -eq $value -or 
            $valueIsEmptyString -or 
            $valueIsEmptyArray -or 
            $valueIsEmptyHashtable -or 
            $valueIsZero) {
            continue
        }

        # If this is the Headers parameter, convert PSCustomObject to hashtable
        if ((($key -eq 'Headers') -or ($key -eq 'AppsScriptVariables') -or ($key -eq 'USBDriveList')) -and ($value -is [System.Management.Automation.PSCustomObject])) {
            $hashtableValue = [hashtable]::new()
            foreach ($prop in $value.psobject.Properties) {
                $hashtableValue[$prop.Name] = $prop.Value
            }
            $value = $hashtableValue
        }

        # Check if this key matches a parameter in the script
        # and if the user did NOT explicitly supply it on the command line
        if ($MyInvocation.MyCommand.Parameters.ContainsKey($key) -and -not $PSBoundParameters.ContainsKey($key)) {
            # Set the parameter's value to what's in the config file
            Set-Variable -Name $key -Value $value -Scope 0
        }
    }
}

function Get-UnattendSourcePath {
    param(
        [Parameter(Mandatory = $true)]
        [string]$UnattendFolder,
        [Parameter(Mandatory = $true)]
        [string]$WindowsArch,
        [string]$UnattendX64FilePath,
        [string]$UnattendArm64FilePath
    )

    $resolvedArch = if ($WindowsArch -ieq 'arm64') { 'arm64' } else { 'x64' }
    $resolvedSourcePath = if ($resolvedArch -eq 'arm64') {
        if ([string]::IsNullOrWhiteSpace($UnattendArm64FilePath)) {
            Join-Path $UnattendFolder 'unattend_arm64.xml'
        }
        else {
            $UnattendArm64FilePath
        }
    }
    else {
        if ([string]::IsNullOrWhiteSpace($UnattendX64FilePath)) {
            Join-Path $UnattendFolder 'unattend_x64.xml'
        }
        else {
            $UnattendX64FilePath
        }
    }

    WriteLog "Resolved unattend source path for ${resolvedArch}: $resolvedSourcePath"
    return $resolvedSourcePath
}

function Initialize-UnattendComputerNamePath {
    param(
        [Parameter(Mandatory = $true)]
        [xml]$UnattendXml,
        [Parameter(Mandatory = $true)]
        [string]$WindowsArch
    )

    $unattendRoot = $UnattendXml.DocumentElement
    if (($null -eq $unattendRoot) -or ($unattendRoot.LocalName -ne 'unattend')) {
        throw 'Unattend XML is missing the unattend root element.'
    }

    $unattendNamespace = $unattendRoot.NamespaceURI
    if ([string]::IsNullOrWhiteSpace($unattendNamespace)) {
        throw 'Unattend XML is missing the default unattend namespace.'
    }

    $namespaceManager = New-Object System.Xml.XmlNamespaceManager($UnattendXml.NameTable)
    $namespaceManager.AddNamespace('un', $unattendNamespace)

    $specializeSettings = $unattendRoot.SelectSingleNode("un:settings[@pass='specialize']", $namespaceManager)
    $createdSpecializeSettings = $false
    if ($null -eq $specializeSettings) {
        $specializeSettings = $UnattendXml.CreateElement('settings', $unattendNamespace)
        $null = $specializeSettings.SetAttribute('pass', 'specialize')
        $firstSettingsNode = $unattendRoot.SelectSingleNode('un:settings', $namespaceManager)
        if ($null -ne $firstSettingsNode) {
            $null = $unattendRoot.InsertBefore($specializeSettings, $firstSettingsNode)
        }
        else {
            $null = $unattendRoot.AppendChild($specializeSettings)
        }
        $createdSpecializeSettings = $true
    }

    $shellSetupComponent = $specializeSettings.SelectSingleNode("un:component[@name='Microsoft-Windows-Shell-Setup']", $namespaceManager)
    $createdShellSetupComponent = $false
    if ($null -eq $shellSetupComponent) {
        $processorArchitecture = if ($WindowsArch -ieq 'arm64') { 'arm64' } else { 'amd64' }
        $shellSetupComponent = $UnattendXml.CreateElement('component', $unattendNamespace)
        $null = $shellSetupComponent.SetAttribute('name', 'Microsoft-Windows-Shell-Setup')
        $null = $shellSetupComponent.SetAttribute('processorArchitecture', $processorArchitecture)
        $null = $shellSetupComponent.SetAttribute('publicKeyToken', '31bf3856ad364e35')
        $null = $shellSetupComponent.SetAttribute('language', 'neutral')
        $null = $shellSetupComponent.SetAttribute('versionScope', 'nonSxS')
        $null = $shellSetupComponent.SetAttribute('xmlns:wcm', 'http://www.w3.org/2000/xmlns/', 'http://schemas.microsoft.com/WMIConfig/2002/State')
        $null = $shellSetupComponent.SetAttribute('xmlns:xsi', 'http://www.w3.org/2000/xmlns/', 'http://www.w3.org/2001/XMLSchema-instance')

        $firstComponentNode = $specializeSettings.SelectSingleNode('un:component', $namespaceManager)
        if ($null -ne $firstComponentNode) {
            $null = $specializeSettings.InsertBefore($shellSetupComponent, $firstComponentNode)
        }
        else {
            $null = $specializeSettings.AppendChild($shellSetupComponent)
        }
        $createdShellSetupComponent = $true
    }

    $computerNameElement = $shellSetupComponent.SelectSingleNode('un:ComputerName', $namespaceManager)
    $createdComputerNameElement = $false
    if ($null -eq $computerNameElement) {
        $computerNameElement = $UnattendXml.CreateElement('ComputerName', $unattendNamespace)
        $null = $shellSetupComponent.AppendChild($computerNameElement)
        $createdComputerNameElement = $true
    }

    return [PSCustomObject]@{
        ComputerNameElement        = $computerNameElement
        CreatedSpecializeSettings  = $createdSpecializeSettings
        CreatedShellSetupComponent = $createdShellSetupComponent
        CreatedComputerNameElement = $createdComputerNameElement
    }
}

function Save-StagedUnattendFile {
    param(
        [Parameter(Mandatory = $true)]
        [string]$SourcePath,
        [Parameter(Mandatory = $true)]
        [string]$DestinationPath,
        [Parameter(Mandatory = $true)]
        [ValidateSet('Legacy', 'None', 'Prompt', 'Template', 'Prefixes', 'SerialComputerNames')]
        [string]$DeviceNamingMode,
        [string]$DeviceNameTemplate,
        [Parameter(Mandatory = $true)]
        [string]$WindowsArch,
        [bool]$LegacyPrefixesWillBeStaged = $false
    )

    if ($DeviceNamingMode -eq 'None') {
        Copy-Item -Path $SourcePath -Destination $DestinationPath -Force | Out-Null
        return
    }

    [xml]$unattendXml = Get-Content -Path $SourcePath
    $computerNamePath = Initialize-UnattendComputerNamePath -UnattendXml $unattendXml -WindowsArch $WindowsArch

    if ($computerNamePath.CreatedSpecializeSettings -or $computerNamePath.CreatedShellSetupComponent -or $computerNamePath.CreatedComputerNameElement) {
        $createdParts = @()
        if ($computerNamePath.CreatedSpecializeSettings) {
            $createdParts += 'specialize settings'
        }
        if ($computerNamePath.CreatedShellSetupComponent) {
            $createdParts += 'Microsoft-Windows-Shell-Setup component'
        }
        if ($computerNamePath.CreatedComputerNameElement) {
            $createdParts += 'ComputerName element'
        }
        WriteLog "Created $($createdParts -join ', ') while staging unattend file $DestinationPath"
    }

    if ($DeviceNamingMode -eq 'Prompt') {
        $computerNamePath.ComputerNameElement.InnerText = 'MyComputer'
    }
    elseif ($DeviceNamingMode -eq 'Template') {
        $computerNamePath.ComputerNameElement.InnerText = $DeviceNameTemplate
    }
    elseif ($DeviceNamingMode -eq 'Prefixes') {
        if ($computerNamePath.CreatedComputerNameElement) {
            $computerNamePath.ComputerNameElement.InnerText = '*'
        }
    }
    elseif ($DeviceNamingMode -eq 'SerialComputerNames') {
        if ($computerNamePath.CreatedComputerNameElement) {
            $computerNamePath.ComputerNameElement.InnerText = '*'
        }
    }
    elseif (($DeviceNamingMode -eq 'Legacy') -and $computerNamePath.CreatedComputerNameElement) {
        $computerNamePath.ComputerNameElement.InnerText = if ($LegacyPrefixesWillBeStaged) { '*' } else { 'MyComputer' }
    }

    $unattendXml.Save($DestinationPath)
}

$vmSwitchWasExplicitlyBound = $PSBoundParameters.ContainsKey('VMSwitchName')
$enableVmNetworkingWasExplicitlyBound = $PSBoundParameters.ContainsKey('EnableVMNetworking')
if (-not $EnableVMNetworking -and $vmSwitchWasExplicitlyBound -and -not $enableVmNetworkingWasExplicitlyBound) {
    $EnableVMNetworking = $true
    WriteLog 'EnableVMNetworking not explicitly set. Enabling VM networking because -VMSwitchName was supplied on the command line.'
}

$normalizedDeviceNameTemplate = if ($null -ne $DeviceNameTemplate) { $DeviceNameTemplate.Trim() } else { $null }
$effectiveDeviceNamePrefixes = @($DeviceNamePrefixes | Where-Object { -not [string]::IsNullOrWhiteSpace($_) } | ForEach-Object { $_.Trim() })
$resolvedDeviceNamePrefixesPath = if ([string]::IsNullOrWhiteSpace($DeviceNamePrefixesPath)) {
    Join-Path (Join-Path $FFUDevelopmentPath 'Unattend') 'prefixes.txt'
}
else {
    $DeviceNamePrefixesPath
}
$effectiveDeviceNameSerialComputerNames = @($DeviceNameSerialComputerNames | Where-Object { -not [string]::IsNullOrWhiteSpace($_) } | ForEach-Object { $_.Trim() })
$resolvedDeviceNameSerialComputerNamesPath = if ([string]::IsNullOrWhiteSpace($DeviceNameSerialComputerNamesPath)) {
    Join-Path (Join-Path $FFUDevelopmentPath 'Unattend') 'SerialComputerNames.csv'
}
else {
    $DeviceNameSerialComputerNamesPath
}

if (($DeviceNamingMode -eq 'Prefixes') -and ($effectiveDeviceNamePrefixes.Count -eq 0) -and (Test-Path -Path $resolvedDeviceNamePrefixesPath -PathType Leaf)) {
    $effectiveDeviceNamePrefixes = @(Get-Content -Path $resolvedDeviceNamePrefixesPath | Where-Object { -not [string]::IsNullOrWhiteSpace($_) } | ForEach-Object { $_.Trim() })
    WriteLog "Loaded device name prefixes from $resolvedDeviceNamePrefixesPath"
}

if (($DeviceNamingMode -eq 'SerialComputerNames') -and ($effectiveDeviceNameSerialComputerNames.Count -eq 0) -and (Test-Path -Path $resolvedDeviceNameSerialComputerNamesPath -PathType Leaf)) {
    $effectiveDeviceNameSerialComputerNames = @(Get-Content -Path $resolvedDeviceNameSerialComputerNamesPath | Where-Object { -not [string]::IsNullOrWhiteSpace($_) } | ForEach-Object { $_.Trim() })
    WriteLog "Loaded serial computer-name mappings from $resolvedDeviceNameSerialComputerNamesPath"
}

if ($CopyUnattend -and $InjectUnattend) {
    throw 'CopyUnattend and InjectUnattend cannot both be set to `$true. Select only one unattend delivery method.'
}

if ($DeviceNamingMode -eq 'Template') {
    if ([string]::IsNullOrWhiteSpace($normalizedDeviceNameTemplate)) {
        throw 'DeviceNamingMode Template requires DeviceNameTemplate.'
    }

    $templateWithoutSupportedVariables = $normalizedDeviceNameTemplate -replace '(?i)%serial%', ''
    if ($templateWithoutSupportedVariables -match '%') {
        throw 'Only the %serial% device name variable is supported.'
    }

    if (-not ($CopyUnattend -or $InjectUnattend)) {
        throw 'DeviceNamingMode Template requires either CopyUnattend or InjectUnattend.'
    }

    if ($InjectUnattend -and (-not $CopyUnattend) -and $normalizedDeviceNameTemplate -match '(?i)%serial%') {
        throw 'The %serial% device name variable is only supported when CopyUnattend is used.'
    }
}
elseif ($DeviceNamingMode -eq 'Prompt') {
    if (-not $CopyUnattend) {
        throw 'DeviceNamingMode Prompt requires CopyUnattend. Prompt-based naming is not supported with InjectUnattend.'
    }
}
elseif ($DeviceNamingMode -eq 'Prefixes') {
    if (-not $CopyUnattend) {
        throw 'DeviceNamingMode Prefixes requires CopyUnattend. Prefix-based naming is not supported with InjectUnattend.'
    }

    if ($effectiveDeviceNamePrefixes.Count -eq 0) {
        throw 'DeviceNamingMode Prefixes requires at least one DeviceNamePrefixes entry or a valid DeviceNamePrefixesPath.'
    }
}
elseif ($DeviceNamingMode -eq 'SerialComputerNames') {
    if (-not $CopyUnattend) {
        throw 'DeviceNamingMode SerialComputerNames requires CopyUnattend. Serial-to-computer-name mapping is not supported with InjectUnattend.'
    }

    if ($effectiveDeviceNameSerialComputerNames.Count -eq 0) {
        throw 'DeviceNamingMode SerialComputerNames requires DeviceNameSerialComputerNames content or a valid DeviceNameSerialComputerNamesPath.'
    }

    try {
        $serialComputerNameMappings = @($effectiveDeviceNameSerialComputerNames | ConvertFrom-Csv -ErrorAction Stop)
    }
    catch {
        throw "DeviceNamingMode SerialComputerNames requires valid CSV content with SerialNumber and ComputerName headers. $($_.Exception.Message)"
    }

    if ($serialComputerNameMappings.Count -eq 0) {
        throw 'DeviceNamingMode SerialComputerNames requires at least one CSV data row.'
    }

    $serialComputerNameHeaders = @($serialComputerNameMappings[0].PSObject.Properties.Name)
    if ((-not ($serialComputerNameHeaders -contains 'SerialNumber')) -or (-not ($serialComputerNameHeaders -contains 'ComputerName'))) {
        throw 'DeviceNamingMode SerialComputerNames requires SerialNumber and ComputerName headers.'
    }

    $validSerialComputerNameMappings = @($serialComputerNameMappings | Where-Object {
            -not [string]::IsNullOrWhiteSpace([string]$_.SerialNumber) -and -not [string]::IsNullOrWhiteSpace([string]$_.ComputerName)
        })
    if ($validSerialComputerNameMappings.Count -eq 0) {
        throw 'DeviceNamingMode SerialComputerNames requires at least one row with both SerialNumber and ComputerName values.'
    }
}

# Validate that the selected Windows SKU is compatible with the chosen Windows release and ensure an ISO is provided for unsupported releases
$clientSKUs = @(
    'Home',
    'Home N',
    'Home Single Language',
    'Education',
    'Education N',
    'Pro',
    'Pro N',
    'Pro Education',
    'Pro Education N',
    'Pro for Workstations',
    'Pro N for Workstations',
    'Enterprise',
    'Enterprise N'
)
$LTSCSKUs = @(
    'Enterprise 2016 LTSB',
    'Enterprise N 2016 LTSB',
    'Enterprise LTSC',
    'Enterprise N LTSC',
    'IoT Enterprise LTSC',
    'IoT Enterprise N LTSC'
)
$ServerSKUs = @(
    'Standard',
    'Standard (Desktop Experience)',
    'Datacenter',
    'Datacenter (Desktop Experience)'
)
$releaseToSKUMapping = @{
    10   = $clientSKUs
    11   = $clientSKUs
    2016 = $LTSCSKUs + $ServerSKUs
    2019 = $LTSCSKUs + $ServerSKUs
    2021 = $LTSCSKUs
    2022 = $ServerSKUs
    2024 = $LTSCSKUs
    2025 = $ServerSKUs
}
if ($releaseToSKUMapping.ContainsKey($WindowsRelease) -and $WindowsSKU -notin $releaseToSKUMapping[$WindowsRelease]) {
    throw "Selected SKU is $WindowsSKU. Windows $WindowsRelease requires one of these SKUs: $($releaseToSKUMapping[$WindowsRelease] -join ', ')"
}
if ($WindowsRelease -notin 10, 11 -and -not $ISOPath) {
    throw "Windows $WindowsRelease cannot automatically be downloaded. Please specify your own ISO using the -ISOPath parameter."
}

#Class definition for vhdx cache
$partitionLayoutSignatureVersion = 2

class VhdxCacheUpdateItem {
    [string]$Name
    VhdxCacheUpdateItem([string]$Name) {
        $this.Name = $Name
    }
}

class VhdxCacheItem {
    [string]$VhdxFileName = ""
    [uint32]$LogicalSectorSizeBytes = ""
    [uint64]$Disksize = ""
    [bool]$CreateRecoveryPartition = $true
    [uint32]$PartitionLayoutSignatureVersion = 0
    [string]$PartitionLayoutSignature = ""
    [string]$WindowsSKU = ""
    [string]$WindowsRelease = ""
    [string]$WindowsVersion = ""
    [string]$OptionalFeatures = ""
    [VhdxCacheUpdateItem[]]$IncludedUpdates = @()
}

#Support for ini reading
$definition = @'
[DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
public static extern uint GetPrivateProfileString(
    string lpAppName,
    string lpKeyName,
    string lpDefault,
    System.Text.StringBuilder lpReturnedString,
    uint nSize,
    string lpFileName);

[DllImport("kernel32.dll", CharSet = CharSet.Auto)]
public static extern uint GetPrivateProfileSection(
    string lpAppName,
    byte[] lpReturnedString,
    uint nSize,
    string lpFileName);
'@
Add-Type -MemberDefinition $definition -Namespace Win32 -Name Kernel32 -PassThru | Out-Null

function Get-WindowsTargetRuntimeState {
    param(
        [Parameter(Mandatory = $true)]
        [int]$WindowsRelease,

        [Parameter(Mandatory = $true)]
        [string]$WindowsSKU,

        [Parameter(Mandatory = $true)]
        [string]$CurrentWindowsVersion,

        [Parameter(Mandatory = $true)]
        [bool]$UpdateLatestCU
    )

    $localInstallationType = if ($WindowsSKU -like 'Standard*' -or $WindowsSKU -like 'Datacenter*') { 'Server' } else { 'Client' }
    $localWindowsVersion = $CurrentWindowsVersion
    $localIsLTSC = $false

    if ($localInstallationType -eq 'Server') {
        switch ($WindowsRelease) {
            2016 { $localWindowsVersion = '1607' }
            2019 { $localWindowsVersion = '1809' }
            2022 { $localWindowsVersion = '21H2' }
            2025 { $localWindowsVersion = '24H2' }
        }
    }

    if ($WindowsSKU -like '*LTS*') {
        switch ($WindowsRelease) {
            2016 { $localWindowsVersion = '1607' }
            2019 { $localWindowsVersion = '1809' }
            2021 { $localWindowsVersion = '21H2' }
            2024 { $localWindowsVersion = '24H2' }
        }
        $localIsLTSC = $true
    }

    $localIsWindows10LtscClient = ($localInstallationType -eq 'Client') -and ($WindowsRelease -in 2016, 2019, 2021) -and $localIsLTSC

    return [pscustomobject]@{
        InstallationType      = $localInstallationType
        WindowsVersion        = $localWindowsVersion
        IsLTSC                = $localIsLTSC
        IsWindows10LtscClient = $localIsWindows10LtscClient
        InstallLatestCuInVm   = ($UpdateLatestCU -and $localIsWindows10LtscClient)
    }
}

#Check if Hyper-V feature is installed (requires only checks the module)
$osInfo = Get-CimInstance -ClassName win32_OperatingSystem
$isServer = $osInfo.Caption -match 'server'

if ($isServer) {
    $hyperVFeature = Get-WindowsFeature -Name Hyper-V
    if ($hyperVFeature.InstallState -ne "Installed") {
        Write-Host "Hyper-V feature is not installed. Please install it before running this script."
        exit
    }
}
else {
    $hyperVFeature = Get-WindowsOptionalFeature -Online -FeatureName Microsoft-Hyper-V-All
    if ($hyperVFeature.State -ne "Enabled") {
        Write-Host "Hyper-V feature is not enabled. Please enable it before running this script."
        exit
    }
}

# Set default values for variables that depend on other parameters
if (-not $AppsISO) { $AppsISO = "$FFUDevelopmentPath\Apps.iso" }
if (-not $AppsPath) { $AppsPath = "$FFUDevelopmentPath\Apps" }
if (-not $AppListPath) { $AppListPath = "$AppsPath\AppList.json" }
if (-not $UserAppListPath) { $UserAppListPath = "$AppsPath\UserAppList.json" }
if (-not $OrchestrationPath) { $OrchestrationPath = "$AppsPath\Orchestration" }
if (-not $appInstallConfigPath) { $appInstallConfigPath = "$OrchestrationPath\AppInstallConfig.json" }
if (-not $wingetWin32jsonFile) { $wingetWin32jsonFile = "$OrchestrationPath\WinGetWin32Apps.json" }
if (-not $InstallOfficePath) { $InstallOfficePath = "$OrchestrationPath\Install-Office.ps1" }
if (-not $InstallDefenderPath) { $InstallDefenderPath = "$OrchestrationPath\Update-Defender.ps1" }
if (-not $InstallMSRTPath) { $InstallMSRTPath = "$OrchestrationPath\Update-MSRT.ps1" }
if (-not $InstallODPath) { $InstallODPath = "$OrchestrationPath\Update-OneDrive.ps1" }
if (-not $InstallEdgePath) { $InstallEdgePath = "$OrchestrationPath\Update-Edge.ps1" }

# Set default LTSC CU in-VM orchestration paths
if (-not $InstallLTSCUpdatePath) { $InstallLTSCUpdatePath = "$OrchestrationPath\Install-LTSCUpdate.ps1" }
if (-not $LtscCUStagePath) { $LtscCUStagePath = "$AppsPath\LTSCUpdate" }

if (-not $AppsScriptVarsJsonPath) { $AppsScriptVarsJsonPath = "$OrchestrationPath\AppsScriptVariables.json" }
if (-not $DeployISO) { $DeployISO = "$FFUDevelopmentPath\WinPE_FFU_Deploy_$WindowsArch.iso" }
if (-not $OfficePath) { $OfficePath = "$AppsPath\Office" }
if (-not $OfficeDownloadXML) { $OfficeDownloadXML = "$OfficePath\DownloadFFU.xml" }
if (-not $OfficeInstallXML) { $OfficeInstallXML = "DeployFFU.xml" }
if (-not $rand) { $rand = Get-Random }
if (-not $VMLocation) { $VMLocation = "$FFUDevelopmentPath\VM" }
if (-not $VMName) { $VMName = "$FFUPrefix-$rand" }
if (-not $VMPath) { $VMPath = "$VMLocation\$VMName" }
if (-not $VHDXPath) { $VHDXPath = "$VMPath\$VMName.vhdx" }
if (-not $FFUCaptureLocation) { $FFUCaptureLocation = "$FFUDevelopmentPath\FFU" }
if (-not $LogFile) { $LogFile = "$FFUDevelopmentPath\FFUDevelopment.log" }
if (-not $KBPath) { $KBPath = "$FFUDevelopmentPath\KB" }
if (-not $MicrocodePath) { $MicrocodePath = "$KBPath\Microcode" }
if (-not $DefenderPath) { $DefenderPath = "$AppsPath\Defender" }
if (-not $MSRTPath) { $MSRTPath = "$AppsPath\MSRT" }
if (-not $OneDrivePath) { $OneDrivePath = "$AppsPath\OneDrive" }
if (-not $EdgePath) { $EdgePath = "$AppsPath\Edge" }
if (-not $DriversFolder) { $DriversFolder = "$FFUDevelopmentPath\Drivers" }
if (-not $PPKGFolder) { $PPKGFolder = "$FFUDevelopmentPath\PPKG" }
if (-not $UnattendFolder) { $UnattendFolder = "$FFUDevelopmentPath\Unattend" }
if ([string]::IsNullOrWhiteSpace($UnattendX64FilePath)) { $UnattendX64FilePath = Join-Path $UnattendFolder 'unattend_x64.xml' }
if ([string]::IsNullOrWhiteSpace($UnattendArm64FilePath)) { $UnattendArm64FilePath = Join-Path $UnattendFolder 'unattend_arm64.xml' }
if (-not $AutopilotFolder) { $AutopilotFolder = "$FFUDevelopmentPath\Autopilot" }
if (-not $PEDriversFolder) { $PEDriversFolder = "$FFUDevelopmentPath\PEDrivers" }
if (-not $VHDXCacheFolder) { $VHDXCacheFolder = "$FFUDevelopmentPath\VHDXCache" }
if (-not $AppListPath) { $AppListPath = "$AppsPath\AppList.json" }
$windowsTargetRuntimeState = Get-WindowsTargetRuntimeState -WindowsRelease $WindowsRelease -WindowsSKU $WindowsSKU -CurrentWindowsVersion $WindowsVersion -UpdateLatestCU:$UpdateLatestCU
$installationType = $windowsTargetRuntimeState.InstallationType
$WindowsVersion = $windowsTargetRuntimeState.WindowsVersion
$isLTSC = $windowsTargetRuntimeState.IsLTSC
$isWindows10LtscClient = $windowsTargetRuntimeState.IsWindows10LtscClient
$installLatestCuInVm = $windowsTargetRuntimeState.InstallLatestCuInVm
$refreshAppsIsoForLtscCu = $false

# Set the log path for the common logger
Set-CommonCoreLogPath -Path $LogFile
Set-BitsTransferPriority -Priority $BitsPriority

#FUNCTIONS

function Get-EffectiveDriverWindowsRelease {
    param(
        [Parameter(Mandatory = $true)]
        [int]$WindowsRelease,
        [string]$WindowsSKU
    )

    # Normalize LTSC/LTSB client releases to OEM driver release values.
    if (-not [string]::IsNullOrWhiteSpace($WindowsSKU) -and $WindowsSKU -like '*LTS*') {
        if ($WindowsRelease -in 2016, 2019, 2021) {
            return 10
        }
        if ($WindowsRelease -eq 2024) {
            return 11
        }
    }

    return $WindowsRelease
}

function Get-Parameters {
    [CmdletBinding()]
    param (
        [Parameter()]
        $ParamNames
    )
    # Define unwanted parameters
    $excludedParams = 'Debug', 'ErrorAction', 'ErrorVariable', 'InformationAction', 'InformationVariable', 'OutBuffer', 'OutVariable', 'PipelineVariable', 'Verbose', 'WarningAction', 'WarningVariable', 'ProgressAction'

    # Filter out the unwanted parameters
    $filteredParamNames = $paramNames | Where-Object { $excludedParams -notcontains $_ }
    return $filteredParamNames
}

function LogVariableValues {
    $excludedVariables = @(
        'PSBoundParameters', 
        'PSScriptRoot', 
        'PSCommandPath', 
        'MyInvocation', 
        '?', 
        'ConsoleFileName', 
        'ExecutionContext',
        'false',
        'HOME',
        'Host',
        'hyperVFeature',
        'input',
        'MaximumAliasCount',
        'MaximumDriveCount',
        'MaximumErrorCount',
        'MaximumFunctionCount',
        'MaximumVariableCount',
        'null',
        'PID',
        'PSCmdlet',
        'PSCulture',
        'PSUICulture',
        'PSVersionTable',
        'ShellId',
        'true'
    )

    $allVariables = Get-Variable -Scope Script | Where-Object { $_.Name -notin $excludedVariables }
    Writelog "Script version: $version"
    WriteLog 'Logging variables'
    foreach ($variable in $allVariables) {
        $variableName = $variable.Name
        $variableValue = $variable.Value
        if ($null -ne $variableValue) {
            WriteLog "[VAR]$variableName`: $variableValue"
        }
        else {
            WriteLog "[VAR]Variable $variableName not found or not set"
        }
    }
    WriteLog 'End logging variables'
}

function Get-ChildProcesses($parentId) {
    $result = @()
    $children = Get-CimInstance Win32_Process -Filter "ParentProcessId = $parentId"
    foreach ($child in $children) {
        $result += $child
        $result += Get-ChildProcesses $child.ProcessId
    }
    return $result
}

function Test-Url {
    param (
        [Parameter(Mandatory = $true)]
        [string]$Url
    )
    try {
        # Create a web request and check the response
        $request = [System.Net.WebRequest]::Create($Url)
        $request.Method = 'HEAD'
        $response = $request.GetResponse()
        return $true
    }
    catch {
        return $false
    }
}

function Get-MicrosoftDrivers {
    param (
        [string]$Make,
        [string]$Model,
        [int]$WindowsRelease
    )

    # Use the shared Learn-based Microsoft model index so the CLI stays aligned with the UI.
    WriteLog "Getting Surface driver information from the shared Microsoft model index"
    $models = @(Get-SurfaceDriverModelIndex -DriversFolder $DriversFolder | Select-Object Model, Link)
    WriteLog "Complete"
    WriteLog "Loaded $($models.Count) Surface models from the shared Microsoft model index."

    ### FIND THE MODEL IN THE LIST OF MODELS
    $selectedModel = $models | Where-Object { $_.Model -eq $Model }

    if ($null -eq $selectedModel) {
        if ($VerbosePreference -ne 'Continue') {
            Write-Host "The model '$Model' was not found in the list of available models."
            Write-Host "Please run the script with the -Verbose switch to see the list of available models."
        }
        WriteLog "The model '$Model' was not found in the list of available models."
        WriteLog "Please select a model from the list below by number:"

        for ($i = 0; $i -lt $models.Count; $i++) {
            if ($VerbosePreference -ne 'Continue') {
                Write-Host "$($i + 1). $($models[$i].Model)"
            }
            WriteLog "$($i + 1). $($models[$i].Model)"
        }

        do {
            $selection = Read-Host "Enter the number of the model you want to select"
            WriteLog "User selected model number: $selection"

            if ($selection -match '^\d+$' -and [int]$selection -ge 1 -and [int]$selection -le $models.Count) {
                $selectedModel = $models[$selection - 1]
            }
            else {
                if ($VerbosePreference -ne 'Continue') {
                    Write-Host "Invalid selection. Please try again."
                }
                WriteLog "Invalid selection. Please try again."
            }
        } while ($null -eq $selectedModel)
    }

    $Model = $selectedModel.Model
    WriteLog "Model: $Model"
    WriteLog "Download Page: $($selectedModel.Link)"

    ### GET THE DOWNLOAD LINK FOR THE SELECTED MODEL
    WriteLog "Getting download page content"
    $OriginalVerbosePreference = $VerbosePreference
    $VerbosePreference = 'SilentlyContinue'
    $downloadPageContent = Invoke-WebRequest -Uri $selectedModel.Link -UseBasicParsing -Headers $Headers -UserAgent $UserAgent
    $VerbosePreference = $OriginalVerbosePreference
    WriteLog "Complete"
    WriteLog "Parsing download page for file"
    $scriptPattern = '<script>window.__DLCDetails__={(.*?)}<\/script>'
    $scriptMatch = [regex]::Match($downloadPageContent.Content, $scriptPattern)

    if ($scriptMatch.Success) {
        $scriptContent = $scriptMatch.Groups[1].Value

        # Extract the download file information from the script tag
        $downloadFilePattern = '"name":"(.*?)",.*?"url":"(.*?)"'
        $downloadFileMatches = [regex]::Matches($scriptContent, $downloadFilePattern)

        $downloadLink = $null
        foreach ($downloadFile in $downloadFileMatches) {
            $fileName = $downloadFile.Groups[1].Value
            $fileUrl = $downloadFile.Groups[2].Value

            if ($fileName -match "Win$WindowsRelease") {
                $downloadLink = $fileUrl
                break
            }
        }


        ### CREATE FOLDER STRUCTURE AND DOWNLOAD AND EXTRACT THE FILE
        if ($downloadLink) {
            WriteLog "Download Link for Windows ${WindowsRelease}: $downloadLink"

            # Create directory structure
            if (-not (Test-Path -Path $DriversFolder)) {
                WriteLog "Creating Drivers folder: $DriversFolder"
                New-Item -Path $DriversFolder -ItemType Directory -Force | Out-Null
                WriteLog "Drivers folder created"
            }
            $sanitizedModel = ConvertTo-SafeName -Name $Model
            if ($sanitizedModel -ne $Model) { WriteLog "Sanitized model name: '$Model' -> '$sanitizedModel'" }
            $surfaceDriversPath = Join-Path -Path $DriversFolder -ChildPath $Make
            $modelPath = Join-Path -Path $surfaceDriversPath -ChildPath $sanitizedModel
            if (-Not (Test-Path -Path $modelPath)) {
                WriteLog "Creating model folder: $modelPath"
                New-Item -Path $modelPath -ItemType Directory | Out-Null
                WriteLog "Complete"
            }

            ### DOWNLOAD THE FILE
            $filePath = Join-Path -Path $surfaceDriversPath -ChildPath ($fileName)
            WriteLog "Downloading $Model driver file to $filePath"
            Mark-DownloadInProgress -FFUDevelopmentPath $FFUDevelopmentPath -TargetPath $filePath
            Start-BitsTransferWithRetry -Source $downloadLink -Destination $filePath
            Clear-DownloadInProgress -FFUDevelopmentPath $FFUDevelopmentPath -TargetPath $filePath
            WriteLog "Download complete"

            # Determine file extension
            $fileExtension = [System.IO.Path]::GetExtension($filePath).ToLower()

            ### EXTRACT THE FILE
            if ($fileExtension -eq ".msi") {
                # Extract the MSI file using an administrative install
                WriteLog "Extracting MSI file to $modelPath"
                $arguments = "/a `"$($filePath)`" /qn TARGETDIR=`"$($modelPath)`""
                Invoke-Process -FilePath "msiexec.exe" -ArgumentList $arguments | Out-Null
                WriteLog "Extraction complete"
            }
            elseif ($fileExtension -eq ".zip") {
                # Extract the ZIP file
                WriteLog "Extracting ZIP file to $modelPath"
                # $ProgressPreference = 'SilentlyContinue'
                Expand-Archive -Path $filePath -DestinationPath $modelPath -Force
                # $ProgressPreference = 'Continue'
                WriteLog "Extraction complete"
            }
            else {
                WriteLog "Unsupported file type: $fileExtension"
            }
            # Remove the downloaded file
            WriteLog "Removing $filePath"
            Remove-Item -Path $filePath -Force
            WriteLog "Complete"
        }
        else {
            WriteLog "No download link found for Windows $WindowsRelease."
        }
    }
    else {
        WriteLog "Failed to parse the download page for the MSI file."
    }
}

function Get-HPDrivers {
    [CmdletBinding()]
    param (
        [Parameter()]
        [string]$Make,
        [Parameter()]
        [string]$Model,
        [Parameter()]
        [ValidateSet("x64", "x86", "ARM64")]
        [string]$WindowsArch,
        [Parameter()]
        [ValidateSet(10, 11)]
        [int]$WindowsRelease,
        [Parameter()]
        [string]$WindowsVersion
    )

    # Download and extract the PlatformList.cab
    $PlatformListUrl = 'https://hpia.hpcloud.hp.com/ref/platformList.cab'
    $DriversFolder = "$DriversFolder\$Make"
    $PlatformListCab = "$DriversFolder\platformList.cab"
    $PlatformListXml = "$DriversFolder\PlatformList.xml"

    if (-not (Test-Path -Path $DriversFolder)) {
        WriteLog "Creating Drivers folder: $DriversFolder"
        New-Item -Path $DriversFolder -ItemType Directory -Force | Out-Null
        WriteLog "Drivers folder created"
    }
    WriteLog "Downloading $PlatformListUrl to $PlatformListCab"
    Start-BitsTransferWithRetry -Source $PlatformListUrl -Destination $PlatformListCab
    WriteLog "Download complete"
    WriteLog "Expanding $PlatformListCab to $PlatformListXml"
    Invoke-Process -FilePath expand.exe -ArgumentList "$PlatformListCab $PlatformListXml" | Out-Null
    WriteLog "Expansion complete"

    # Parse the PlatformList.xml to find the SystemID based on the ProductName
    [xml]$PlatformListContent = Get-Content -Path $PlatformListXml
    $ProductNodes = $PlatformListContent.ImagePal.Platform | Where-Object { $_.ProductName.'#text' -match $Model }

    # Create a list of unique ProductName entries
    $ProductNames = @()
    foreach ($node in $ProductNodes) {
        foreach ($productName in $node.ProductName) {
            if ($productName.'#text' -match $Model) {
                $ProductNames += [PSCustomObject]@{
                    ProductName = $productName.'#text'
                    SystemID    = $node.SystemID
                    OSReleaseID = $node.OS.OSReleaseIdFileName -replace 'H', 'h'
                    IsWindows11 = $node.OS.IsWindows11 -contains 'true'
                }
            }
        }
    }

    if ($ProductNames.Count -gt 1) {
        Write-Output "More than one model found matching '$Model':"
        WriteLog "More than one model found matching '$Model':"
        $ProductNames | ForEach-Object -Begin { $i = 1 } -Process {
            if ($VerbosePreference -ne 'Continue') {
                Write-Output "$i. $($_.ProductName)"
            }
            WriteLog "$i. $($_.ProductName)"
            $i++
        }
        $selection = Read-Host "Please select the number corresponding to the correct model"
        WriteLog "User selected model number: $selection"
        if ($selection -match '^\d+$' -and [int]$selection -le $ProductNames.Count) {
            $SelectedProduct = $ProductNames[[int]$selection - 1]
            $ProductName = $SelectedProduct.ProductName
            WriteLog "Selected model: $ProductName"
            $SystemID = $SelectedProduct.SystemID
            WriteLog "SystemID: $SystemID"
            $ValidOSReleaseIDs = $SelectedProduct.OSReleaseID
            WriteLog "Valid OSReleaseIDs: $ValidOSReleaseIDs"
            $IsWindows11 = $SelectedProduct.IsWindows11
            WriteLog "IsWindows11 supported: $IsWindows11"
        }
        else {
            WriteLog "Invalid selection. Exiting."
            if ($VerbosePreference -ne 'Continue') {
                Write-Host "Invalid selection. Exiting."
            }
            exit
        }
    }
    elseif ($ProductNames.Count -eq 1) {
        $SelectedProduct = $ProductNames[0]
        $ProductName = $SelectedProduct.ProductName
        WriteLog "Selected model: $ProductName"
        $SystemID = $SelectedProduct.SystemID
        WriteLog "SystemID: $SystemID"
        $ValidOSReleaseIDs = $SelectedProduct.OSReleaseID
        WriteLog "OSReleaseID: $ValidOSReleaseIDs"
        $IsWindows11 = $SelectedProduct.IsWindows11
        WriteLog "IsWindows11: $IsWindows11"
    }
    else {
        WriteLog "No models found matching '$Model'. Exiting."
        if ($VerbosePreference -ne 'Continue') {
            Write-Host "No models found matching '$Model'. Exiting."
        }
        exit
    }

    if (-not $SystemID) {
        WriteLog "SystemID not found for model: $Model Exiting."
        if ($VerbosePreference -ne 'Continue') {
            Write-Host "SystemID not found for model: $Model Exiting."
        }
        exit
    }

    # Validate if WindowsRelease is 11 and there is no IsWindows11 element set to true
    if ($WindowsRelease -eq 11 -and -not $IsWindows11) {
        WriteLog "WindowsRelease is set to 11, but no drivers are available for this Windows release. Please set the -WindowsRelease parameter to 10, or provide your own drivers to the FFUDevelopment\Drivers folder."
        Write-Output "WindowsRelease is set to 11, but no drivers are available for this Windows release. Please set the -WindowsRelease parameter to 10, or provide your own drivers to the FFUDevelopment\Drivers folder."
        exit
    }

    # Validate WindowsVersion against OSReleaseID
    $OSReleaseIDs = $ValidOSReleaseIDs -split ' '
    $MatchingReleaseID = $OSReleaseIDs | Where-Object { $_ -eq "$WindowsVersion" }

    if (-not $MatchingReleaseID) {
        Write-Output "The specified WindowsVersion value '$WindowsVersion' is not valid for the selected model. Please select a valid OSReleaseID:"
        $OSReleaseIDs | ForEach-Object -Begin { $i = 1 } -Process {
            Write-Output "$i. $_"
            $i++
        }
        $selection = Read-Host "Please select the number corresponding to the correct OSReleaseID"
        WriteLog "User selected OSReleaseID number: $selection"
        if ($selection -match '^\d+$' -and [int]$selection -le $OSReleaseIDs.Count) {
            $WindowsVersion = $OSReleaseIDs[[int]$selection - 1]
            WriteLog "Selected OSReleaseID: $WindowsVersion"
        }
        else {
            WriteLog "Invalid selection. Exiting."
            exit
        }
    }

    # Modify WindowsArch for URL
    $Arch = $WindowsArch -replace "^x", ""

    # Construct the URL to download the driver XML cab for the model
    # The HPcloud reference site is case sensitve so we must convert the Windowsversion to lower 'h' first
    $WindowsVersionHP = $WindowsVersion -replace 'H', 'h'
    $ModelRelease = $SystemID + "_$Arch" + "_$WindowsRelease" + ".0.$WindowsVersionHP"
    $DriverCabUrl = "https://hpia.hpcloud.hp.com/ref/$SystemID/$ModelRelease.cab"
    $DriverCabFile = "$DriversFolder\$ModelRelease.cab"
    $DriverXmlFile = "$DriversFolder\$ModelRelease.xml"

    if (-not (Test-Url -Url $DriverCabUrl)) {
        WriteLog "HP Driver cab URL is not accessible: $DriverCabUrl Exiting"
        if ($VerbosePreference -ne 'Continue') {
            Write-Host "HP Driver cab URL is not accessible: $DriverCabUrl Exiting"
        }
        exit
    }

    # Download and extract the driver XML cab
    Writelog "Downloading HP Driver cab from $DriverCabUrl to $DriverCabFile"
    Start-BitsTransferWithRetry -Source $DriverCabUrl -Destination $DriverCabFile
    WriteLog "Expanding HP Driver cab to $DriverXmlFile"
    Invoke-Process -FilePath expand.exe -ArgumentList "$DriverCabFile $DriverXmlFile" | Out-Null

    # Parse the extracted XML file to download individual drivers
    [xml]$DriverXmlContent = Get-Content -Path $DriverXmlFile
    $baseUrl = "https://ftp.hp.com/pub/softpaq/sp"

    WriteLog "Downloading drivers for $ProductName"
    foreach ($update in $DriverXmlContent.ImagePal.Solutions.UpdateInfo) {
        if ($update.Category -notmatch '^Driver') {
            continue
        }
    
        $Name = $update.Name
        # Fix the name for drivers that contain illegal characters for folder name purposes
        $Name = $Name -replace '[\\\/\:\*\?\"\<\>\|]', '_'
        WriteLog "Downloading driver: $Name"
        $Category = $update.Category
        $Category = $Category -replace '[\\\/\:\*\?\"\<\>\|]', '_'
        $Version = $update.Version
        $Version = $Version -replace '[\\\/\:\*\?\"\<\>\|]', '_'
        $DriverUrl = "https://$($update.URL)"
        WriteLog "Driver URL: $DriverUrl"
        $DriverFileName = [System.IO.Path]::GetFileName($DriverUrl)
        $downloadFolder = "$DriversFolder\$ProductName\$Category"
        $DriverFilePath = Join-Path -Path $downloadFolder -ChildPath $DriverFileName

        if (Test-Path -Path $DriverFilePath) {
            WriteLog "Driver already downloaded: $DriverFilePath, skipping"
            continue
        }

        if (-not (Test-Path -Path $downloadFolder)) {
            WriteLog "Creating download folder: $downloadFolder"
            New-Item -Path $downloadFolder -ItemType Directory -Force | Out-Null
            WriteLog "Download folder created"
        }

        # Download the driver with retry
        WriteLog "Downloading driver to: $DriverFilePath"
        Mark-DownloadInProgress -FFUDevelopmentPath $FFUDevelopmentPath -TargetPath $DriverFilePath
        Start-BitsTransferWithRetry -Source $DriverUrl -Destination $DriverFilePath
        Clear-DownloadInProgress -FFUDevelopmentPath $FFUDevelopmentPath -TargetPath $DriverFilePath
        WriteLog 'Driver downloaded'

        # Make folder for extraction
        $extractFolder = "$downloadFolder\$Name\$Version\" + $DriverFileName.TrimEnd('.exe')
        Writelog "Creating extraction folder: $extractFolder"
        New-Item -Path $extractFolder -ItemType Directory -Force | Out-Null
        WriteLog 'Extraction folder created'
    
        # Extract the driver
        $arguments = "/s /e /f `"$extractFolder`""
        WriteLog "Extracting driver"
        Invoke-Process -FilePath $DriverFilePath -ArgumentList $arguments | Out-Null
        WriteLog "Driver extracted to: $extractFolder"

        # Delete the .exe driver file after extraction
        Remove-Item -Path $DriverFilePath -Force
        WriteLog "Driver installation file deleted: $DriverFilePath"
    }
    # Clean up the downloaded cab and xml files
    Remove-Item -Path $DriverCabFile, $DriverXmlFile, $PlatformListCab, $PlatformListXml -Force
    WriteLog "Driver cab and xml files deleted"
}
function Get-LenovoDrivers {
    param (
        [Parameter()]
        [string]$Model,
        [Parameter()]
        [ValidateSet("x64", "x86", "ARM64")]
        [string]$WindowsArch,
        [Parameter()]
        [ValidateSet(10, 11)]
        [int]$WindowsRelease
    )

    function Get-LenovoPSREF {
        param (
            [string]$ModelName
        )

        $normalizedModelName = $ModelName.Trim()
        $searchType = if ($normalizedModelName.Length -eq 4) { 'MT' } elseif ($normalizedModelName.Length -eq 10) { 'Model' } else { 'Normal' }
        $cacheBuster = ([System.Random]::new().NextDouble()).ToString('0.################', [Globalization.CultureInfo]::InvariantCulture)
        $url = "https://psref.lenovo.com/api/search/DefinitionFilterAndSearch/Suggest?kw=$([uri]::EscapeDataString($normalizedModelName))&limit=6&IsPreviewProduct=true&SearchType=$searchType&t=$cacheBuster"
        WriteLog "Querying Lenovo PSREF API for model: $normalizedModelName"
        $requestHeaders = Get-LenovoPSREFRequestHeaders -Method GET -Uri $url -Headers $Headers -UserAgent $UserAgent
        $originalVerbosePreference = $VerbosePreference
        try {
            $VerbosePreference = 'SilentlyContinue'
            $jsonResponse = Invoke-RestMethod -Uri $url -Headers $requestHeaders -UserAgent $UserAgent -ErrorAction Stop
        }
        finally {
            $VerbosePreference = $originalVerbosePreference
        }
        WriteLog "Lenovo PSREF API query complete."

        $products = @()
        foreach ($item in $jsonResponse.data) {
            if (-not [string]::IsNullOrEmpty($item.MachineType) -and -not [string]::IsNullOrEmpty($item.ProductName)) {
                $productName = $item.ProductName
                $machineTypes = $item.MachineType -split '\s*/\s*'

                foreach ($machineType in $machineTypes) {
                    if ($machineType -eq $normalizedModelName) {
                        WriteLog "Model name entered is a matching machine type"
                        $products = @()
                        $products += [pscustomobject]@{
                            ProductName = $productName
                            MachineType = $machineType
                        }
                        WriteLog "Product Name: $productName Machine Type: $machineType"
                        return $products
                    }
                    $products += [pscustomobject]@{
                        ProductName = $productName
                        MachineType = $machineType
                    }
                }
            }
        }

        return , $products
    }
    
    # Parse the Lenovo PSREF page for the model
    $machineTypes = Get-LenovoPSREF -ModelName $Model
    if ($machineTypes.ProductName.Count -eq 0) {
        WriteLog "No machine types found for model: $Model"
        WriteLog "Enter a valid model or machine type in the -model parameter"
        exit
    }
    elseif ($machineTypes.ProductName.Count -eq 1) {
        $machineType = $machineTypes[0].MachineType
        $model = $machineTypes[0].ProductName
    }
    else {
        if ($VerbosePreference -ne 'Continue') {
            Write-Output "Multiple machine types found for model: $Model"
        }
        WriteLog "Multiple machine types found for model: $Model"
        for ($i = 0; $i -lt $machineTypes.ProductName.Count; $i++) {
            if ($VerbosePreference -ne 'Continue') {
                Write-Output "$($i + 1). $($machineTypes[$i].ProductName) ($($machineTypes[$i].MachineType))"
            }
            WriteLog "$($i + 1). $($machineTypes[$i].ProductName) ($($machineTypes[$i].MachineType))"
        }
        $selection = Read-Host "Enter the number of the model you want to select"
        $machineType = $machineTypes[$selection - 1].MachineType
        WriteLog "Selected machine type: $machineType"
        $model = $machineTypes[$selection - 1].ProductName
        WriteLog "Selected model: $model"
    }
    

    # Construct the catalog URL based on Windows release and machine type
    $ModelRelease = $machineType + "_Win" + $WindowsRelease
    $CatalogUrl = "https://download.lenovo.com/catalog/$ModelRelease.xml"
    WriteLog "Lenovo Driver catalog URL: $CatalogUrl"

    if (-not (Test-Url -Url $catalogUrl)) {
        Write-Error "Lenovo Driver catalog URL is not accessible: $catalogUrl"
        WriteLog "Lenovo Driver catalog URL is not accessible: $catalogUrl"
        exit
    }

    # Create the folder structure for the Lenovo drivers
    $driversFolder = "$DriversFolder\$Make"
    if (-not (Test-Path -Path $DriversFolder)) {
        WriteLog "Creating Drivers folder: $DriversFolder"
        New-Item -Path $DriversFolder -ItemType Directory -Force | Out-Null
        WriteLog "Drivers folder created"
    }

    # Download and parse the Lenovo catalog XML
    $LenovoCatalogXML = "$DriversFolder\$ModelRelease.xml"
    WriteLog "Downloading $catalogUrl to $LenovoCatalogXML"
    Start-BitsTransferWithRetry -Source $catalogUrl -Destination $LenovoCatalogXML
    WriteLog "Download Complete"
    $xmlContent = [xml](Get-Content -Path $LenovoCatalogXML)

    WriteLog "Parsing Lenovo catalog XML"
    # Process each package in the catalog
    foreach ($package in $xmlContent.packages.package) {
        $packageUrl = $package.location
        $category = $package.category

        #If category starts with BIOS, skip the package
        if ($category -like 'BIOS*') {
            continue
        }

        #If category name is 'Motherboard Devices Backplanes core chipset onboard video PCIe switches', truncate to 'Motherboard Devices' to shorten path
        if ($category -eq 'Motherboard Devices Backplanes core chipset onboard video PCIe switches') {
            $category = 'Motherboard Devices'
        }

        $packageName = [System.IO.Path]::GetFileName($packageUrl)
        #Remove the filename from the $packageURL
        $baseURL = $packageUrl -replace $packageName, "" 

        # Download the package XML
        $packageXMLPath = "$DriversFolder\$packageName"
        WriteLog "Downloading $category package XML $packageUrl to $packageXMLPath"
        If ((Start-BitsTransferWithRetry -Source $packageUrl -Destination $packageXMLPath) -eq $false) {
            Write-Output "Failed to download $category package XML: $packageXMLPath"
            WriteLog "Failed to download $category package XML: $packageXMLPath"
            continue
        }

        # Load the package XML content
        $packageXmlContent = [xml](Get-Content -Path $packageXMLPath)
        $packageType = $packageXmlContent.Package.PackageType.type
        $packageTitle = $packageXmlContent.Package.title.InnerText

        # Fix the name for drivers that contain illegal characters for folder name purposes
        $packageTitle = $packageTitle -replace '[\\\/\:\*\?\"\<\>\|]', '_'

        # If ' - ' is in the package title, truncate the title to the first part of the string.
        $packageTitle = $packageTitle -replace ' - .*', ''

        #Check if packagetype = 2. If packagetype is not 2, skip the package. $packageType is a System.Xml.XmlElement.
        #This filters out Firmware, BIOS, and other non-INF drivers
        if ($packageType -ne 2) {
            Remove-Item -Path $packageXMLPath -Force
            continue
        }

        # Extract the driver file name and the extract command
        $driverFileName = $packageXmlContent.Package.Files.Installer.File.Name
        $extractCommand = $packageXmlContent.Package.ExtractCommand

        #if extract command is empty/missing, skip the package
        if (!($extractCommand)) {
            Remove-Item -Path $packageXMLPath -Force
            continue
        }

        # Create the download URL and folder structure
        $driverUrl = $baseUrl + $driverFileName
        $downloadFolder = "$DriversFolder\$Model\$Category\$packageTitle"
        $driverFilePath = Join-Path -Path $downloadFolder -ChildPath $driverFileName

        # Check if file has already been downloaded
        if (Test-Path -Path $driverFilePath) {
            Write-Output "Driver already downloaded: $driverFilePath skipping"
            WriteLog "Driver already downloaded: $driverFilePath skipping"
            continue
        }

        if (-not (Test-Path -Path $downloadFolder)) {
            WriteLog "Creating download folder: $downloadFolder"
            New-Item -Path $downloadFolder -ItemType Directory -Force | Out-Null
            WriteLog "Download folder created"
        }

        # Download the driver with retry
        WriteLog "Downloading driver: $driverUrl to $driverFilePath"
        Mark-DownloadInProgress -FFUDevelopmentPath $FFUDevelopmentPath -TargetPath $driverFilePath
        Start-BitsTransferWithRetry -Source $driverUrl -Destination $driverFilePath
        Clear-DownloadInProgress -FFUDevelopmentPath $FFUDevelopmentPath -TargetPath $driverFilePath
        WriteLog "Driver downloaded"

        # Make folder for extraction
        $extractFolder = $downloadFolder + "\" + $driverFileName.TrimEnd($driverFileName[-4..-1])
        WriteLog "Creating extract folder: $extractFolder"
        New-Item -Path $extractFolder -ItemType Directory -Force | Out-Null
        WriteLog "Extract folder created"

        # Modify the extract command
        $modifiedExtractCommand = $extractCommand -replace '%PACKAGEPATH%', "`"$extractFolder`""

        # Extract the driver
        # Start-Process -FilePath $driverFilePath -ArgumentList $modifiedExtractCommand -Wait -NoNewWindow
        WriteLog "Extracting driver: $driverFilePath to $extractFolder"
        Invoke-Process -FilePath $driverFilePath -ArgumentList $modifiedExtractCommand | Out-Null
        WriteLog "Driver extracted"

        # Delete the .exe driver file after extraction
        WriteLog "Deleting driver installation file: $driverFilePath"
        Remove-Item -Path $driverFilePath -Force
        WriteLog "Driver installation file deleted: $driverFilePath"

        # Delete the package XML file after extraction
        WriteLog "Deleting package XML file: $packageXMLPath"
        Remove-Item -Path $packageXMLPath -Force
        WriteLog "Package XML file deleted"
    }

    #Delete the catalog XML file after processing
    WriteLog "Deleting catalog XML file: $LenovoCatalogXML"
    Remove-Item -Path $LenovoCatalogXML -Force
    WriteLog "Catalog XML file deleted"
}

function Get-DellDrivers {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Model,
        [Parameter(Mandatory = $true)]
        [ValidateSet('x64', 'x86', 'ARM64')]
        [string]$WindowsArch,
        [Parameter(Mandatory = $true)]
        [int]$WindowsRelease
    )

    if (-not (Test-Path -Path $DriversFolder)) {
        New-Item -Path $DriversFolder -ItemType Directory -Force | Out-Null
    }
    $DriversFolder = Join-Path $DriversFolder $Make
    if (-not (Test-Path $DriversFolder)) {
        New-Item -Path $DriversFolder -ItemType Directory -Force | Out-Null
    }

    # Client pathway (<=11): use CatalogIndexPC + per-model cab.
    if ($WindowsRelease -le 11) {
        $indexXml = Get-DellCatalogIndex -DriversFolder (Split-Path $DriversFolder -Parent)
        $allModels = Get-DellClientModels -CatalogIndexXmlPath $indexXml
        $target = $allModels | Where-Object { $_.ModelDisplay -eq $Model }
        if (-not $target) { throw "Requested Dell model '$Model' not found in index." }

        $cabUrl = $target.CabUrl
        if ([string]::IsNullOrWhiteSpace($cabUrl)) {
            WriteLog "CabUrl missing for '$($target.Model)'; resolving via CatalogIndexPC."
            $resolved = Resolve-DellCabUrlFromModel -DriversFolder $DriversFolder -ModelDisplay $target.Model
            if ($null -eq $resolved -or [string]::IsNullOrWhiteSpace($resolved.CabUrl)) {
                throw "Unable to resolve CabUrl for $($target.Model) from CatalogIndexPC."
            }
            $cabUrl = $resolved.CabUrl
            $target.CabUrl = $cabUrl
        }
        $modelCabName = [IO.Path]::GetFileName($cabUrl)
        $modelCabPath = Join-Path $DriversFolder $modelCabName
        $modelXmlPath = Join-Path $DriversFolder ($modelCabName -replace '\.cab$', '.xml')

        if (Test-Path $modelCabPath) { Remove-Item $modelCabPath -Force -ErrorAction SilentlyContinue }
        if (Test-Path $modelXmlPath) { Remove-Item $modelXmlPath -Force -ErrorAction SilentlyContinue }

        Start-BitsTransferWithRetry -Source $cabUrl -Destination $modelCabPath
        Invoke-Process -FilePath Expand.exe -ArgumentList """$modelCabPath"" ""$modelXmlPath""" | Out-Null
        Remove-Item $modelCabPath -Force -ErrorAction SilentlyContinue
        if (-not (Test-Path $modelXmlPath)) { throw "Failed to extract model cab XML: $modelXmlPath" }

        $pkgs = Get-DellLatestDriverPackages -ModelXmlPath $modelXmlPath -WindowsArch $WindowsArch -WindowsRelease $WindowsRelease
        if (-not $pkgs) {
            WriteLog "No drivers found for '$Model'."
            return
        }

        foreach ($pkg in $pkgs) {
            $categorySafe = ($pkg.Category -replace '[\\\/\:\*\?\"\<\>\| ]', '_')
            $downloadFolder = Join-Path $DriversFolder (Join-Path $Model $categorySafe)
            if (-not (Test-Path $downloadFolder)) { New-Item -Path $downloadFolder -ItemType Directory -Force | Out-Null }
            $driverFilePath = Join-Path $downloadFolder $pkg.DriverFileName
            $extractFolder = Join-Path $downloadFolder ($pkg.DriverFileName.TrimEnd($pkg.DriverFileName[-4..-1]))

            if (Test-Path $extractFolder) {
                $sz = (Get-ChildItem -Path $extractFolder -Recurse -Exclude *.log | Measure-Object -Property Length -Sum -ErrorAction SilentlyContinue).Sum
                if ($sz -gt 1KB) { continue }
            }
            if (-not (Test-Path $driverFilePath)) {
                try { Start-BitsTransferWithRetry -Source $pkg.DownloadUrl -Destination $driverFilePath }
                catch { WriteLog "Download failed: $($pkg.DownloadUrl) $($_.Exception.Message)"; continue }
            }

            if (-not (Test-Path $extractFolder)) { New-Item -Path $extractFolder -ItemType Directory -Force | Out-Null }

            $arg1 = "/s /e=`"$extractFolder`" /l=`"$extractFolder\log.log`""
            $arg2 = "/s /drivers=`"$extractFolder`" /l=`"$extractFolder\log.log`""
            $ok = $false
            try {
                Invoke-Process -FilePath $driverFilePath -ArgumentList $arg1 | Out-Null
                $sz = (Get-ChildItem -Path $extractFolder -Recurse -Exclude *.log | Measure-Object -Property Length -Sum -ErrorAction SilentlyContinue).Sum
                if ($sz -gt 1KB) { $ok = $true }
                if (-not $ok) {
                    Remove-Item $extractFolder -Recurse -Force -ErrorAction SilentlyContinue
                    New-Item -Path $extractFolder -ItemType Directory -Force | Out-Null
                    Invoke-Process -FilePath $driverFilePath -ArgumentList $arg2 | Out-Null
                    $sz = (Get-ChildItem -Path $extractFolder -Recurse -Exclude *.log | Measure-Object -Property Length -Sum -ErrorAction SilentlyContinue).Sum
                    if ($sz -gt 1KB) { $ok = $true }
                }
            }
            catch {
                WriteLog "Extraction error: $($_.Exception.Message)"
            }
            if ($ok) { Remove-Item $driverFilePath -Force -ErrorAction SilentlyContinue }
        }
        return
    }

    # Server pathway (unchanged legacy)
    $catalogUrl = "https://downloads.dell.com/catalog/Catalog.cab"
    $DellCabFile = Join-Path $DriversFolder 'Catalog.cab'
    $DellCatalogXML = Join-Path $DriversFolder 'Catalog.xml'

    Start-BitsTransferWithRetry -Source $catalogUrl -Destination $DellCabFile
    Invoke-Process -FilePath Expand.exe -ArgumentList "$DellCabFile $DellCatalogXML" | Out-Null

    $xmlContent = [xml](Get-Content -Path $DellCatalogXML)
    $baseLocation = "https://" + $xmlContent.manifest.baseLocation + "/"
    $latestDrivers = @{}
    $softwareComponents = $xmlContent.Manifest.SoftwareComponent | Where-Object { $_.ComponentType.value -eq 'DRVR' }

    foreach ($component in $softwareComponents) {
        $models = $component.SupportedSystems.Brand.Model
        foreach ($item in $models) {
            if ($item.Display.'#cdata-section' -match $Model) {
                $validOS = $component.SupportedOperatingSystems.OperatingSystem | Where-Object { $_.osArch -eq $WindowsArch }
                if (-not $validOS) { continue }
                $driverPath = $component.path
                $downloadUrl = $baseLocation + $driverPath
                $driverFileName = [System.IO.Path]::GetFileName($driverPath)
                $name = $component.Name.Display.'#cdata-section' -replace '[\\\/\:\*\?\"\<\>\| ]', '_' -replace '[\,]', '-'
                $category = $component.Category.Display.'#cdata-section' -replace '[\\\/\:\*\?\"\<\>\| ]', '_'
                $version = [version]$component.vendorVersion
                $namePrefix = ($name -split '-')[0]
                if (-not $latestDrivers[$category]) { $latestDrivers[$category] = @{} }
                if (-not $latestDrivers[$category][$namePrefix] -or $latestDrivers[$category][$namePrefix].Version -lt $version) {
                    $latestDrivers[$category][$namePrefix] = [pscustomobject]@{
                        Name = $name; DownloadUrl = $downloadUrl; DriverFileName = $driverFileName; Version = $version; Category = $category
                    }
                }
            }
        }
    }

    foreach ($category in $latestDrivers.Keys) {
        foreach ($driver in $latestDrivers[$category].Values) {
            $downloadFolder = Join-Path $DriversFolder (Join-Path $Model $driver.Category)
            if (-not (Test-Path $downloadFolder)) { New-Item -Path $downloadFolder -ItemType Directory -Force | Out-Null }
            $driverFilePath = Join-Path $downloadFolder $driver.DriverFileName
            if (Test-Path $driverFilePath) { continue }
            Start-BitsTransferWithRetry -Source $driver.DownloadUrl -Destination $driverFilePath
            $extractFolder = Join-Path $downloadFolder ($driver.DriverFileName.TrimEnd($driver.DriverFileName[-4..-1]))
            $arguments = "/s /drivers=`"$extractFolder`""
            try { Invoke-Process -FilePath $driverFilePath -ArgumentList $arguments | Out-Null } catch {}
            Remove-Item $driverFilePath -Force -ErrorAction SilentlyContinue
        }
    }
}
function Get-ADKURL {
    param (
        [ValidateSet("Windows ADK", "WinPE add-on")]
        [string]$ADKOption
    )

    # Define base pattern for URL scraping
    $basePattern = '<li><a href="(https://[^"]+)" data-linktype="external">Download the '

    # Define specific URL patterns based on ADK options
    $ADKUrlPattern = @{
        "Windows ADK"  = $basePattern + "Windows ADK"
        "WinPE add-on" = $basePattern + "Windows PE add-on for the Windows ADK"
    }[$ADKOption]

    try {
        # Retrieve content of Microsoft documentation page
        $OriginalVerbosePreference = $VerbosePreference
        $VerbosePreference = 'SilentlyContinue'
        $ADKWebPage = Invoke-RestMethod "https://learn.microsoft.com/en-us/windows-hardware/get-started/adk-install" -Headers $Headers -UserAgent $UserAgent
        $VerbosePreference = $OriginalVerbosePreference
        
        # Extract download URL based on specified pattern
        $ADKMatch = [regex]::Match($ADKWebPage, $ADKUrlPattern)

        if (-not $ADKMatch.Success) {
            WriteLog "Failed to retrieve ADK download URL. Pattern match failed."
            return
        }

        # Extract FWlink from the matched pattern
        $ADKFWLink = $ADKMatch.Groups[1].Value

        if ($null -eq $ADKFWLink) {
            WriteLog "FWLink for $ADKOption not found."
            return
        }

        # Let Invoke-WebRequest handle the redirect and get the final URL.
        try {
            $OriginalVerbosePreference = $VerbosePreference
            $VerbosePreference = 'SilentlyContinue'
            # Allow one redirection to get the final URL from the fwlink
            $response = Invoke-WebRequest -Uri $ADKFWLink -Method Head -MaximumRedirection 1 -Headers $Headers -UserAgent $UserAgent
            $VerbosePreference = $OriginalVerbosePreference
            
            # The final URL after redirection is in the ResponseUri property of the BaseResponse's RequestMessage.
            $ADKUrl = $response.BaseResponse.RequestMessage.RequestUri.AbsoluteUri
            
            if ($null -eq $ADKUrl) {
                WriteLog "Could not determine final ADK download URL after redirection."
                return $null
            }
            
            WriteLog "Resolved ADK download URL to: $ADKUrl"
            return $ADKUrl
        }
        catch {
            WriteLog "An error occurred while resolving the ADK FWLink: $($_.Exception.Message)"
            throw
        }
    }
    catch {
        WriteLog $_
        Write-Error "Error occurred while retrieving ADK download URL"
        throw $_
    }
}
function Install-ADK {
    param (
        [ValidateSet("Windows ADK", "WinPE add-on")]
        [string]$ADKOption
    )

    try {
        $ADKUrl = Get-ADKURL -ADKOption $ADKOption
        
        if ($null -eq $ADKUrl) {
            throw "Failed to retrieve URL for $ADKOption. Please manually install it."
        }

        # Select the installer based on the ADK option specified
        $installer = @{
            "Windows ADK"  = "adksetup.exe"
            "WinPE add-on" = "adkwinpesetup.exe"
        }[$ADKOption]

        # Select the feature based on the ADK option specified
        $feature = @{
            "Windows ADK"  = "OptionId.DeploymentTools"
            "WinPE add-on" = "OptionId.WindowsPreinstallationEnvironment"
        }[$ADKOption]

        $installerLocation = Join-Path $env:TEMP $installer

        WriteLog "Downloading $ADKOption from $ADKUrl to $installerLocation"
        Start-BitsTransferWithRetry -Source $ADKUrl -Destination $installerLocation -ErrorAction Stop
        WriteLog "$ADKOption downloaded to $installerLocation"
        
        WriteLog "Installing $ADKOption with $feature enabled"
        Invoke-Process $installerLocation "/quiet /installpath ""%ProgramFiles(x86)%\Windows Kits\10"" /features $feature" | Out-Null
        
        WriteLog "$ADKOption installation completed."
        WriteLog "Removing $installer from $installerLocation"
        # Clean up downloaded installation file
        Remove-Item -Path $installerLocation -Force -ErrorAction SilentlyContinue
    }
    catch {
        WriteLog $_
        Write-Error "Error occurred while installing $ADKOption. Please manually install it."
        throw $_
    }
}
function Get-InstalledProgramRegKey {
    param (
        [ValidateSet("Windows ADK", "WinPE add-on")]
        [string]$ADKOption
    )

    # Match ADK entries using installer executable names in BundleCachePath
    $bundleExecutableName = switch ($ADKOption) {
        "Windows ADK" { "adksetup.exe" }
        "WinPE add-on" { "adkwinpesetup.exe" }
    }

    # Check both uninstall roots to support different install contexts
    $uninstallRegPaths = @(
        "HKLM:\SOFTWARE\WOW6432Node\Microsoft\Windows\CurrentVersion\Uninstall",
        "HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Uninstall"
    )

    foreach ($uninstallRegPath in $uninstallRegPaths) {
        if (-not (Test-Path -Path $uninstallRegPath)) {
            continue
        }

        $uninstallRegKeys = Get-ChildItem -Path $uninstallRegPath -Recurse -ErrorAction SilentlyContinue

        foreach ($regKey in $uninstallRegKeys) {
            try {
                $bundleCachePath = $regKey.GetValue("BundleCachePath")
                if ($null -ne $bundleCachePath) {
                    $bundleFileName = Split-Path -Path $bundleCachePath -Leaf
                    if ($bundleFileName -ieq $bundleExecutableName) {
                        return $regKey
                    }
                }
            }
            catch {
                WriteLog $_
                throw "Error retrieving installed program info for $ADKOption : $_"
            }
        }
    }

    return $null
}

function Uninstall-ADK {
    param (
        [ValidateSet("Windows ADK", "WinPE add-on")]
        [string]$ADKOption
    )

    try {
        # Find ADK entry by bundle executable path instead of display name
        $adkRegKey = Get-InstalledProgramRegKey -ADKOption $ADKOption

        if (-not $adkRegKey) {
            WriteLog "$ADKOption is not installed."
            return
        }

        $adkBundleCachePath = $adkRegKey.GetValue("BundleCachePath")
        WriteLog "Uninstalling $ADKOption..."
        Invoke-Process $adkBundleCachePath "/uninstall /quiet" | Out-Null
        WriteLog "$ADKOption uninstalled successfully."
    }
    catch {
        WriteLog $_
        Write-Error "Error occurred while uninstalling $ADKOption. Please manually uninstall it."
        throw $_
    }
}

function Confirm-ADKVersionIsLatest {
    param (
        [ValidateSet("Windows ADK", "WinPE add-on")]
        [string]$ADKOption
    )

    try {
        # Find ADK entry by bundle executable path instead of display name
        $adkRegKey = Get-InstalledProgramRegKey -ADKOption $ADKOption

        if (-not $adkRegKey) {
            return $false
        }

        $installedADKVersion = $adkRegKey.GetValue("DisplayVersion")

        # Retrieve content of Microsoft documentation page
        $adkWebPage = Invoke-RestMethod "https://learn.microsoft.com/en-us/windows-hardware/get-started/adk-install" -Headers $Headers -UserAgent $UserAgent
        # Specify regex pattern for ADK version
        $adkVersionPattern = 'ADK\s+(\d+(\.\d+)+)'
        # Check for regex pattern match
        $adkVersionMatch = [regex]::Match($adkWebPage, $adkVersionPattern)

        if (-not $adkVersionMatch.Success) {
            WriteLog "Failed to retrieve latest ADK version from web page."
            return $false
        }

        # Extract ADK version from the matched pattern
        $latestADKVersion = $adkVersionMatch.Groups[1].Value

        if ($installedADKVersion -eq $latestADKVersion) {
            WriteLog "Installed $ADKOption version $installedADKVersion is the latest."
            return $true
        }
        else {
            WriteLog "Installed $ADKOption version $installedADKVersion is not the latest ($latestADKVersion)"
            return $false
        }
    }
    catch {
        WriteLog "An error occurred while confirming the ADK version: $_"
        return $false
    }
}

function Get-ADK {
    # Check if latest ADK and WinPE add-on are installed
    if ($UpdateADK) {
        WriteLog "Checking if latest ADK and WinPE add-on are installed"
        $latestADKInstalled = Confirm-ADKVersionIsLatest -ADKOption "Windows ADK"
        $latestWinPEInstalled = Confirm-ADKVersionIsLatest -ADKOption "WinPE add-on"

        # Uninstall older versions and install latest versions if necessary
        if (-not $latestADKInstalled) {
            Uninstall-ADK -ADKOption "Windows ADK"
            Install-ADK -ADKOption "Windows ADK"
        }

        if (-not $latestWinPEInstalled) {
            Uninstall-ADK -ADKOption "WinPE add-on"
            Install-ADK -ADKOption "WinPE add-on"
        }
    }

    # Define registry path
    $adkPathKey = "HKLM:\SOFTWARE\WOW6432Node\Microsoft\Windows Kits\Installed Roots"
    $adkPathName = "KitsRoot10"

    # Check if ADK installation path exists in registry
    $adkPathNameExists = (Get-ItemProperty -Path $adkPathKey -Name $adkPathName -ErrorAction SilentlyContinue)

    if ($adkPathNameExists) {
        # Get the ADK installation path
        WriteLog 'Get ADK Path'
        $adkPath = (Get-ItemProperty -Path $adkPathKey -Name $adkPathName).$adkPathName
        WriteLog "ADK located at $adkPath"
    }
    else {
        throw "Windows ADK installation path could not be found."
    }

    # If ADK is installed, validate Deployment Tools by checking DandISetEnv.bat
    $deploymentToolsEnvPath = Join-Path $adkPath "Assessment and Deployment Kit\Deployment Tools\DandISetEnv.bat"

    if (-not (Test-Path -Path $deploymentToolsEnvPath)) {
        WriteLog "ADK is installed, but the Windows Deployment Tools feature is not installed."
        $adkRegKey = Get-InstalledProgramRegKey -ADKOption "Windows ADK"

        if ($null -eq $adkRegKey) {
            throw "Failed to locate the installed ADK package to install Windows Deployment Tools."
        }

        $adkBundleCachePath = $adkRegKey.GetValue("BundleCachePath")
        if ($adkBundleCachePath) {
            WriteLog "Installing Windows Deployment Tools..."
            $adkInstallPath = $adkPath.TrimEnd('\')
            Invoke-Process $adkBundleCachePath "/quiet /installpath ""$adkInstallPath"" /features OptionId.DeploymentTools" | Out-Null
            WriteLog "Windows Deployment Tools installed successfully."
        }
        else {
            throw "Failed to retrieve path to adksetup.exe to install the Windows Deployment Tools. Please manually install it."
        }
    }
    return $adkPath
}
function Get-ProductsCab {
    param(
        [Parameter(Mandatory = $true)]
        [string]$OutFile,
        [Parameter(Mandatory = $true)]
        [ValidateSet('x64', 'arm64')]
        [string]$Architecture,
        [Parameter(Mandatory = $true)]
        [string]$BuildVersion
    )

    # Always request the amd64 products catalog because it contains both x64 and arm64 media entries.
    $productsParam = "PN=Windows.Products.Cab.amd64&V=$BuildVersion"
    $deviceAttributes = "DUScan=1;OSVersion=10.0.26100.1"

    $bodyObj = [ordered]@{
        Products         = $productsParam
        DeviceAttributes = $deviceAttributes
    }
    $bodyJson = $bodyObj | ConvertTo-Json -Compress

    $searchUri = 'https://fe3.delivery.mp.microsoft.com/UpdateMetadataService/updates/search/v1/bydeviceinfo'

    WriteLog "Requesting products.cab location from Windows Update service..."
    try {
        $searchResponse = Invoke-RestMethod -Uri $searchUri -Method Post -ContentType 'application/json' -Headers @{ Accept = '*/*' } -Body $bodyJson
    }
    catch {
        WriteLog "Failed to retrieve products.cab metadata: $($_.Exception.Message)"
        throw
    }

    if ($searchResponse -is [System.Array]) { $searchResponse = $searchResponse[0] }
    if (-not $searchResponse.FileLocations) { throw "Search response did not include FileLocations." }

    $fileRec = $searchResponse.FileLocations | Where-Object { $_.FileName -eq 'products.cab' } | Select-Object -First 1
    if (-not $fileRec) { throw "products.cab entry not found in FileLocations." }

    $downloadUrl = $fileRec.Url
    $serverDigestB64 = $fileRec.Digest
    $serverSize = [int64]$fileRec.Size
    $updateId = $searchResponse.UpdateIds[0]

    try {
        $metaUri = "https://fe3.delivery.mp.microsoft.com/UpdateMetadataService/updates/v1/$updateId"
        $meta = Invoke-RestMethod -Uri $metaUri -Method Get -Headers @{ Accept = '*/*' }
        if ($meta.LocalizedProperties.Count -gt 0) {
            $title = $meta.LocalizedProperties[0].Title
            WriteLog "Resolved update: $title"
        }
        else {
            WriteLog "Resolved update id: $updateId"
        }
    }
    catch {
        WriteLog "Resolved update id: $updateId"
    }

    $destDir = Split-Path -Path $OutFile -Parent
    if ($destDir -and -not (Test-Path $destDir)) {
        [void](New-Item -ItemType Directory -Path $destDir)
    }

    WriteLog "Downloading products.cab to $OutFile ..."
    $downloadHeaders = @{ Accept = '*/*' }
    Invoke-WebRequest -Uri $downloadUrl -OutFile $OutFile -Headers $downloadHeaders -UserAgent $UserAgent

    $actualSize = (Get-Item $OutFile).Length
    if ($actualSize -ne $serverSize) {
        throw "Size check failed. Expected $serverSize bytes. Got $actualSize bytes."
    }

    $sha256 = [System.Security.Cryptography.SHA256]::Create()
    $fs = [System.IO.File]::OpenRead($OutFile)
    try {
        $hashBytes = $sha256.ComputeHash($fs)
    }
    finally {
        $fs.Dispose()
    }
    $actualDigestB64 = [Convert]::ToBase64String($hashBytes)

    if ($actualDigestB64 -ne $serverDigestB64) {
        throw "Digest check failed. Expected $serverDigestB64. Got $actualDigestB64."
    }

    WriteLog "products.cab downloaded and verified successfully."
    return $OutFile
}

function Get-WindowsESDMetadata {
    param(
        [Parameter(Mandatory = $false)]
        [ValidateSet(10, 11)]
        [int]$WindowsRelease,

        [Parameter(Mandatory = $false)]
        [ValidateSet('x86', 'x64', 'ARM64')]
        [string]$WindowsArch,

        [Parameter(Mandatory = $false)]
        [string]$WindowsLang,

        [Parameter(Mandatory = $false)]
        [ValidateSet('consumer', 'business')]
        [string]$MediaType
    )
    WriteLog "Resolving Windows $WindowsRelease ESD metadata"
    $cabFilePath = Join-Path $PSScriptRoot "tempCabFile.cab"
    $xmlFilePath = Join-Path $PSScriptRoot "products.xml"
    $esdMetadata = $null
    $OriginalVerbosePreference = $VerbosePreference
    $VerbosePreference = 'SilentlyContinue'
    try {
        if ($WindowsRelease -eq 10) {
            WriteLog "Downloading Cab file"
            Invoke-WebRequest -Uri 'https://go.microsoft.com/fwlink/?LinkId=841361' -OutFile $cabFilePath -Headers $Headers -UserAgent $UserAgent
        }
        elseif ($WindowsRelease -eq 11) {
            WriteLog "Downloading Cab file"
            $buildVersionMap = @{
                '22H2' = '22621.0.0.0'
                '23H2' = '22631.0.0.0'
                '24H2' = '26100.0.0.0'
                '25H2' = '26100.0.0.0'
            }
            $normalizedVersion = $WindowsVersion.ToUpper()
            if ($buildVersionMap.ContainsKey($normalizedVersion)) {
                $buildVersion = $buildVersionMap[$normalizedVersion]
            }
            else {
                WriteLog "No explicit build mapping found for Windows 11 version '$WindowsVersion'. Defaulting products.cab build token to 26100.0.0.0."
                $buildVersion = '26100.0.0.0'
            }

            $cabArchitecture = if ($WindowsArch -eq 'ARM64') { 'arm64' } else { 'x64' }
            Get-ProductsCab -OutFile $cabFilePath -Architecture $cabArchitecture -BuildVersion $buildVersion | Out-Null
        }
        else {
            throw "Downloading Windows $WindowsRelease is not supported. Please use the -ISOPath parameter to specify the path to the Windows $WindowsRelease ISO file."
        }
        WriteLog "products.cab download succeeded"
    }
    finally {
        $VerbosePreference = $OriginalVerbosePreference
    }

    WriteLog "Extracting Products XML from cab"
    Invoke-Process Expand "-F:*.xml $cabFilePath $xmlFilePath" | Out-Null
    WriteLog "Products XML extracted"

    [xml]$xmlContent = Get-Content -Path $xmlFilePath
    $clientType = if ($MediaType -eq 'consumer') { 'CLIENTCONSUMER' } else { 'CLIENTBUSINESS' }

    foreach ($file in $xmlContent.MCT.Catalogs.Catalog.PublishedMedia.Files.File) {
        if ($file.Architecture -eq $WindowsArch -and $file.LanguageCode -eq $WindowsLang -and $file.FilePath -like "*$clientType*") {
            $fileName = Split-Path $file.FilePath -Leaf
            $esdFilePath = Join-Path $PSScriptRoot $fileName
            $esdVersion = $null
            if ($file.FileName -match '^([0-9]+\.[0-9]+)') {
                $esdVersion = $matches[1]
            }

            $esdMetadata = [pscustomobject]@{
                FileUrl   = $file.FilePath
                FileName  = $fileName
                LocalPath = $esdFilePath
                Version   = $esdVersion
            }
            break
        }
    }

    if ($esdMetadata) {
        WriteLog "Resolved ESD metadata: $($esdMetadata.FileName) (Version: $($esdMetadata.Version))"
    }
    else {
        WriteLog "No matching ESD entry found in products.xml."
    }

    WriteLog "Cleaning up temporary cab and xml files"
    Remove-Item -Path $cabFilePath -Force -ErrorAction SilentlyContinue
    Remove-Item -Path $xmlFilePath -Force -ErrorAction SilentlyContinue

    return $esdMetadata
}

function Get-WindowsESD {
    param(
        [Parameter(Mandatory = $false)]
        [ValidateSet(10, 11)]
        [int]$WindowsRelease,

        [Parameter(Mandatory = $false)]
        [ValidateSet('x86', 'x64', 'ARM64')]
        [string]$WindowsArch,

        [Parameter(Mandatory = $false)]
        [string]$WindowsLang,

        [Parameter(Mandatory = $false)]
        [ValidateSet('consumer', 'business')]
        [string]$MediaType,

        [Parameter(Mandatory = $false)]
        [pscustomobject]$Metadata
    )
    WriteLog "Downloading Windows $WindowsRelease ESD file"
    WriteLog "Windows Architecture: $WindowsArch"
    WriteLog "Windows Language: $WindowsLang"
    WriteLog "Windows Media Type: $MediaType"

    $esdMetadata = $Metadata
    if (-not $esdMetadata) {
        $esdMetadata = Get-WindowsESDMetadata -WindowsRelease $WindowsRelease -WindowsArch $WindowsArch -WindowsLang $WindowsLang -MediaType $MediaType
    }
    if (-not $esdMetadata) {
        throw "Unable to resolve Windows ESD metadata."
    }

    $esdFilePath = $esdMetadata.LocalPath
    $latestEsdFileName = $esdMetadata.FileName
    $existingEsdMatch = $null

    # Reuse an existing ESD only when it matches the latest metadata filename
    if (-not [string]::IsNullOrWhiteSpace($latestEsdFileName)) {
        $existingEsdMatch = Get-ChildItem -Path $PSScriptRoot -Filter *.esd -File -ErrorAction SilentlyContinue | Where-Object { $_.Name -ieq $latestEsdFileName } | Select-Object -First 1
    }

    if ($null -ne $existingEsdMatch) {
        $esdFilePath = $existingEsdMatch.FullName
        WriteLog "Found existing ESD matching metadata filename '$latestEsdFileName' at $esdFilePath. Skipping download."
    }
    else {
        WriteLog "No existing ESD matching metadata filename '$latestEsdFileName' was found. Downloading latest ESD to $esdFilePath."
        WriteLog "Downloading $($esdMetadata.FileUrl) to $esdFilePath"
        $OriginalVerbosePreference = $VerbosePreference
        $VerbosePreference = 'SilentlyContinue'
        Mark-DownloadInProgress -FFUDevelopmentPath $FFUDevelopmentPath -TargetPath $esdFilePath
        Start-BitsTransferWithRetry -Source $esdMetadata.FileUrl -Destination $esdFilePath
        Clear-DownloadInProgress -FFUDevelopmentPath $FFUDevelopmentPath -TargetPath $esdFilePath
        $VerbosePreference = $OriginalVerbosePreference
        WriteLog "ESD download succeeded"
    }

    return $esdFilePath
}

function Get-ODTURL {
    try {
        [String]$ODTPage = Invoke-WebRequest 'https://www.microsoft.com/en-us/download/details.aspx?id=49117' -Headers $Headers -UserAgent $UserAgent -ErrorAction Stop

        # Extract JSON data from the webpage
        if ($ODTPage -match '<script>window\.__DLCDetails__=(.*?)<\/script>') {
            # Parse JSON content
            $jsonContent = $matches[1] | ConvertFrom-Json
            $ODTURL = $jsonContent.dlcDetailsView.downloadFile[0].url

            if ($ODTURL) {
                return $ODTURL
            }
            else {
                WriteLog 'Cannot find the ODT download URL in the JSON content'
                throw 'Cannot find the ODT download URL in the JSON content'
            }
        }
        else {
            WriteLog 'Failed to extract JSON content from the ODT webpage'
            throw 'Failed to extract JSON content from the ODT webpage'
        }
    }
    catch {
        WriteLog $_.Exception.Message
        throw 'An error occurred while retrieving the ODT URL.'
    }
}

function Get-Office {
    # If a custom Office Config XML is provided via config file, use its filename for the installation.
    # The UI script is responsible for copying the file itself to the OfficePath.
    if ((Get-Variable -Name 'OfficeConfigXMLFile' -ErrorAction SilentlyContinue) -and -not([string]::IsNullOrEmpty($OfficeConfigXMLFile))) {
        $script:OfficeInstallXML = Split-Path -Path $OfficeConfigXMLFile -Leaf
        WriteLog "A custom Office configuration file was specified. Using '$($script:OfficeInstallXML)' for installation."
    }
    #Download ODT
    $ODTUrl = Get-ODTURL
    $ODTInstallFile = "$OfficePath\odtsetup.exe"
    WriteLog "Downloading Office Deployment Toolkit from $ODTUrl to $ODTInstallFile"
    $OriginalVerbosePreference = $VerbosePreference
    $VerbosePreference = 'SilentlyContinue'
    Invoke-WebRequest -Uri $ODTUrl -OutFile $ODTInstallFile -Headers $Headers -UserAgent $UserAgent
    $VerbosePreference = $OriginalVerbosePreference

    # Extract ODT
    WriteLog "Extracting ODT to $OfficePath"
    Invoke-Process $ODTInstallFile "/extract:$OfficePath /quiet" | Out-Null

    # Run setup.exe with config.xml and modify xml file to download to $OfficePath
    $xmlContent = [xml](Get-Content $OfficeDownloadXML)
    $xmlContent.Configuration.Add.SourcePath = $OfficePath
    # Back up existing XML content before changing it so cancel cleanup can restore pre-run state.
    Backup-RunFile -FFUDevelopmentPath $FFUDevelopmentPath -Path $OfficeDownloadXML
    $xmlContent.Save($OfficeDownloadXML)
    Mark-DownloadInProgress -FFUDevelopmentPath $FFUDevelopmentPath -TargetPath $OfficePath
    WriteLog "Downloading M365 Apps/Office to $OfficePath"
    Invoke-Process $OfficePath\setup.exe "/download $OfficeDownloadXML" | Out-Null
    Clear-DownloadInProgress -FFUDevelopmentPath $FFUDevelopmentPath -TargetPath $OfficePath

    WriteLog "Cleaning up ODT default config files"
    #Clean up default configuration files
    Remove-Item -Path "$OfficePath\configuration*" -Force

    #Create Install-Office.ps1 in $orchestrationpath
    WriteLog "Creating $orchestrationpath\Install-Office.ps1"   
    $installOfficePath = Join-Path -Path $orchestrationpath -ChildPath "Install-Office.ps1"
    # Create the Install-Office.ps1 file
    $installOfficeCommand = "& `"`$env:FFUAppsRoot\Office\setup.exe`" /configure `"`$env:FFUAppsRoot\Office\$OfficeInstallXML`""
    # Back up any pre-existing script with the same name before overwrite.
    Backup-RunFile -FFUDevelopmentPath $FFUDevelopmentPath -Path $installOfficePath
    Set-Content -Path $installOfficePath -Value $installOfficeCommand -Force
    WriteLog "Install-Office.ps1 created successfully at $installOfficePath"

    #Remove the ODT setup file
    WriteLog "Removing ODT setup file"
    Remove-Item -Path $ODTInstallFile -Force
    WriteLog "ODT setup file removed"
}

function Get-KBLink {
    param(
        [Parameter(Mandatory)]
        [string]$Name
    )
    $OriginalVerbosePreference = $VerbosePreference
    $VerbosePreference = 'SilentlyContinue'
    $results = Invoke-WebRequest -Uri "http://www.catalog.update.microsoft.com/Search.aspx?q=$Name" -Headers $Headers -UserAgent $UserAgent
    $VerbosePreference = $OriginalVerbosePreference

    # Extract the first KB article ID and Windows version (if present) from the HTML content and store globally
    # Edge and Defender do not have KB article IDs
    if ($Name -notmatch 'Defender|Edge') {
        $global:LastKBArticleID = $null
        $global:LastKBWindowsVersion = $null
        if ($results.Content -match '\(KB(\d+)\)[^(<]*\(([0-9]+\.[0-9]+)\)\s*<') {
            $kbArticleID = "KB$($matches[1])"
            $global:LastKBArticleID = $kbArticleID
            $global:LastKBWindowsVersion = $matches[2]
            WriteLog "Found KB article ID: $kbArticleID with Windows version $($matches[2])"
        }
        elseif ($results.Content -match '>\s*([^\(<]+)\(KB(\d+)\)(?:\s*\([^)]+\))*\s*<') {
            $kbArticleID = "KB$($matches[2])"
            $global:LastKBArticleID = $kbArticleID
            WriteLog "Found KB article ID: $kbArticleID (no Windows version found)"
        }
        else {
            WriteLog "No KB article ID found in search results."
        }
    }

    $kbids = $results.InputFields |
    Where-Object { $_.type -eq 'Button' -and $_.Value -eq 'Download' } |
    Select-Object -ExpandProperty  ID

    if (-not $kbids) {
        Write-Warning -Message "No results found for $Name"
        return
    }

    $guids = $results.Links |
    Where-Object ID -match '_link' |
    Where-Object { $_.OuterHTML -match ( "(?=.*" + ( $Filter -join ")(?=.*" ) + ")" ) } |
    Select-Object -First 1 |
    ForEach-Object { $_.id.replace('_link', '') } |
    Where-Object { $_ -in $kbids }

    if (-not $guids) {
        Write-Warning -Message "No file found for $Name"
        return
    }

    foreach ($guid in $guids) {
        # Write-Verbose -Message "Downloading information for $guid"
        $post = @{ size = 0; updateID = $guid; uidInfo = $guid } | ConvertTo-Json -Compress
        $body = @{ updateIDs = "[$post]" }
        $OriginalVerbosePreference = $VerbosePreference
        $VerbosePreference = 'SilentlyContinue'
        $links = Invoke-WebRequest -Uri 'https://www.catalog.update.microsoft.com/DownloadDialog.aspx' -Method Post -Body $body -Headers $Headers -UserAgent $UserAgent |
        Select-Object -ExpandProperty Content |
        Select-String -AllMatches -Pattern "http[s]?://[^']*\.microsoft\.com/[^']*|http[s]?://[^']*\.windowsupdate\.com/[^']*" |
        Select-Object -Unique
        $VerbosePreference = $OriginalVerbosePreference

        foreach ($link in $links) {
            $link.matches.value
            #Filter out cab files
            # #if ($link -notmatch '\.cab') {
            #     $link.matches.value
            # }
                    
        }
    }  
}


function Get-UpdateFileInfo {
    [CmdletBinding()]
    param(
        [string[]]$Name
    )
    $updateFileInfos = [System.Collections.Generic.List[pscustomobject]]::new()

    foreach ($kb in $Name) {
        $links = Get-KBLink -Name $kb
        foreach ($link in $links) {
            $fileName = ($link -split '/')[-1]

            $architectureMatch = $false
            if ($link -match 'x64' -or $link -match 'amd64') {
                if ($WindowsArch -eq 'x64') { $architectureMatch = $true }
            }
            elseif ($link -match 'arm64') {
                if ($WindowsArch -eq 'arm64') { $architectureMatch = $true }
            }
            elseif ($link -match 'x86') {
                if ($WindowsArch -eq 'x86') { $architectureMatch = $true }
            }
            else {
                # If no architecture is specified in the URL, we assume the search query was specific enough.
                # The alternative is to download the file to check, which defeats the purpose of this function.
                $architectureMatch = $true
            }

            if ($architectureMatch) {
                # Check for duplicates before adding
                if (-not ($updateFileInfos.Name -contains $fileName)) {
                    $updateFileInfos.Add([pscustomobject]@{
                            Name = $fileName
                            Url  = $link
                        })
                }
            }
        }
    }
    return $updateFileInfos
}

function Save-KB {
    [CmdletBinding()]
    param(
        [string[]]$Name,
        [string]$Path
    )
    foreach ($kb in $name) {
        $links = Get-KBLink -Name $kb
        foreach ($link in $links) {
            # if (!($link -match 'x64' -or $link -match 'amd64' -or $link -match 'x86' -or $link -match 'arm64')) {
            #     WriteLog "No architecture found in $link, skipping"
            #     continue
            # }

            if ($link -match 'x64' -or $link -match 'amd64') {
                if ($WindowsArch -eq 'x64') {
                    Writelog "Downloading $link for $WindowsArch to $Path"
                    Start-BitsTransferWithRetry -Source $link -Destination $Path
                    $fileName = ($link -split '/')[-1]
                    Writelog "Returning $fileName"
                }
                
            }
            elseif ($link -match 'arm64') {
                if ($WindowsArch -eq 'arm64') {
                    Writelog "Downloading $Link for $WindowsArch to $Path"
                    Start-BitsTransferWithRetry -Source $link -Destination $Path
                    $fileName = ($link -split '/')[-1]
                    Writelog "Returning $fileName"
                }
            }
            elseif ($link -match 'x86') {
                if ($WindowsArch -eq 'x86') {
                    Writelog "Downloading $link for $WindowsArch to $Path"
                    Start-BitsTransferWithRetry -Source $link -Destination $Path
                    $fileName = ($link -split '/')[-1]
                    Writelog "Returning $fileName"
                }

            }
            else {
                WriteLog "No architecture found in $link"
                
                #If no architecture is found, download the file and run it through Get-PEArchitecture to determine the architecture
                Writelog "Downloading $link to $Path and analyzing file for architecture"
                Start-BitsTransferWithRetry -Source $link -Destination $Path

                #Take the file and run it through Get-PEArchitecture to determine the architecture
                $fileName = ($link -split '/')[-1]
                $filePath = Join-Path -Path $Path -ChildPath $fileName
                $arch = Get-PEArchitecture -FilePath $filePath
                Writelog "$fileName is $arch"
                #If the architecture matches $WindowsArch, keep the file, otherwise delete it
                if ($arch -eq $WindowsArch) {
                    Writelog "Architecture for $fileName matches $WindowsArch, keeping file"
                    return $fileName
                }
                else {
                    Writelog "Deleting $fileName, architecture does not match"
                    Remove-Item -Path $filePath -Force
                }
            }
             
        }
    }
    return $fileName
}

function Sync-UserAppListForOrchestration {
    param(
        [Parameter(Mandatory)]
        [string]$SourcePath,
        [Parameter(Mandatory)]
        [string]$AppsPath,
        [Parameter(Mandatory)]
        [string]$OrchestrationPath,
        [Parameter(Mandatory)]
        [string]$AppInstallConfigPath
    )

    # Ensure the orchestration folder exists before writing runtime configuration.
    if (-not (Test-Path -Path $OrchestrationPath -PathType Container)) {
        New-Item -Path $OrchestrationPath -ItemType Directory -Force | Out-Null
    }

    # Persist the runtime BYO app list path so Orchestrator can honor custom file names.
    $appInstallConfig = [ordered]@{
        UserAppListPath = $null
    }

    if (-not [string]::IsNullOrWhiteSpace($SourcePath) -and (Test-Path -Path $SourcePath -PathType Leaf)) {
        $stagedUserAppListName = Split-Path -Path $SourcePath -Leaf
        $stagedUserAppListPath = Join-Path -Path $AppsPath -ChildPath $stagedUserAppListName

        if (-not [string]::Equals([System.IO.Path]::GetFullPath($SourcePath), [System.IO.Path]::GetFullPath($stagedUserAppListPath), [System.StringComparison]::OrdinalIgnoreCase)) {
            Copy-Item -Path $SourcePath -Destination $stagedUserAppListPath -Force | Out-Null
            WriteLog "Staged BYO app list for orchestration: $SourcePath -> $stagedUserAppListPath"
        }
        else {
            WriteLog "Using BYO app list already staged at $stagedUserAppListPath"
        }

        $appInstallConfig.UserAppListPath = "%FFUAppsRoot%\$stagedUserAppListName"
    }
    elseif (Test-Path -Path (Join-Path -Path $AppsPath -ChildPath 'UserAppList.json') -PathType Leaf) {
        $appInstallConfig.UserAppListPath = "%FFUAppsRoot%\UserAppList.json"
        WriteLog "Using default BYO app list path for orchestration."
    }
    else {
        WriteLog "No BYO app list found at configured path '$SourcePath'."
    }

    $appInstallConfig | ConvertTo-Json | Set-Content -Path $AppInstallConfigPath -Encoding UTF8
    WriteLog "Wrote app install config to $AppInstallConfigPath"
}

function New-AppsISO {
    #Create Apps ISO file
    $OSCDIMG = "$adkpath`Assessment and Deployment Kit\Deployment Tools\amd64\Oscdimg\oscdimg.exe"
    #Adding Long Path support for AppsPath to prevent issues with oscdimg
    $AppsPath = '\\?\' + $AppsPath
    Invoke-Process $OSCDIMG "-n -m -d $Appspath $AppsISO" | Out-Null  
}
function Test-AppsIsoRefreshRequired {
    param(
        [string]$AppsISOPath,
        [string]$AppsPath,
        [string[]]$InputPaths = @()
    )

    if ([string]::IsNullOrWhiteSpace($AppsISOPath) -or -not (Test-Path -Path $AppsISOPath -PathType Leaf)) {
        return $true
    }

    $appsIsoLastWriteUtc = (Get-Item -Path $AppsISOPath).LastWriteTimeUtc

    foreach ($inputPath in $InputPaths) {
        if ([string]::IsNullOrWhiteSpace($inputPath) -or -not (Test-Path -Path $inputPath)) { continue }

        $inputItem = Get-Item -Path $inputPath -Force
        if ($inputItem.LastWriteTimeUtc -gt $appsIsoLastWriteUtc) {
            WriteLog "Apps ISO refresh required because $($inputItem.FullName) is newer than $AppsISOPath"
            return $true
        }
    }

    if (-not [string]::IsNullOrWhiteSpace($AppsPath) -and (Test-Path -Path $AppsPath -PathType Container)) {
        $newerAppInput = Get-ChildItem -Path $AppsPath -File -Force -Recurse -ErrorAction SilentlyContinue | Where-Object {
            $_.LastWriteTimeUtc -gt $appsIsoLastWriteUtc
        } | Select-Object -First 1

        if ($null -ne $newerAppInput) {
            WriteLog "Apps ISO refresh required because $($newerAppInput.FullName) is newer than $AppsISOPath"
            return $true
        }
    }

    return $false
}
function Dismount-WindowsIsoSource {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $true)]
        [string]$ImagePath
    )

    $diskImage = Get-DiskImage -ImagePath $ImagePath -ErrorAction SilentlyContinue
    if ($null -eq $diskImage -or -not $diskImage.Attached) {
        return
    }

    WriteLog 'Dismounting Windows ISO'
    Dismount-DiskImage -ImagePath $ImagePath -ErrorAction Stop | Out-Null
    WriteLog 'Done'
}
function Get-WimFromISO {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $true)]
        [string]$ImagePath,

        [string[]]$ExcludedDriveLetters = @()
    )

    $existingDiskImage = Get-DiskImage -ImagePath $ImagePath -ErrorAction Stop
    if ($existingDiskImage.Attached) {
        throw "Windows ISO '$ImagePath' is already mounted. Dismount it before starting the build."
    }

    $mountedByBuild = $false
    try {
        $mountResult = Mount-DiskImage -ImagePath $ImagePath -PassThru -ErrorAction Stop
        $mountedByBuild = $true
        $mountedVolumes = @($mountResult | Get-Volume -ErrorAction Stop)
        if ($mountedVolumes.Count -ne 1) {
            throw "Expected one mounted volume for Windows ISO '$ImagePath', found $($mountedVolumes.Count)."
        }

        $mountedVolume = $mountedVolumes[0]
        $isoDriveLetter = ([string]$mountedVolume.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
        if ($isoDriveLetter -notmatch '^[A-Z]$') {
            throw "Windows ISO '$ImagePath' was mounted without a usable drive letter."
        }

        $normalizedExcludedDriveLetters = @($ExcludedDriveLetters | ForEach-Object {
                ([string]$_).Trim().TrimEnd(':').ToUpperInvariant()
            } | Where-Object { $_ -match '^[A-Z]$' } | Select-Object -Unique)
        if ($normalizedExcludedDriveLetters -contains $isoDriveLetter) {
            $replacementDriveLetter = Get-AvailableDriveLetter -ExcludedDriveLetters $normalizedExcludedDriveLetters
            if ($null -eq $replacementDriveLetter) {
                throw 'No drive letter is available for the mounted Windows ISO outside the configured build partition letters.'
            }

            $isoVolume = Get-CimInstance -ClassName Win32_Volume -Filter "DriveLetter = '${isoDriveLetter}:' AND DriveType = 5" -ErrorAction Stop
            if ($null -eq $isoVolume) {
                throw "Unable to resolve the mounted Windows ISO volume at ${isoDriveLetter}:."
            }

            WriteLog "Windows ISO drive ${isoDriveLetter}: conflicts with the configured build partition layout. Moving it to ${replacementDriveLetter}:."
            Set-CimInstance -InputObject $isoVolume -Property @{ DriveLetter = "${replacementDriveLetter}:" } -ErrorAction Stop | Out-Null

            $mountedVolumes = @($mountResult | Get-Volume -ErrorAction Stop)
            $verifiedDriveLetter = if ($mountedVolumes.Count -eq 1) {
                ([string]$mountedVolumes[0].DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
            }
            if ($verifiedDriveLetter -ne $replacementDriveLetter) {
                throw "Windows ISO drive-letter verification failed after assigning ${replacementDriveLetter}:."
            }
            $isoDriveLetter = $verifiedDriveLetter
        }

        $sourcesFolder = "${isoDriveLetter}:\sources"
        $windowsImage = Get-ChildItem -LiteralPath $sourcesFolder -Filter 'install.*' -File -ErrorAction Stop | Where-Object {
            $_.Name -match '^install\.(wim|esd)$'
        } | Select-Object -First 1
        if ($null -eq $windowsImage) {
            throw "No install.wim or install.esd file was found in '$sourcesFolder'."
        }

        WriteLog "The path to the install file is: $($windowsImage.FullName)"
        return $windowsImage.FullName
    }
    catch {
        $mountError = $_
        if ($mountedByBuild) {
            try {
                Dismount-WindowsIsoSource -ImagePath $ImagePath
            }
            catch {
                WriteLog "Failed to dismount Windows ISO after a mount error: $($_.Exception.Message)"
            }
        }
        throw $mountError
    }
}

function Get-ResolvedWindowsSKUFromImage {
    param(
        [Parameter(Mandatory = $true)]
        [string]$EditionId,

        [string]$InstallationType,

        [string]$ImageName,

        [Parameter(Mandatory = $true)]
        [int]$WindowsRelease
    )

    $normalizedInstallationType = if ([string]::IsNullOrWhiteSpace($InstallationType)) { '' } else { $InstallationType.Trim() }

    switch ($EditionId) {
        'Core' { return 'Home' }
        'CoreN' { return 'Home N' }
        'CoreSingleLanguage' { return 'Home Single Language' }
        'Education' { return 'Education' }
        'EducationN' { return 'Education N' }
        'Professional' { return 'Pro' }
        'ProfessionalN' { return 'Pro N' }
        'ProfessionalEducation' { return 'Pro Education' }
        'ProfessionalEducationN' { return 'Pro Education N' }
        'ProfessionalWorkstation' { return 'Pro for Workstations' }
        'ProfessionalWorkstationN' { return 'Pro N for Workstations' }
        'Enterprise' { return 'Enterprise' }
        'EnterpriseN' { return 'Enterprise N' }
        'EnterpriseS' {
            if ($WindowsRelease -eq 2016 -or $ImageName -match 'LTSB') {
                return 'Enterprise 2016 LTSB'
            }
            return 'Enterprise LTSC'
        }
        'EnterpriseSN' {
            if ($WindowsRelease -eq 2016 -or $ImageName -match 'LTSB') {
                return 'Enterprise N 2016 LTSB'
            }
            return 'Enterprise N LTSC'
        }
        'IoTEnterpriseS' { return 'IoT Enterprise LTSC' }
        'IoTEnterpriseSN' { return 'IoT Enterprise N LTSC' }
        'ServerStandard' {
            if ($normalizedInstallationType -eq 'Server') {
                return 'Standard (Desktop Experience)'
            }
            return 'Standard'
        }
        'ServerDatacenter' {
            if ($normalizedInstallationType -eq 'Server') {
                return 'Datacenter (Desktop Experience)'
            }
            return 'Datacenter'
        }
    }

    return $null
}

function Get-WindowsImageSelection {
    param(
        [Parameter(Mandatory = $true)]
        [string]$WindowsImagePath,

        [Parameter(Mandatory = $true)]
        [string]$WindowsSKU,

        [Parameter(Mandatory = $true)]
        [int]$WindowsRelease
    )

    # Get the available indexes in the WIM/ESD
    $imageIndexes = Get-WindowsImage -ImagePath $WindowsImagePath

    # Normalize SKU and determine if Desktop Experience is explicitly requested (Server only)
    $normalizedWindowsSKU = $WindowsSKU.Trim()
    $isDesktopExperienceRequested = $normalizedWindowsSKU -match '\(Desktop Experience\)'
    $normalizedWindowsSKU = $normalizedWindowsSKU -replace '\s*\(Desktop Experience\)\s*', ''

    # Map user-selected SKU to language-independent EditionId values
    # Notes:
    # - Client: EditionId values are stable across languages (e.g. Professional, Core, Education)
    # - Server: Desktop Experience vs Core is differentiated by InstallationType (EditionId is the same)
    $editionIdCandidates = switch ($normalizedWindowsSKU) {
        'Home' { @('Core') }
        'Core' { @('Core') }

        'Home N' { @('CoreN') }
        'CoreN' { @('CoreN') }

        'Home Single Language' { @('CoreSingleLanguage') }
        'CoreSingleLanguage' { @('CoreSingleLanguage') }

        'Education' { @('Education') }
        'Education N' { @('EducationN') }
        'EducationN' { @('EducationN') }

        'Pro' { @('Professional') }
        'Professional' { @('Professional') }

        'Pro N' { @('ProfessionalN') }
        'ProfessionalN' { @('ProfessionalN') }

        'Pro Education' { @('ProfessionalEducation') }
        'ProfessionalEducation' { @('ProfessionalEducation') }

        'Pro Education N' { @('ProfessionalEducationN') }
        'ProfessionalEducationN' { @('ProfessionalEducationN') }

        'Pro for Workstations' { @('ProfessionalWorkstation') }
        'ProfessionalWorkstation' { @('ProfessionalWorkstation') }

        'Pro N for Workstations' { @('ProfessionalWorkstationN') }
        'ProfessionalWorkstationN' { @('ProfessionalWorkstationN') }

        'Enterprise' { @('Enterprise') }
        'Enterprise N' { @('EnterpriseN') }
        'EnterpriseN' { @('EnterpriseN') }

        'Enterprise LTSC' { @('EnterpriseS') }
        'Enterprise 2016 LTSB' { @('EnterpriseS') }
        'EnterpriseS' { @('EnterpriseS') }

        'Enterprise N LTSC' { @('EnterpriseSN') }
        'Enterprise N 2016 LTSB' { @('EnterpriseSN') }
        'EnterpriseSN' { @('EnterpriseSN') }

        'IoT Enterprise LTSC' { @('IoTEnterpriseS') }
        'IoTEnterpriseS' { @('IoTEnterpriseS') }

        'IoT Enterprise N LTSC' { @('IoTEnterpriseSN') }
        'IoTEnterpriseSN' { @('IoTEnterpriseSN') }

        'Standard' { @('ServerStandard') }
        'ServerStandard' { @('ServerStandard') }

        'Datacenter' { @('ServerDatacenter') }
        'ServerDatacenter' { @('ServerDatacenter') }

        default { @() }
    }

    # Determine preferred InstallationType for Server images
    $preferredInstallationType = $null
    if ($normalizedWindowsSKU -in @('Standard', 'Datacenter', 'ServerStandard', 'ServerDatacenter')) {
        if ($isDesktopExperienceRequested) {
            $preferredInstallationType = 'Server'
        }
        else {
            $preferredInstallationType = 'Server Core'
        }
    }

    # Build per-index metadata (EditionId, InstallationType, resolved SKU) once for deterministic matching and fallback prompts.
    $imageMetadata = @(foreach ($imageIndex in $imageIndexes) {
            try {
                $details = Get-WindowsImage -ImagePath $WindowsImagePath -Index $imageIndex.ImageIndex
                [pscustomobject]@{
                    ImageIndex        = $details.ImageIndex
                    ImageName         = $details.ImageName
                    ImageSize         = $details.ImageSize
                    EditionId         = $details.EditionId
                    InstallationType  = $details.InstallationType
                    ResolvedWindowsSKU = Get-ResolvedWindowsSKUFromImage -EditionId $details.EditionId -InstallationType $details.InstallationType -ImageName $details.ImageName -WindowsRelease $WindowsRelease
                }
            }
            catch {
                $null
            }
        }) | Where-Object { $null -ne $_ }

    # If we can map SKU -> EditionId, attempt a non-interactive match
    if ($editionIdCandidates.Count -gt 0) {
        # Match by EditionId first
        $imageMatches = $imageMetadata | Where-Object { $_.EditionId -in $editionIdCandidates }

        # If this is a Server SKU, prefer the requested InstallationType (Server vs Server Core)
        if ($null -ne $preferredInstallationType -and $imageMatches.Count -gt 0) {
            $preferredImageMatches = $imageMatches | Where-Object { $_.InstallationType -eq $preferredInstallationType }
            if ($preferredImageMatches.Count -gt 0) {
                $imageMatches = $preferredImageMatches
            }
        }

        # If multiple matches remain, pick the largest image (Desktop Experience tends to be larger)
        if ($imageMatches.Count -gt 0) {
            $bestMatch = $imageMatches | Sort-Object -Property ImageSize -Descending | Select-Object -First 1
            WriteLog "Selected Windows image index $($bestMatch.ImageIndex) (RequestedSKU='$WindowsSKU', ResolvedSKU='$($bestMatch.ResolvedWindowsSKU)', EditionId='$($bestMatch.EditionId)', InstallationType='$($bestMatch.InstallationType)'): $($bestMatch.ImageName)"
            return $bestMatch
        }
    }

    # Final fallback: prompt the user to select an ImageName
    # Look for the numbers 10, 11, 2016, 2019, 2022+ in the ImageName
    $relevantImageIndexes = @($imageMetadata | Where-Object { $_.ImageName -match "(10|11|2016|2019|202\d)" })
    if ($relevantImageIndexes.Count -eq 0) {
        $relevantImageIndexes = $imageMetadata
    }

    WriteLog "No matching image index found for SKU '$WindowsSKU' in '$WindowsImagePath'. Prompting user to select an ImageName."

    while ($true) {
        # Present list of ImageNames to the end user if no matching ImageIndex is found
        Write-Host "No matching ImageIndex found for Windows SKU '$WindowsSKU'. Please select an ImageName from the list below:"

        $i = 1
        $relevantImageIndexes | ForEach-Object {
            Write-Host "$i. $($_.ImageName)"
            $i++
        }

        # Ask for user input
        $inputValue = Read-Host "Enter the number of the ImageName you want to use"

        # Get selected ImageName based on user input
        $selectedImage = $relevantImageIndexes[$inputValue - 1]

        if ($selectedImage) {
            WriteLog "User selected Windows image index $($selectedImage.ImageIndex) (RequestedSKU='$WindowsSKU', ResolvedSKU='$($selectedImage.ResolvedWindowsSKU)'): $($selectedImage.ImageName)"
            return $selectedImage
        }
        else {
            Write-Host "Invalid selection, please try again."
        }
    }
}

#Create VHDX
function New-ScratchVhdx {
    param(
        [Parameter(Mandatory = $true)]
        [string]$VhdxPath,
        [uint64]$SizeBytes = 50GB,
        [uint32]$LogicalSectorSizeBytes,
        [switch]$Dynamic,
        [Microsoft.PowerShell.Cmdletization.GeneratedTypes.Disk.PartitionStyle]$PartitionStyle = [Microsoft.PowerShell.Cmdletization.GeneratedTypes.Disk.PartitionStyle]::GPT
    )

    WriteLog "Creating new Scratch VHDX..."

    $newVHDX = New-VHD -Path $VhdxPath -SizeBytes $disksize -LogicalSectorSizeBytes $LogicalSectorSizeBytes -Dynamic:($Dynamic.IsPresent) 
    $toReturn = $newVHDX | Mount-VHD -Passthru | Initialize-Disk -PassThru -PartitionStyle GPT

    #Remove auto-created partition so we can create the correct partition layout
    remove-partition $toreturn.DiskNumber -PartitionNumber 1 -Confirm:$False 

    Writelog "Done."
    return $toReturn
}
function Get-NormalizedPartitionDriveLetters {
    param(
        [string]$SystemPartitionDriveLetter,
        [string]$WindowsPartitionDriveLetter,
        [string]$RecoveryPartitionDriveLetter,
        [bool]$CreateRecoveryPartition = $true,
        [object[]]$AdditionalDataPartitions = @(),
        [switch]$ValidateAvailable
    )

    $requestedLetters = [ordered]@{
        SystemPartitionDriveLetter  = $SystemPartitionDriveLetter
        WindowsPartitionDriveLetter = $WindowsPartitionDriveLetter
    }
    if ($CreateRecoveryPartition) {
        $requestedLetters['RecoveryPartitionDriveLetter'] = $RecoveryPartitionDriveLetter
    }
    foreach ($dataPartition in @($AdditionalDataPartitions)) {
        if ($null -eq $dataPartition) { continue }
        $partitionName = [string]$dataPartition.Name
        if ([string]::IsNullOrWhiteSpace($partitionName)) {
            $partitionName = [string]$dataPartition.Label
        }
        if ([string]::IsNullOrWhiteSpace($partitionName)) {
            $partitionName = "DataPartition$($requestedLetters.Count - 2)"
        }
        $requestedLetters["AdditionalDataPartition:$partitionName"] = $dataPartition.DriveLetter
    }
    $normalizedLetters = [ordered]@{}

    foreach ($entry in $requestedLetters.GetEnumerator()) {
        $driveLetter = ([string]$entry.Value).Trim().TrimEnd(':').ToUpperInvariant()
        if ([string]::IsNullOrWhiteSpace($driveLetter) -or $driveLetter -notmatch '^[A-Z]$') {
            throw "$($entry.Key) must be a single drive letter from A to Z without a colon."
        }

        $normalizedLetters[$entry.Key] = $driveLetter
    }

    $duplicateLetters = @($normalizedLetters.Values | Group-Object | Where-Object { $_.Count -gt 1 })
    if ($duplicateLetters.Count -gt 0) {
        $duplicateLetterList = ($duplicateLetters | ForEach-Object { $_.Name }) -join ', '
        throw "Build partition drive letters must be unique. Duplicate value(s): $duplicateLetterList."
    }

    if ($ValidateAvailable) {
        foreach ($entry in $normalizedLetters.GetEnumerator()) {
            $existingDrive = Get-PSDrive -Name $entry.Value -PSProvider FileSystem -ErrorAction SilentlyContinue
            if ($null -ne $existingDrive) {
                throw "$($entry.Key) uses drive letter $($entry.Value), but that letter is already assigned on this host. Choose an unused drive letter."
            }
        }
    }

    return [pscustomobject]$normalizedLetters
}
function ConvertTo-UInt64SizeBytes {
    param(
        [object]$Value,
        [string]$PropertyName
    )

    if ($null -eq $Value -or [string]::IsNullOrWhiteSpace([string]$Value)) {
        return 0
    }

    [uint64]$sizeBytes = 0
    if (-not [uint64]::TryParse([string]$Value, [ref]$sizeBytes)) {
        throw "$PropertyName must be a whole number."
    }

    return $sizeBytes
}
function ConvertTo-NormalizedDataPartitions {
    param(
        [object[]]$DataPartitions = @()
    )

    $normalizedPartitions = [System.Collections.Generic.List[pscustomobject]]::new()
    $partitionIndex = 0
    foreach ($dataPartition in @($DataPartitions)) {
        if ($null -eq $dataPartition) { continue }
        $partitionIndex++

        $name = [string]$dataPartition.Name
        $label = [string]$dataPartition.Label
        if ([string]::IsNullOrWhiteSpace($name)) { $name = "Data$partitionIndex" }
        if ([string]::IsNullOrWhiteSpace($label)) { $label = $name }

        $driveLetter = ([string]$dataPartition.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
        if ([string]::IsNullOrWhiteSpace($driveLetter) -or $driveLetter -notmatch '^[D-Z]$') {
            throw "Additional data partition '$name' must specify a single drive letter from D to Z without a colon."
        }

        $fileSystem = [string]$dataPartition.FileSystem
        if ([string]::IsNullOrWhiteSpace($fileSystem)) { $fileSystem = 'NTFS' }
        if ($fileSystem -notin @('NTFS', 'ReFS', 'exFAT', 'FAT32')) {
            throw "Additional data partition '$name' uses unsupported file system '$fileSystem'."
        }

        $sizeBytes = ConvertTo-UInt64SizeBytes -Value $dataPartition.SizeBytes -PropertyName "Additional data partition '$name' SizeBytes"
        if ($sizeBytes -eq 0 -and $null -ne $dataPartition.SizeGB -and -not [string]::IsNullOrWhiteSpace([string]$dataPartition.SizeGB)) {
            [decimal]$sizeGb = 0
            if (-not [decimal]::TryParse([string]$dataPartition.SizeGB, [ref]$sizeGb)) {
                throw "Additional data partition '$name' SizeGB must be a number."
            }
            $sizeBytes = [uint64]($sizeGb * 1GB)
        }

        $fillRemaining = $false
        if ($dataPartition.PSObject.Properties.Name -contains 'FillRemaining') {
            $fillRemaining = [System.Convert]::ToBoolean($dataPartition.FillRemaining)
        }

        $persistDriveLetter = $false
        if ($dataPartition.PSObject.Properties.Name -contains 'PersistDriveLetter') {
            $persistDriveLetter = [System.Convert]::ToBoolean($dataPartition.PersistDriveLetter)
        }

        if (($sizeBytes -eq 0) -and (-not $fillRemaining)) {
            throw "Additional data partition '$name' must specify SizeBytes, SizeGB, or FillRemaining."
        }

        $normalizedPartitions.Add([pscustomobject]@{
            Name               = $name
            Label              = $label
            DriveLetter        = $driveLetter
            FileSystem         = $fileSystem
            SizeBytes          = $sizeBytes
            FillRemaining      = $fillRemaining
            PersistDriveLetter = $persistDriveLetter
        })
    }

    $fillRemainingPartitions = @($normalizedPartitions | Where-Object { $_.FillRemaining })
    if ($fillRemainingPartitions.Count -gt 1) {
        throw 'Only one additional data partition can use FillRemaining.'
    }

    if ($fillRemainingPartitions.Count -eq 1) {
        $orderedPartitions = [System.Collections.Generic.List[pscustomobject]]::new()
        foreach ($fixedSizePartition in @($normalizedPartitions | Where-Object { -not $_.FillRemaining })) {
            $orderedPartitions.Add($fixedSizePartition)
        }
        $orderedPartitions.Add($fillRemainingPartitions[0])
        return @($orderedPartitions)
    }

    return @($normalizedPartitions)
}
function Get-PartitionLayoutSignature {
    param(
        [uint64]$OSPartitionSize,
        [uint64]$RecoveryPartitionSize,
        [bool]$CreateRecoveryPartition = $true,
        [object[]]$DataPartitions = @()
    )

    $dataPartitionSignatures = @($DataPartitions | ForEach-Object {
        "$($_.Name)|$($_.Label)|$($_.FileSystem)|$($_.SizeBytes)|$($_.FillRemaining)"
    })

    return "OS=$OSPartitionSize;CreateRecovery=$CreateRecoveryPartition;Recovery=$RecoveryPartitionSize;Data=$($dataPartitionSignatures -join ';')"
}
#Add System Partition
function New-SystemPartition {
    param(
        [Parameter(Mandatory = $true)]
        [ciminstance]$VhdxDisk,
        [string]$DriveLetter = 'S',
        [uint64]$SystemPartitionSize = 260MB
    )

    WriteLog "Creating System partition..."

    $sysPartition = $VhdxDisk | New-Partition -DriveLetter $DriveLetter -Size $SystemPartitionSize -GptType "{c12a7328-f81f-11d2-ba4b-00a0c93ec93b}" -IsHidden 
    $sysPartition | Format-Volume -FileSystem FAT32 -Force -NewFileSystemLabel "System" | Out-Null

    WriteLog 'Done.'
    return [string]$sysPartition.DriveLetter
}
#Add MSRPartition
function New-MSRPartition {
    param(
        [Parameter(Mandatory = $true)]
        [ciminstance]$VhdxDisk
    )

    WriteLog "Creating MSR partition..."

    # $toReturn = $VhdxDisk | New-Partition -AssignDriveLetter -Size 16MB -GptType "{e3c9e316-0b5c-4db8-817d-f92df00215ae}" -IsHidden | Out-Null
    $toReturn = $VhdxDisk | New-Partition -Size 16MB -GptType "{e3c9e316-0b5c-4db8-817d-f92df00215ae}" -IsHidden | Out-Null

    WriteLog "Done."

    return $toReturn
}
#Add OS Partition
function New-OSPartition {
    param(
        [Parameter(Mandatory = $true)]
        [ciminstance]$VhdxDisk,
        [Parameter(Mandatory = $true)]
        [string]$WimPath,
        [uint32]$WimIndex,
        [string]$DriveLetter = 'W',
        [uint64]$OSPartitionSize = 0
    )

    WriteLog "Creating OS partition..."

    if ($OSPartitionSize -gt 0) {
        $osPartition = $vhdxDisk | New-Partition -DriveLetter $DriveLetter -Size $OSPartitionSize -GptType "{ebd0a0a2-b9e5-4433-87c0-68b6b72699c7}" 
    }
    else {
        $osPartition = $vhdxDisk | New-Partition -DriveLetter $DriveLetter -UseMaximumSize -GptType "{ebd0a0a2-b9e5-4433-87c0-68b6b72699c7}" 
    }

    $osPartition | Format-Volume -FileSystem NTFS -Confirm:$false -Force -NewFileSystemLabel "Windows" 
    WriteLog 'Done'
    Writelog "OS partition at drive $($osPartition.DriveLetter):"

    WriteLog "Writing Windows at $WimPath to OS partition at drive $($osPartition.DriveLetter):..."
    
    #Server 2019 is missing the Windows Overlay Filter (wof.sys), likely other Server SKUs are missing it as well. Script will error if trying to use the -compact switch on Server OSes
    if ((Get-CimInstance Win32_OperatingSystem).Caption -match "Server") {
        WriteLog (Expand-WindowsImage -ImagePath $WimPath -Index $WimIndex -ApplyPath "$($osPartition.DriveLetter):\")
    }
    elseif ($CompactOS) {
        WriteLog '$CompactOS is set to true, using -Compact switch to apply the WIM file to the OS partition.'
        WriteLog (Expand-WindowsImage -ImagePath $WimPath -Index $WimIndex -ApplyPath "$($osPartition.DriveLetter):\" -Compact)
    }
    else {
        WriteLog (Expand-WindowsImage -ImagePath $WimPath -Index $WimIndex -ApplyPath "$($osPartition.DriveLetter):\")
    }
    
    WriteLog 'Done'    
    return $osPartition
}
function Resize-OSPartitionForDataPartitions {
    param(
        [Parameter(Mandatory = $true)]
        [ciminstance]$OsPartition,
        [object[]]$DataPartitions = @()
    )

    $fixedDataPartitions = @($DataPartitions | Where-Object { -not $_.FillRemaining })
    if ($fixedDataPartitions.Count -eq 0) {
        return $OsPartition
    }

    [uint64]$dataPartitionReserveSize = 0
    foreach ($dataPartition in $fixedDataPartitions) {
        [uint64]$dataPartitionSize = $dataPartition.SizeBytes
        if ($dataPartitionSize -le 0) {
            throw "Data partition '$($dataPartition.Name)' must have a positive fixed size before space can be reserved."
        }

        $dataPartitionReserveSize += $dataPartitionSize + 1MB
    }

    WriteLog "Reserving $dataPartitionReserveSize bytes from the maximum-size Windows partition for $($fixedDataPartitions.Count) fixed data partition(s)."
    if ($OsPartition.Size -le $dataPartitionReserveSize) {
        throw "Windows partition is too small to reserve $dataPartitionReserveSize bytes for fixed data partitions. Current Windows partition size is $($OsPartition.Size) bytes."
    }

    $newOsPartitionSize = [math]::Floor(($OsPartition.Size - $dataPartitionReserveSize) / 4096) * 4096
    $supportedOsPartitionSize = $OsPartition | Get-PartitionSupportedSize -ErrorAction Stop
    if ($newOsPartitionSize -lt $supportedOsPartitionSize.SizeMin) {
        throw "Windows partition cannot be shrunk enough to reserve $dataPartitionReserveSize bytes for fixed data partitions. Minimum supported Windows partition size is $($supportedOsPartitionSize.SizeMin) bytes."
    }

    $OsPartition | Resize-Partition -Size $newOsPartitionSize -ErrorAction Stop | Out-Null
    WriteLog "Windows partition shrunk to $newOsPartitionSize bytes to reserve space for fixed data partitions."

    return Get-Partition -DiskNumber $OsPartition.DiskNumber -PartitionNumber $OsPartition.PartitionNumber -ErrorAction Stop
}
#Add Recovery partition
function New-RecoveryPartition {
    param(
        [Parameter(Mandatory = $true)]
        [ciminstance]$VhdxDisk,
        [Parameter(Mandatory = $true)]
        $OsPartition,
        [uint64]$RecoveryPartitionSize = 0,
        [string]$DriveLetter = 'R',
        [ciminstance]$DataPartition,
        [bool]$OsPartitionUsesMaximumSize = $true
    )

    WriteLog "Creating empty Recovery partition (to be filled on first boot automatically)..."

    $calculatedRecoverySize = 0
    $recoveryPartition = $null

    try {
        if ($RecoveryPartitionSize -gt 0) {
            $calculatedRecoverySize = $RecoveryPartitionSize
            WriteLog "Using configured Recovery partition size in bytes: $calculatedRecoverySize"
        }
        else {
            $winReWim = Get-ChildItem "$($OsPartition.DriveLetter):\Windows\System32\Recovery\Winre.wim" -Attributes Hidden -ErrorAction SilentlyContinue

            if (($null -ne $winReWim) -and ($winReWim.Count -eq 1)) {
                # Wim size + 100MB is minimum WinRE partition size.
                # NTFS and other partitioning size differences account for about 17MB of space that's unavailable.
                # Adding 32MB as a buffer to ensure there's enough space to account for NTFS file system overhead.
                # Adding 250MB as per recommendations from
                # https://learn.microsoft.com/en-us/windows-hardware/manufacture/desktop/configure-uefigpt-based-hard-drive-partitions?view=windows-11#recovery-tools-partition
                $calculatedRecoverySize = $winReWim.Length + 250MB + 32MB

                WriteLog "Calculated space needed for recovery in bytes: $calculatedRecoverySize"
            }
            else {
                WriteLog "No WinRE.WIM found in the OS partition under \Windows\System32\Recovery."
                WriteLog "Skipping creating the Recovery partition."
                WriteLog "If a Recovery partition is desired, please re-run the script setting the -RecoveryPartitionSize flag as appropriate."
            }
        }

        if ($calculatedRecoverySize -gt 0) {
            $calculatedRecoverySize = [uint64]([math]::Ceiling($calculatedRecoverySize / 4096) * 4096)
            WriteLog "Aligned Recovery partition size in bytes: $calculatedRecoverySize"
            $recoveryPartitionReserveSize = $calculatedRecoverySize + 1MB
            WriteLog "Recovery partition reserve size in bytes: $recoveryPartitionReserveSize"

            if ($null -ne $DataPartition) {
                if ($DataPartition.Size -le $recoveryPartitionReserveSize) {
                    throw "Data partition is too small to reserve $recoveryPartitionReserveSize bytes for the Recovery partition."
                }

                $newDataPartitionSize = [math]::Floor(($DataPartition.Size - $recoveryPartitionReserveSize) / 4096) * 4096
                $supportedDataPartitionSize = $DataPartition | Get-PartitionSupportedSize -ErrorAction Stop

                if ($newDataPartitionSize -lt $supportedDataPartitionSize.SizeMin) {
                    throw "Data partition cannot be shrunk enough to reserve $recoveryPartitionReserveSize bytes for the Recovery partition. Minimum supported data partition size is $($supportedDataPartitionSize.SizeMin) bytes."
                }

                $DataPartition | Resize-Partition -Size $newDataPartitionSize -ErrorAction Stop
                WriteLog "Data partition shrunk by $recoveryPartitionReserveSize bytes for Recovery partition."
            }
            elseif ($OsPartitionUsesMaximumSize) {
                if ($OsPartition.Size -le $recoveryPartitionReserveSize) {
                    throw "OS partition is too small to reserve $recoveryPartitionReserveSize bytes for the Recovery partition."
                }

                $newOsPartitionSize = [math]::Floor(($OsPartition.Size - $recoveryPartitionReserveSize) / 4096) * 4096
                $supportedOsPartitionSize = $OsPartition | Get-PartitionSupportedSize -ErrorAction Stop

                if ($newOsPartitionSize -lt $supportedOsPartitionSize.SizeMin) {
                    throw "OS partition cannot be shrunk enough to reserve $recoveryPartitionReserveSize bytes for the Recovery partition. Minimum supported OS partition size is $($supportedOsPartitionSize.SizeMin) bytes."
                }

                $OsPartition | Resize-Partition -Size $newOsPartitionSize -ErrorAction Stop
                WriteLog "OS partition shrunk by $recoveryPartitionReserveSize bytes for Recovery partition."
            }
            else {
                WriteLog 'Using free space reserved after the Windows partition for Recovery partition.'
            }

            $currentVhdxDisk = Get-Disk -Number $VhdxDisk.Number -ErrorAction Stop
            $largestFreeExtent = [uint64]$currentVhdxDisk.LargestFreeExtent
            WriteLog "Largest free extent available for Recovery partition in bytes: $largestFreeExtent"

            if ($largestFreeExtent -lt $calculatedRecoverySize) {
                throw "Not enough available capacity for Recovery partition. Required: $calculatedRecoverySize bytes. Largest free extent: $largestFreeExtent bytes."
            }

            $partition = $currentVhdxDisk | New-Partition -DriveLetter $DriveLetter -Size $calculatedRecoverySize -GptType "{de94bba4-06d1-4d40-a16a-bfd50179d6ac}" -ErrorAction Stop

            if ([string]::IsNullOrWhiteSpace([string]$partition.DriveLetter)) {
                Set-Partition -DiskNumber $partition.DiskNumber -PartitionNumber $partition.PartitionNumber -NewDriveLetter $DriveLetter -ErrorAction Stop
                $partition = Get-Partition -DiskNumber $partition.DiskNumber -PartitionNumber $partition.PartitionNumber -ErrorAction Stop
            }

            if (([string]$partition.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant() -ne $DriveLetter.Trim().TrimEnd(':').ToUpperInvariant()) {
                throw "Recovery partition was created, but drive letter $DriveLetter`: was not assigned."
            }

            $partition | Format-Volume -FileSystem NTFS -Confirm:$false -Force -NewFileSystemLabel 'Recovery' -ErrorAction Stop | Out-Null
            $recoveryPartition = Get-Partition -DiskNumber $partition.DiskNumber -PartitionNumber $partition.PartitionNumber -ErrorAction Stop

            WriteLog "Done. Recovery partition at drive $($recoveryPartition.DriveLetter):"
        }
    }
    catch {
        WriteLog "Failed to create Recovery partition. $($_.Exception.Message)"
        throw
    }

    return $recoveryPartition
}
#Add Data partition
function New-DataPartition {
    param(
        [Parameter(Mandatory = $true)]
        [ciminstance]$VhdxDisk,
        [Parameter(Mandatory = $true)]
        [pscustomobject]$DataPartition
    )

    WriteLog "Creating data partition '$($DataPartition.Name)'..."

    if ($DataPartition.FillRemaining) {
        $partition = $VhdxDisk | New-Partition -DriveLetter $DataPartition.DriveLetter -UseMaximumSize -GptType "{ebd0a0a2-b9e5-4433-87c0-68b6b72699c7}" -ErrorAction Stop
    }
    else {
        $partition = $VhdxDisk | New-Partition -DriveLetter $DataPartition.DriveLetter -Size $DataPartition.SizeBytes -GptType "{ebd0a0a2-b9e5-4433-87c0-68b6b72699c7}" -ErrorAction Stop
    }

    if ([string]::IsNullOrWhiteSpace([string]$partition.DriveLetter)) {
        Set-Partition -DiskNumber $partition.DiskNumber -PartitionNumber $partition.PartitionNumber -NewDriveLetter $DataPartition.DriveLetter -ErrorAction Stop
        $partition = Get-Partition -DiskNumber $partition.DiskNumber -PartitionNumber $partition.PartitionNumber -ErrorAction Stop
    }

    if (([string]$partition.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant() -ne $DataPartition.DriveLetter) {
        throw "Data partition '$($DataPartition.Name)' was created, but drive letter $($DataPartition.DriveLetter): was not assigned."
    }

    $partition | Format-Volume -FileSystem $DataPartition.FileSystem -Confirm:$false -Force -NewFileSystemLabel $DataPartition.Label -ErrorAction Stop | Out-Null
    WriteLog "Done. Data partition '$($DataPartition.Name)' at drive $($partition.DriveLetter):"

    return $partition
}
function Resolve-VhdxPartitionLayout {
    param(
        [Parameter(Mandatory = $true)]
        [object]$Disk,
        [bool]$CreateRecoveryPartition = $true,
        [object[]]$AdditionalDataPartitions = @()
    )

    $systemGptType = '{c12a7328-f81f-11d2-ba4b-00a0c93ec93b}'
    $msrGptType = '{e3c9e316-0b5c-4db8-817d-f92df00215ae}'
    $basicDataGptType = '{ebd0a0a2-b9e5-4433-87c0-68b6b72699c7}'
    $recoveryGptType = '{de94bba4-06d1-4d40-a16a-bfd50179d6ac}'
    $partitions = @($Disk | Get-Partition -ErrorAction Stop | Sort-Object -Property PartitionNumber)

    $systemPartitions = @($partitions | Where-Object { ([string]$_.GptType).ToLowerInvariant() -eq $systemGptType })
    $msrPartitions = @($partitions | Where-Object { ([string]$_.GptType).ToLowerInvariant() -eq $msrGptType })
    $basicDataPartitions = @($partitions | Where-Object { ([string]$_.GptType).ToLowerInvariant() -eq $basicDataGptType })
    $recoveryPartitions = @($partitions | Where-Object { ([string]$_.GptType).ToLowerInvariant() -eq $recoveryGptType })

    if ($systemPartitions.Count -ne 1) {
        throw "Expected one EFI System partition on VHDX disk $($Disk.Number), found $($systemPartitions.Count)."
    }
    if ($msrPartitions.Count -ne 1) {
        throw "Expected one Microsoft Reserved partition on VHDX disk $($Disk.Number), found $($msrPartitions.Count)."
    }
    if ($CreateRecoveryPartition -and $recoveryPartitions.Count -ne 1) {
        throw "Expected one Recovery partition on VHDX disk $($Disk.Number), found $($recoveryPartitions.Count)."
    }
    if (-not $CreateRecoveryPartition -and $recoveryPartitions.Count -ne 0) {
        throw "Expected no Recovery partition on VHDX disk $($Disk.Number), found $($recoveryPartitions.Count)."
    }

    $expectedBasicDataPartitionCount = 1 + @($AdditionalDataPartitions).Count
    if ($basicDataPartitions.Count -ne $expectedBasicDataPartitionCount) {
        throw "Expected $expectedBasicDataPartitionCount basic data partition(s) on VHDX disk $($Disk.Number), found $($basicDataPartitions.Count)."
    }

    $windowsPartition = $basicDataPartitions[0]
    $dataPartitions = @($basicDataPartitions | Select-Object -Skip 1)
    if ($systemPartitions[0].PartitionNumber -ge $msrPartitions[0].PartitionNumber -or $msrPartitions[0].PartitionNumber -ge $windowsPartition.PartitionNumber) {
        throw "VHDX disk $($Disk.Number) does not use the expected System, MSR, Windows partition order."
    }
    if ($CreateRecoveryPartition) {
        $firstDataPartitionNumber = if ($dataPartitions.Count -gt 0) { $dataPartitions[0].PartitionNumber } else { [int]::MaxValue }
        if ($recoveryPartitions[0].PartitionNumber -le $windowsPartition.PartitionNumber -or $recoveryPartitions[0].PartitionNumber -ge $firstDataPartitionNumber) {
            throw "VHDX disk $($Disk.Number) does not use the expected Windows, Recovery, data partition order."
        }
    }

    $windowsVolumes = @($windowsPartition | Get-Volume -ErrorAction SilentlyContinue)
    if ($windowsVolumes.Count -ne 1 -or ([string]$windowsVolumes[0].FileSystemLabel).Trim() -ine 'Windows') {
        throw "Unable to verify the Windows volume on VHDX disk $($Disk.Number) partition $($windowsPartition.PartitionNumber)."
    }

    $resolvedDataPartitions = [System.Collections.Generic.List[pscustomobject]]::new()
    for ($dataPartitionIndex = 0; $dataPartitionIndex -lt @($AdditionalDataPartitions).Count; $dataPartitionIndex++) {
        $dataPartitionConfig = @($AdditionalDataPartitions)[$dataPartitionIndex]
        $dataPartition = $dataPartitions[$dataPartitionIndex]
        $dataVolumes = @($dataPartition | Get-Volume -ErrorAction SilentlyContinue)
        if ($dataVolumes.Count -ne 1) {
            throw "Unable to resolve data partition '$($dataPartitionConfig.Name)' on VHDX disk $($Disk.Number) partition $($dataPartition.PartitionNumber)."
        }
        if (([string]$dataVolumes[0].FileSystemLabel).Trim() -ine ([string]$dataPartitionConfig.Label).Trim()) {
            throw "Data partition '$($dataPartitionConfig.Name)' label mismatch on VHDX disk $($Disk.Number) partition $($dataPartition.PartitionNumber)."
        }
        if (([string]$dataVolumes[0].FileSystem).Trim() -ine ([string]$dataPartitionConfig.FileSystem).Trim()) {
            throw "Data partition '$($dataPartitionConfig.Name)' file system mismatch on VHDX disk $($Disk.Number) partition $($dataPartition.PartitionNumber)."
        }

        $resolvedDataPartitions.Add([pscustomobject]@{
                Config    = $dataPartitionConfig
                Partition = $dataPartition
                Volume    = $dataVolumes[0]
            })
    }

    return [pscustomobject]@{
        Disk              = $Disk
        SystemPartition   = $systemPartitions[0]
        MsrPartition      = $msrPartitions[0]
        WindowsPartition  = $windowsPartition
        WindowsVolume     = $windowsVolumes[0]
        RecoveryPartition = if ($CreateRecoveryPartition) { $recoveryPartitions[0] } else { $null }
        DataPartitions    = $resolvedDataPartitions.ToArray()
    }
}

function Set-VhdxBuildPartitionDriveLetters {
    param(
        [Parameter(Mandatory = $true)]
        [object]$Layout,
        [Parameter(Mandatory = $true)]
        [string]$SystemPartitionDriveLetter,
        [Parameter(Mandatory = $true)]
        [string]$WindowsPartitionDriveLetter,
        [string]$RecoveryPartitionDriveLetter,
        [bool]$CreateRecoveryPartition = $true,
        [object[]]$AdditionalDataPartitions = @()
    )

    $assignments = [System.Collections.Generic.List[pscustomobject]]::new()
    $assignments.Add([pscustomobject]@{ Name = 'System'; Partition = $Layout.SystemPartition; DriveLetter = $SystemPartitionDriveLetter })
    $assignments.Add([pscustomobject]@{ Name = 'Windows'; Partition = $Layout.WindowsPartition; DriveLetter = $WindowsPartitionDriveLetter })
    if ($CreateRecoveryPartition) {
        $assignments.Add([pscustomobject]@{ Name = 'Recovery'; Partition = $Layout.RecoveryPartition; DriveLetter = $RecoveryPartitionDriveLetter })
    }
    for ($dataPartitionIndex = 0; $dataPartitionIndex -lt @($AdditionalDataPartitions).Count; $dataPartitionIndex++) {
        $assignments.Add([pscustomobject]@{
                Name        = "Data partition '$(@($AdditionalDataPartitions)[$dataPartitionIndex].Name)'"
                Partition   = $Layout.DataPartitions[$dataPartitionIndex].Partition
                DriveLetter = @($AdditionalDataPartitions)[$dataPartitionIndex].DriveLetter
            })
    }

    $resolvedAssignments = [System.Collections.Generic.List[pscustomobject]]::new()
    $changedAssignments = [System.Collections.Generic.List[pscustomobject]]::new()
    foreach ($assignment in $assignments) {
        $requestedDriveLetter = ([string]$assignment.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
        $currentDriveLetter = ([string]$assignment.Partition.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
        $resolvedAssignment = [pscustomobject]@{
            Name                 = $assignment.Name
            Partition            = $assignment.Partition
            RequestedDriveLetter = $requestedDriveLetter
            CurrentDriveLetter   = $currentDriveLetter
        }
        $resolvedAssignments.Add($resolvedAssignment)
        if ($currentDriveLetter -ne $requestedDriveLetter) {
            $changedAssignments.Add($resolvedAssignment)
        }

        $existingPartitions = @(Get-Partition -DriveLetter $requestedDriveLetter -ErrorAction SilentlyContinue)
        $externalPartition = $existingPartitions | Where-Object { $_.DiskNumber -ne $Layout.Disk.Number } | Select-Object -First 1
        if ($null -ne $externalPartition) {
            $externalVolume = @($externalPartition | Get-Volume -ErrorAction SilentlyContinue | Select-Object -First 1)
            $externalVolumeLabel = if ($externalVolume.Count -gt 0 -and -not [string]::IsNullOrWhiteSpace([string]$externalVolume[0].FileSystemLabel)) { "'$($externalVolume[0].FileSystemLabel)'" } else { 'an unlabeled volume' }
            throw "Cannot assign drive ${requestedDriveLetter}: to $($assignment.Name) because it is owned by $externalVolumeLabel on disk $($externalPartition.DiskNumber), partition $($externalPartition.PartitionNumber)."
        }

        $existingFileSystemDrive = Get-PSDrive -Name $requestedDriveLetter -PSProvider FileSystem -ErrorAction SilentlyContinue
        if ($null -ne $existingFileSystemDrive -and $existingPartitions.Count -eq 0) {
            throw "Cannot assign drive ${requestedDriveLetter}: to $($assignment.Name) because it is already mapped to '$($existingFileSystemDrive.Root)'."
        }
    }

    if ($changedAssignments.Count -eq 0) {
        WriteLog "Build partition drive letters already match the requested layout on working VHDX disk $($Layout.Disk.Number)."
    }
    else {
        # Remove changed access paths first so swaps cannot collide with old VHDX assignments.
        foreach ($assignment in $changedAssignments) {
            if ($assignment.CurrentDriveLetter -match '^[A-Z]$') {
                WriteLog "Removing drive $($assignment.CurrentDriveLetter): from $($assignment.Name) on working VHDX disk $($Layout.Disk.Number)."
                Remove-PartitionAccessPath -DiskNumber $assignment.Partition.DiskNumber -PartitionNumber $assignment.Partition.PartitionNumber -AccessPath "$($assignment.CurrentDriveLetter):\" -ErrorAction Stop
            }
        }

        foreach ($assignment in $changedAssignments) {
            WriteLog "Assigning drive $($assignment.RequestedDriveLetter): to $($assignment.Name) on working VHDX disk $($Layout.Disk.Number)."
            Set-Partition -DiskNumber $assignment.Partition.DiskNumber -PartitionNumber $assignment.Partition.PartitionNumber -NewDriveLetter $assignment.RequestedDriveLetter -ErrorAction Stop
        }
    }

    foreach ($assignment in $resolvedAssignments) {
        $verifiedPartition = Get-Partition -DiskNumber $assignment.Partition.DiskNumber -PartitionNumber $assignment.Partition.PartitionNumber -ErrorAction Stop
        $verifiedDriveLetter = ([string]$verifiedPartition.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
        if ($verifiedDriveLetter -ne $assignment.RequestedDriveLetter) {
            throw "Drive letter verification failed for $($assignment.Name) on working VHDX disk $($Layout.Disk.Number)."
        }
    }

    return Resolve-VhdxPartitionLayout -Disk $Layout.Disk -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $AdditionalDataPartitions
}

function New-FFUDataPartitionDriveLetterManifest {
    param(
        [Parameter(Mandatory = $true)]
        [object]$Layout,
        [Parameter(Mandatory = $true)]
        [object[]]$AdditionalDataPartitions,
        [Parameter(Mandatory = $true)]
        [string]$WindowsArch,
        [Parameter(Mandatory = $true)]
        [string]$ManifestPath,
        [ValidateSet('BuildVm', 'Deployment')]
        [string]$ManifestPurpose = 'Deployment',
        [bool]$Optimize = $false,
        [int]$OptimizeFFUPartitionNumber = 0
    )

    $resizablePartitionNumber = 0
    if ($Optimize) {
        $resizablePartitionNumber = Get-FFUOptimizePartitionNumber -Layout $Layout -RequestedPartitionNumber $OptimizeFFUPartitionNumber -AdditionalDataPartitions $AdditionalDataPartitions
    }

    $reservedDeploymentDriveLetters = [System.Collections.Generic.HashSet[string]]::new([System.StringComparer]::OrdinalIgnoreCase)
    if ($ManifestPurpose -eq 'Deployment') {
        foreach ($dataPartitionConfig in $AdditionalDataPartitions) {
            if ([bool]$dataPartitionConfig.PersistDriveLetter) {
                $null = $reservedDeploymentDriveLetters.Add([string]$dataPartitionConfig.DriveLetter)
            }
        }
    }

    $manifestPartitions = [System.Collections.Generic.List[pscustomobject]]::new()
    for ($dataPartitionIndex = 0; $dataPartitionIndex -lt $AdditionalDataPartitions.Count; $dataPartitionIndex++) {
        $dataPartitionConfig = $AdditionalDataPartitions[$dataPartitionIndex]
        $assignmentMode = 'Configured'
        $requestedDriveLetter = [string]$dataPartitionConfig.DriveLetter
        if ($ManifestPurpose -eq 'Deployment' -and -not [bool]$dataPartitionConfig.PersistDriveLetter) {
            $assignmentMode = 'Automatic'
            $requestedDriveLetter = $null
            foreach ($driveLetterCode in ([int][char]'D')..([int][char]'Z')) {
                $candidateDriveLetter = [string][char]$driveLetterCode
                if ($reservedDeploymentDriveLetters.Add($candidateDriveLetter)) {
                    $requestedDriveLetter = $candidateDriveLetter
                    break
                }
            }
            if ([string]::IsNullOrWhiteSpace($requestedDriveLetter)) {
                throw "No deployment drive letter is available for data partition '$($dataPartitionConfig.Name)'."
            }
        }

        $resolvedDataPartition = $Layout.DataPartitions[$dataPartitionIndex]
        $partitionGuid = [string]$resolvedDataPartition.Partition.Guid
        if ([string]::IsNullOrWhiteSpace($partitionGuid)) {
            WriteLog "GPT partition GUID was unavailable for data partition '$($dataPartitionConfig.Name)'. The runtime will require strict ordered fallback validation."
        }

        $manifestPartitions.Add([pscustomobject][ordered]@{
                Name                 = [string]$dataPartitionConfig.Name
                RequestedDriveLetter = $requestedDriveLetter
                ConfiguredDriveLetter = [string]$dataPartitionConfig.DriveLetter
                AssignmentMode       = $assignmentMode
                DataOrdinal          = $dataPartitionIndex + 1
                PartitionNumber      = [int]$resolvedDataPartition.Partition.PartitionNumber
                PartitionGuid        = $partitionGuid
                Label                = [string]$dataPartitionConfig.Label
                FileSystem           = [string]$dataPartitionConfig.FileSystem
                SizeBytes            = [uint64]$resolvedDataPartition.Partition.Size
                AllowSizeChange      = ($Optimize -and ([int]$resolvedDataPartition.Partition.PartitionNumber -eq $resizablePartitionNumber))
            })
    }

    if ($manifestPartitions.Count -eq 0) {
        throw 'Cannot create a data partition drive-letter manifest without data partitions.'
    }

    $manifestDirectory = Split-Path -Path $ManifestPath -Parent
    New-Item -Path $manifestDirectory -ItemType Directory -Force | Out-Null
    $manifest = [pscustomobject][ordered]@{
        SchemaVersion         = 1
        ProcessorArchitecture = if ($WindowsArch -ieq 'arm64') { 'arm64' } else { 'amd64' }
        Partitions            = $manifestPartitions.ToArray()
    }
    $manifest | ConvertTo-Json -Depth 5 | Set-Content -LiteralPath $ManifestPath -Encoding UTF8 -Force
    WriteLog "Created data partition drive-letter manifest at $ManifestPath with $($manifestPartitions.Count) partition(s)."

    return $manifest
}

function Add-FFUDataPartitionDriveLetterArtifacts {
    param(
        [Parameter(Mandatory = $true)]
        [string]$WindowsPartitionRoot,
        [Parameter(Mandatory = $true)]
        [object]$Layout,
        [Parameter(Mandatory = $true)]
        [object[]]$AdditionalDataPartitions,
        [Parameter(Mandatory = $true)]
        [string]$WindowsArch,
        [ValidateSet('BuildVm', 'Deployment')]
        [string]$ManifestPurpose = 'Deployment',
        [bool]$Optimize = $false,
        [int]$OptimizeFFUPartitionNumber = 0,
        [Parameter(Mandatory = $true)]
        [string]$FFUDevelopmentPath
    )

    $runtimeSourcePath = Join-Path -Path $FFUDevelopmentPath -ChildPath 'BuildFFUUnattend\Set-FFUDataPartitionDriveLetters.ps1'
    if (-not (Test-Path -LiteralPath $runtimeSourcePath -PathType Leaf)) {
        throw "Data partition drive-letter runtime script was not found at $runtimeSourcePath."
    }

    $runtimeDirectory = Join-Path -Path $WindowsPartitionRoot -ChildPath 'Windows\Setup\Scripts\FFUDL'
    $runtimeScriptPath = Join-Path -Path $runtimeDirectory -ChildPath 'Apply.ps1'
    $manifestPath = Join-Path -Path $runtimeDirectory -ChildPath 'Manifest.json'
    New-Item -Path $runtimeDirectory -ItemType Directory -Force | Out-Null
    Copy-Item -LiteralPath $runtimeSourcePath -Destination $runtimeScriptPath -Force
    $null = New-FFUDataPartitionDriveLetterManifest -Layout $Layout -AdditionalDataPartitions $AdditionalDataPartitions -WindowsArch $WindowsArch -ManifestPath $manifestPath -ManifestPurpose $ManifestPurpose -Optimize $Optimize -OptimizeFFUPartitionNumber $OptimizeFFUPartitionNumber
    Remove-Item -LiteralPath (Join-Path $runtimeDirectory 'Audit.success'), (Join-Path $runtimeDirectory 'Audit.failure'), (Join-Path $runtimeDirectory 'Specialize.failure') -Force -ErrorAction SilentlyContinue

    if (-not (Test-Path -LiteralPath $runtimeScriptPath -PathType Leaf) -or -not (Test-Path -LiteralPath $manifestPath -PathType Leaf)) {
        throw "Failed to stage data partition drive-letter runtime artifacts under $runtimeDirectory."
    }
    WriteLog "Staged data partition drive-letter runtime artifacts under $runtimeDirectory."

    return [pscustomobject]@{
        RuntimeDirectory = $runtimeDirectory
        RuntimeScriptPath = $runtimeScriptPath
        ManifestPath      = $manifestPath
    }
}

function Add-FFUDataPartitionDriveLetterCommandToUnattend {
    param(
        [Parameter(Mandatory = $true)]
        [string]$UnattendPath,
        [Parameter(Mandatory = $true)]
        [string]$ProcessorArchitecture
    )

    $unattendNamespace = 'urn:schemas-microsoft-com:unattend'
    $wcmNamespace = 'http://schemas.microsoft.com/WMIConfig/2002/State'
    $resolvedProcessorArchitecture = if ($ProcessorArchitecture -ieq 'arm64') { 'arm64' } else { 'amd64' }
    $unattendXml = New-Object System.Xml.XmlDocument
    $unattendXml.PreserveWhitespace = $true

    if (Test-Path -LiteralPath $UnattendPath -PathType Leaf) {
        $unattendXml.Load($UnattendPath)
    }
    else {
        $unattendDirectory = Split-Path -Path $UnattendPath -Parent
        New-Item -Path $unattendDirectory -ItemType Directory -Force | Out-Null
        $xmlDeclaration = $unattendXml.CreateXmlDeclaration('1.0', 'utf-8', $null)
        $null = $unattendXml.AppendChild($xmlDeclaration)
        $unattendRoot = $unattendXml.CreateElement('unattend', $unattendNamespace)
        $null = $unattendXml.AppendChild($unattendRoot)
    }

    $unattendRoot = $unattendXml.DocumentElement
    if ($null -eq $unattendRoot -or $unattendRoot.LocalName -ne 'unattend' -or $unattendRoot.NamespaceURI -ne $unattendNamespace) {
        throw "Unattend XML at $UnattendPath does not use the supported unattend root namespace."
    }

    $namespaceManager = New-Object System.Xml.XmlNamespaceManager($unattendXml.NameTable)
    $namespaceManager.AddNamespace('un', $unattendNamespace)
    $specializeSettings = $unattendRoot.SelectSingleNode("un:settings[@pass='specialize']", $namespaceManager)
    if ($null -eq $specializeSettings) {
        $specializeSettings = $unattendXml.CreateElement('settings', $unattendNamespace)
        $null = $specializeSettings.SetAttribute('pass', 'specialize')
        $firstSettingsNode = $unattendRoot.SelectSingleNode('un:settings', $namespaceManager)
        if ($null -ne $firstSettingsNode) {
            $null = $unattendRoot.InsertBefore($specializeSettings, $firstSettingsNode)
        }
        else {
            $null = $unattendRoot.AppendChild($specializeSettings)
        }
    }

    $deploymentComponents = @($specializeSettings.SelectNodes("un:component[@name='Microsoft-Windows-Deployment']", $namespaceManager) |
        Where-Object { $_.GetAttribute('processorArchitecture') -ieq $resolvedProcessorArchitecture })
    if ($deploymentComponents.Count -gt 1) {
        throw "Unattend XML at $UnattendPath contains multiple Microsoft-Windows-Deployment components for $resolvedProcessorArchitecture."
    }
    if ($deploymentComponents.Count -eq 0) {
        $deploymentComponent = $unattendXml.CreateElement('component', $unattendNamespace)
        $null = $deploymentComponent.SetAttribute('name', 'Microsoft-Windows-Deployment')
        $null = $deploymentComponent.SetAttribute('processorArchitecture', $resolvedProcessorArchitecture)
        $null = $deploymentComponent.SetAttribute('publicKeyToken', '31bf3856ad364e35')
        $null = $deploymentComponent.SetAttribute('language', 'neutral')
        $null = $deploymentComponent.SetAttribute('versionScope', 'nonSxS')
        $null = $specializeSettings.AppendChild($deploymentComponent)
    }
    else {
        $deploymentComponent = $deploymentComponents[0]
    }

    $runSynchronous = $deploymentComponent.SelectSingleNode('un:RunSynchronous', $namespaceManager)
    if ($null -eq $runSynchronous) {
        $runSynchronous = $unattendXml.CreateElement('RunSynchronous', $unattendNamespace)
        $null = $deploymentComponent.AppendChild($runSynchronous)
    }

    $existingCommands = @($runSynchronous.SelectNodes('un:RunSynchronousCommand', $namespaceManager))
    $orderedExistingCommands = [System.Collections.Generic.List[pscustomobject]]::new()
    $commandIndex = 0
    foreach ($existingCommand in $existingCommands) {
        $pathNode = $existingCommand.SelectSingleNode('un:Path', $namespaceManager)
        if ($null -ne $pathNode -and $pathNode.InnerText -match '(?i)\\FFUDL\\Apply\.ps1') {
            $null = $runSynchronous.RemoveChild($existingCommand)
            continue
        }

        $orderValue = [int]::MaxValue
        $orderNode = $existingCommand.SelectSingleNode('un:Order', $namespaceManager)
        if ($null -ne $orderNode) {
            $parsedOrder = 0
            if ([int]::TryParse($orderNode.InnerText, [ref]$parsedOrder)) {
                $orderValue = $parsedOrder
            }
        }
        $orderedExistingCommands.Add([pscustomobject]@{ Node = $existingCommand; Order = $orderValue; Index = $commandIndex })
        $commandIndex++
    }

    $nextOrder = 2
    foreach ($existingCommandInfo in @($orderedExistingCommands | Sort-Object -Property Order, Index)) {
        $orderNode = $existingCommandInfo.Node.SelectSingleNode('un:Order', $namespaceManager)
        if ($null -eq $orderNode) {
            $orderNode = $unattendXml.CreateElement('Order', $unattendNamespace)
            $null = $existingCommandInfo.Node.PrependChild($orderNode)
        }
        $orderNode.InnerText = [string]$nextOrder
        $nextOrder++
    }

    $newCommand = $unattendXml.CreateElement('RunSynchronousCommand', $unattendNamespace)
    $actionAttribute = $unattendXml.CreateAttribute('wcm', 'action', $wcmNamespace)
    $actionAttribute.Value = 'add'
    $null = $newCommand.Attributes.Append($actionAttribute)
    $orderElement = $unattendXml.CreateElement('Order', $unattendNamespace)
    $orderElement.InnerText = '1'
    $null = $newCommand.AppendChild($orderElement)
    $descriptionElement = $unattendXml.CreateElement('Description', $unattendNamespace)
    $descriptionElement.InnerText = 'Apply FFU data partition drive letters'
    $null = $newCommand.AppendChild($descriptionElement)
    $pathElement = $unattendXml.CreateElement('Path', $unattendNamespace)
    $specializeCommand = 'C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe -NoProfile -ExecutionPolicy Bypass -WindowStyle Hidden -File "C:\Windows\Setup\Scripts\FFUDL\Apply.ps1" -ManifestPath "C:\Windows\Setup\Scripts\FFUDL\Manifest.json" -Phase Specialize'
    if ($specializeCommand.Length -gt 259) {
        throw "Data partition drive-letter specialize command exceeds the 259-character unattend Path limit. Length: $($specializeCommand.Length)."
    }
    $pathElement.InnerText = $specializeCommand
    $null = $newCommand.AppendChild($pathElement)
    $willRebootElement = $unattendXml.CreateElement('WillReboot', $unattendNamespace)
    $willRebootElement.InnerText = 'OnRequest'
    $null = $newCommand.AppendChild($willRebootElement)
    if ($null -ne $runSynchronous.FirstChild) {
        $null = $runSynchronous.InsertBefore($newCommand, $runSynchronous.FirstChild)
    }
    else {
        $null = $runSynchronous.AppendChild($newCommand)
    }

    $unattendXml.Save($UnattendPath)
    WriteLog "Merged data partition drive-letter specialize command into $UnattendPath."
}

function Test-FFUDataPartitionDriveLetterAuditResult {
    param(
        [Parameter(Mandatory = $true)]
        [string]$WindowsPartitionRoot
    )

    $runtimeDirectory = Join-Path -Path $WindowsPartitionRoot -ChildPath 'Windows\Setup\Scripts\FFUDL'
    $successMarkerPath = Join-Path -Path $runtimeDirectory -ChildPath 'Audit.success'
    $failureMarkerPath = Join-Path -Path $runtimeDirectory -ChildPath 'Audit.failure'
    $auditLogPath = Join-Path -Path $WindowsPartitionRoot -ChildPath 'Windows\Temp\FFUDataPartitionDriveLetters.log'
    if (Test-Path -LiteralPath $failureMarkerPath -PathType Leaf) {
        $failureMessage = Get-Content -LiteralPath $failureMarkerPath -Raw -ErrorAction SilentlyContinue
        throw "Build VM data partition drive-letter enforcement failed. $failureMessage"
    }
    if (-not (Test-Path -LiteralPath $successMarkerPath -PathType Leaf)) {
        throw 'Build VM data partition drive-letter enforcement did not create its audit success marker.'
    }

    Remove-Item -LiteralPath $successMarkerPath, $failureMarkerPath -Force -ErrorAction SilentlyContinue
    Remove-Item -LiteralPath $auditLogPath -Force -ErrorAction SilentlyContinue
    WriteLog 'Validated build VM data partition drive-letter enforcement and removed audit result files.'
}

function Get-FFUOptimizePartitionNumber {
    param(
        [Parameter(Mandatory = $true)]
        [object]$Layout,
        [int]$RequestedPartitionNumber = 0,
        [object[]]$AdditionalDataPartitions = @()
    )

    if ($RequestedPartitionNumber -gt 0) {
        return $RequestedPartitionNumber
    }

    $fillRemainingDataPartition = $null
    $fillRemainingDataPartitionIndex = -1
    $dataPartitionIndex = 0
    foreach ($dataPartition in @($AdditionalDataPartitions)) {
        if ($null -eq $dataPartition) { continue }

        $fillRemaining = $false
        if ($dataPartition.PSObject.Properties.Name -contains 'FillRemaining') {
            $fillRemaining = [System.Convert]::ToBoolean($dataPartition.FillRemaining)
        }

        if ($fillRemaining) {
            $fillRemainingDataPartition = $dataPartition
            $fillRemainingDataPartitionIndex = $dataPartitionIndex
            break
        }

        $dataPartitionIndex++
    }
    if ($null -ne $fillRemainingDataPartition) {
        $dataPartition = if ($fillRemainingDataPartitionIndex -ge 0 -and $fillRemainingDataPartitionIndex -lt $Layout.DataPartitions.Count) { $Layout.DataPartitions[$fillRemainingDataPartitionIndex].Partition } else { $null }
        if ($null -eq $dataPartition) {
            throw "Unable to resolve FillRemaining data partition '$($fillRemainingDataPartition.Name)' for FFU optimization."
        }

        $resolvedDriveLetter = ([string]$dataPartition.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
        $resolvedDriveLetterLog = if ([string]::IsNullOrWhiteSpace($resolvedDriveLetter)) { '' } else { ", drive ${resolvedDriveLetter}:" }
        WriteLog "Using FillRemaining data partition '$($fillRemainingDataPartition.Name)' (partition number $($dataPartition.PartitionNumber)$resolvedDriveLetterLog) for FFU optimization."
        return [int]$dataPartition.PartitionNumber
    }

    if ($null -eq $Layout.WindowsPartition) {
        throw 'Unable to resolve Windows partition for FFU optimization.'
    }

    WriteLog "Using DISM default Windows partition selection for FFU optimization. Resolved Windows partition number $($Layout.WindowsPartition.PartitionNumber)."
    return 0
}
#Add boot files
function Add-BootFiles {
    param(
        [Parameter(Mandatory = $true)]
        [string]$OsPartitionDriveLetter,
        [Parameter(Mandatory = $true)]
        [string]$SystemPartitionDriveLetter,
        [Parameter(Mandatory = $true)]
        [string]$AdkPath,
        [Parameter(Mandatory = $true)]
        [ValidateSet('x86', 'x64', 'arm64')]
        [string]$WindowsArch,
        [string]$FirmwareType = 'UEFI'
    )

    # Use the ADK copy of BCDBoot so the boot binaries come from the validated ADK toolset
    # instead of the local OS installation, which can differ based on Secure Boot servicing state.
    $bcdBootArchitecture = if ($WindowsArch -ieq 'arm64') { 'arm64' } else { 'amd64' }
    $bcdBootPath = Join-Path $AdkPath "Assessment and Deployment Kit\Deployment Tools\$bcdBootArchitecture\BCDBoot\bcdboot.exe"

    if (-not (Test-Path -Path $bcdBootPath)) {
        throw "ADK BCDBoot was not found at $bcdBootPath"
    }

    WriteLog "Adding boot files for `"$($OsPartitionDriveLetter):\Windows`" to System partition `"$($SystemPartitionDriveLetter):`" using ADK BCDBoot at `"$bcdBootPath`"..."
    Invoke-Process $bcdBootPath "$($OsPartitionDriveLetter):\Windows /S $($SystemPartitionDriveLetter): /F $FirmwareType" | Out-Null
    WriteLog "Done."
}

function Enable-WindowsFeaturesByName {
    param (
        [Parameter(Mandatory = $true)]
        [string]$FeatureNames,
        [Parameter(Mandatory = $true)]
        [string]$Source
    )

    $FeaturesArray = $FeatureNames.Split(';')

    # Looping through each feature and enabling it
    foreach ($FeatureName in $FeaturesArray) {
        WriteLog "Enabling Windows Optional feature: $FeatureName"
        Enable-WindowsOptionalFeature -Path $WindowsPartition -FeatureName $FeatureName -All -Source $Source | Out-Null
        WriteLog "Done"
    }
}

#Dismount VHDX
function Dismount-ScratchVhdx {
    param(
        [Parameter(Mandatory = $true)]
        [string]$VhdxPath
    )

    if (Test-Path $VhdxPath) {
        WriteLog "Dismounting scratch VHDX..."
        Dismount-VHD -Path $VhdxPath
        WriteLog "Done."
    }
}

function New-FFUVM {
    #Create new Gen2 VM
    $VM = New-VM -Name $VMName -Path $VMPath -MemoryStartupBytes $memory -VHDPath $VHDXPath -Generation 2
    Set-VMProcessor -VMName $VMName -Count $processors

    # Connect the VM to the requested switch only when the experimental networking flag is enabled.
    if ($EnableVMNetworking) {
        $primaryVmNetworkAdapter = Get-VMNetworkAdapter -VMName $VMName -ErrorAction SilentlyContinue | Select-Object -First 1
        if ($null -ne $primaryVmNetworkAdapter) {
            if ($primaryVmNetworkAdapter.SwitchName -eq $VMSwitchName) {
                WriteLog "VM '$VMName' is already connected to Hyper-V switch '$VMSwitchName'."
            }
            else {
                Connect-VMNetworkAdapter -VMNetworkAdapter $primaryVmNetworkAdapter -SwitchName $VMSwitchName -ErrorAction Stop
                WriteLog "Connected VM '$VMName' to Hyper-V switch '$VMSwitchName'."
            }
        }
        else {
            Add-VMNetworkAdapter -VMName $VMName -SwitchName $VMSwitchName -Name 'FFUNetworkAdapter' -ErrorAction Stop | Out-Null
            WriteLog "Added VM network adapter for '$VMName' on Hyper-V switch '$VMSwitchName'."
        }
    }
    else {
        WriteLog "VM networking is disabled for '$VMName'."
    }

    #Mount AppsISO
    Add-VMDvdDrive -VMName $VMName -Path $AppsISO
   
    #Set Hard Drive as boot device
    $VMHardDiskDrive = Get-VMHarddiskdrive -VMName $VMName 
    Set-VMFirmware -VMName $VMName -FirstBootDevice $VMHardDiskDrive
    Set-VM -Name $VMName -AutomaticCheckpointsEnabled $false -StaticMemory

    #Configure TPM
    New-HgsGuardian -Name $VMName -GenerateCertificates
    $owner = get-hgsguardian -Name $VMName
    $kp = New-HgsKeyProtector -Owner $owner -AllowUntrustedRoot
    Set-VMKeyProtector -VMName $VMName -KeyProtector $kp.RawData
    Enable-VMTPM -VMName $VMName

    #Connect to VM
    WriteLog "Starting vmconnect localhost $VMName"
    & vmconnect localhost "$VMName"

    #Start VM
    try{
        Start-VM -Name $VMName -ErrorAction Stop
    }
    catch{
        WriteLog "Error starting VM: $_"
        throw $_
    }

    return $VM
}

function Get-PrivateProfileString {
    param (
        [Parameter()]
        [string]$FileName,
        [Parameter()]
        [string]$SectionName,
        [Parameter()]
        [string]$KeyName
    )

    # Read key from an INF/INI file. Use a larger buffer and allow it to grow if needed.
    $bufferSize = 4096
    $maxBufferSize = 65536
    $sbuilder = $null
    $charsCopied = 0

    while ($true) {
        $sbuilder = [System.Text.StringBuilder]::new($bufferSize)
        $charsCopied = [Win32.Kernel32]::GetPrivateProfileString($SectionName, $KeyName, "", $sbuilder, [uint32]$sbuilder.Capacity, $FileName)

        if ([int]$charsCopied -lt ($sbuilder.Capacity - 1)) {
            break
        }

        if ($bufferSize -ge $maxBufferSize) {
            break
        }

        $bufferSize = [Math]::Min(($bufferSize * 2), $maxBufferSize)
    }

    return $sbuilder.ToString()
}

function Get-PrivateProfileSection {
    param (
        [Parameter()]
        [string]$FileName,
        [Parameter()]
        [string]$SectionName
    )
    # Read the requested section from an INF/INI file
    # Some INF sections can be large; grow the buffer to avoid truncated results
    $hashTable = @{}
    $bufferSize = 16384
    $buffer = $null
    $charsCopied = 0

    while ($true) {
        $buffer = [byte[]]::new($bufferSize)
        $charsCopied = [Win32.Kernel32]::GetPrivateProfileSection($SectionName, $buffer, $buffer.Length, $FileName)

        # No section found or no content
        if ($charsCopied -eq 0) {
            return $hashTable
        }

        # If the returned data is close to the buffer size, assume truncation and retry bigger
        if (($charsCopied -ge ($bufferSize - 2)) -and ($bufferSize -lt 1048576)) {
            $bufferSize = $bufferSize * 2
            continue
        }

        break
    }

    # Convert only the returned portion of the buffer (Unicode = 2 bytes per char)
    $sectionText = [System.Text.Encoding]::Unicode.GetString($buffer, 0, ($charsCopied * 2))
    $keyValues = $sectionText.TrimEnd("`0").Split("`0")

    foreach ($keyValue in $keyValues) {
        if (![string]::IsNullOrEmpty($keyValue)) {
            $parts = $keyValue -split "=", 2
            if ($parts.Count -eq 2) {
                $hashTable[$parts[0]] = $parts[1]
            }
        }
    }

    return $hashTable
}
    
function Get-AvailableDriveLetter {
    param(
        [string[]]$ExcludedDriveLetters = @()
    )

    # Get an unused drive letter for temporary ISO and SUBST mappings.
    $usedLetters = @((Get-PSDrive -PSProvider FileSystem).Name | ForEach-Object { $_.ToUpperInvariant() })
    $excludedLetters = @($ExcludedDriveLetters | ForEach-Object {
            ([string]$_).Trim().TrimEnd(':').ToUpperInvariant()
        } | Where-Object { $_ -match '^[A-Z]$' } | Select-Object -Unique)
    for ($ascii = [int][char]'Z'; $ascii -ge [int][char]'A'; $ascii--) {
        $candidate = [string][char]$ascii
        if ($usedLetters -notcontains $candidate -and $excludedLetters -notcontains $candidate) {
            return $candidate
        }
    }
    return $null
}
    
function New-DriverSubstMapping {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $true)]
        [string]$SourcePath
    )
    
    # Map a long driver source folder to a short drive root using SUBST
    $resolvedPath = (Resolve-Path -Path $SourcePath -ErrorAction Stop).Path
    $driveLetter = Get-AvailableDriveLetter
    if ($null -eq $driveLetter) {
        throw 'No drive letters are available for SUBST mapping.'
    }
    $driveName = "$driveLetter`:"
    $mappedPath = "$driveLetter`:\"
    WriteLog "Mapping driver folder '$resolvedPath' to $driveName with SUBST."
    $escapedPath = $resolvedPath -replace '"', '""'
    $arguments = "/c subst $driveName `"$escapedPath`""
    Invoke-Process -FilePath cmd.exe -ArgumentList $arguments
    return [PSCustomObject]@{
        DriveLetter = $driveLetter
        DriveName   = $driveName
        DrivePath   = $mappedPath
    }
}
    
function Remove-DriverSubstMapping {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $true)]
        [string]$DriveLetter
    )
    
    # Remove the temporary SUBST mapping
    $driveName = "$DriveLetter`:"
    WriteLog "Removing SUBST drive $driveName"
    try {
        $arguments = "/c subst $driveName /d"
        Invoke-Process -FilePath cmd.exe -ArgumentList $arguments
    }
    catch {
        WriteLog "Failed to remove SUBST drive $($driveName): $_"
    }
}

function Invoke-DismDriverInjectionWithSubstLoop {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $true)]
        [string]$ImagePath,

        [Parameter(Mandatory = $true)]
        [string]$DriverRoot
    )

    # Resolve input paths
    $resolvedImagePath = (Resolve-Path -Path $ImagePath -ErrorAction Stop).Path
    $resolvedDriverRoot = (Resolve-Path -Path $DriverRoot -ErrorAction Stop).Path

    # Discover INF files under the driver root
    WriteLog "Scanning for INF files under: $resolvedDriverRoot"
    $infFiles = Get-ChildItem -Path $resolvedDriverRoot -Filter '*.inf' -File -Recurse -ErrorAction SilentlyContinue
    if ($null -eq $infFiles -or $infFiles.Count -eq 0) {
        WriteLog "No INF files found under: $resolvedDriverRoot"
        return
    }

    # Determine the deepest stable folders we can map with SUBST (SUBST has its own max path constraints)
    # Strategy:
    # - Start at the INF parent folder
    # - If too long for SUBST, walk up until the path is short enough
    # - Deduplicate and avoid redundant child folders when a parent already covers them via DISM /Recurse
    $substTargetMaxLength = 240
    $candidateDirs = [System.Collections.Generic.HashSet[string]]::new([System.StringComparer]::OrdinalIgnoreCase)

    foreach ($infFile in $infFiles) {
        $candidateDir = Split-Path -Path $infFile.FullName -Parent

        while ($candidateDir.Length -gt $substTargetMaxLength) {
            $parentDir = Split-Path -Path $candidateDir -Parent
            if ([string]::IsNullOrWhiteSpace($parentDir) -or $parentDir -eq $candidateDir) {
                break
            }
            $candidateDir = $parentDir
        }

        if ($candidateDir.Length -gt $substTargetMaxLength) {
            WriteLog "Warning: Skipping INF folder due to SUBST length limit (len=$($candidateDir.Length)): $candidateDir"
            continue
        }

        [void]$candidateDirs.Add($candidateDir)
    }

    $sortedCandidates = $candidateDirs | Sort-Object Length, @{ Expression = { $_ }; Ascending = $true }
    $selectedDirs = [System.Collections.Generic.List[string]]::new()

    foreach ($candidateDir in $sortedCandidates) {
        $isCovered = $false
        foreach ($selectedDir in $selectedDirs) {
            if ($candidateDir.Equals($selectedDir, [System.StringComparison]::OrdinalIgnoreCase)) {
                $isCovered = $true
                break
            }

            $prefix = $selectedDir.TrimEnd('\') + '\'
            if ($candidateDir.StartsWith($prefix, [System.StringComparison]::OrdinalIgnoreCase)) {
                $isCovered = $true
                break
            }
        }

        if (-not $isCovered) {
            [void]$selectedDirs.Add($candidateDir)
        }
    }

    $infDirs = $selectedDirs | Sort-Object
    WriteLog "Driver injection will process $($infDirs.Count) SUBST-safe folders (candidateFolders=$($candidateDirs.Count), INF total=$($infFiles.Count), substMaxLen=$substTargetMaxLength)."

    # Use a single SUBST drive letter and reuse it in a loop (map -> dism -> unmap)
    $driveLetter = Get-AvailableDriveLetter
    if ($null -eq $driveLetter) {
        throw 'No drive letters are available for SUBST mapping.'
    }

    $driveName = "$driveLetter`:"
    $drivePath = "$driveLetter`:\"
    WriteLog "Using SUBST drive $driveName for driver injection loop."

    $currentIndex = 0
    foreach ($infDir in $infDirs) {
        $currentIndex++
        $escapedPath = $infDir -replace '"', '""'

        try {
            WriteLog "[$currentIndex/$($infDirs.Count)] Mapping '$infDir' to $driveName with SUBST."
            $mapArgs = "/c subst $driveName `"$escapedPath`""
            Invoke-Process -FilePath cmd.exe -ArgumentList $mapArgs | Out-Null

            # Inject drivers (do not use \\?\ with DISM)
            $dismArgs = @(
                "/Image:`"$resolvedImagePath`""
                '/Add-Driver'
                "/Driver:$drivePath"
                '/Recurse'
            )

            WriteLog "dism.exe $($dismArgs -join ' ')"
            Invoke-Process -FilePath dism.exe -ArgumentList $dismArgs | Out-Null
        }
        catch {
            WriteLog "Warning: Driver injection failed for '$infDir': $($_.Exception.Message)"
        }
        finally {
            try {
                WriteLog "Removing SUBST drive $driveName"
                $unmapArgs = "/c subst $driveName /d"
                Invoke-Process -FilePath cmd.exe -ArgumentList $unmapArgs | Out-Null
            }
            catch {
                WriteLog "Warning: Failed removing SUBST drive $($driveName): $($_.Exception.Message)"
            }
        }
    }

    WriteLog "Driver injection loop complete for $resolvedDriverRoot"
}
    
function Copy-Drivers {
    param (
        [Parameter()]
        [string]$Path,
        [Parameter()]
        [string]$Output
    )
    # Find more information about device classes here:
    # https://learn.microsoft.com/en-us/windows-hardware/drivers/install/system-defined-device-setup-classes-available-to-vendors
    # For now, included are system devices, scsi and raid controllers, keyboards, mice and HID devices for touch support
    # 4D36E97D-E325-11CE-BFC1-08002BE10318 = System devices
    # 4D36E97B-E325-11CE-BFC1-08002BE10318 = SCSI, RAID, and NVMe Controllers
    # 4d36e96b-e325-11ce-bfc1-08002be10318 = Keyboards
    # 4d36e96f-e325-11ce-bfc1-08002be10318 = Mice and other pointing devices
    # 745a17a0-74d3-11d0-b6fe-00a0c90f57da = Human Interface Devices
    $filterGUIDs = @("{4D36E97D-E325-11CE-BFC1-08002BE10318}", "{4D36E97B-E325-11CE-BFC1-08002BE10318}", "{4d36e96b-e325-11ce-bfc1-08002be10318}", "{4d36e96f-e325-11ce-bfc1-08002be10318}", "{745a17a0-74d3-11d0-b6fe-00a0c90f57da}")
    $exclusionList = "wdmaudio.inf|Sound|Machine Learning|Camera|Firmware"

    # Log start and validate paths
    WriteLog "Copying PE drivers from '$Path' to '$Output' (WindowsArch: $WindowsArch)"
    if (-not (Test-Path -Path $Path)) {
        WriteLog "ERROR: Drivers source path not found: $Path"
        return
    }
    [void](New-Item -Path $Output -ItemType Directory -Force)
    
    $driverSourcePath = $Path
    $pathLength = $Path.Length
    
    # Determine common arch-specific SourceDisksFiles section names
    # Many INFs use 'amd64' rather than 'x64' for 64-bit paths (e.g. SourceDisksFiles.amd64)
    $sourceDisksFileSections = @("SourceDisksFiles")
    if ($WindowsArch -eq 'x64') {
        $sourceDisksFileSections += "SourceDisksFiles.amd64"
    }
    elseif ($WindowsArch -eq 'arm64') {
        $sourceDisksFileSections += "SourceDisksFiles.arm64"
    }
    
    $infFiles = Get-ChildItem -Path $Path -Recurse -Filter "*.inf"
    WriteLog "Found $($infFiles.Count) INF files under: $driverSourcePath"
    
    $matchedInfCount = 0
    $skippedInfCount = 0
    $copiedFileCount = 0
    $errorCount = 0
    
    for ($i = 0; $i -lt $infFiles.Count; $i++) {
        $infFullName = $infFiles[$i].FullName
        # Add long path prefix to handle long paths
        $longInfFullName = "\\?\$infFullName"
        $infPath = Split-Path -Path $infFullName
        $childPath = $infPath.Substring($pathLength).TrimStart('\')
        $targetPath = Join-Path -Path $Output -ChildPath $childPath
    
        # Log the INF files found
        WriteLog "Examining PE driver INF ($($i + 1)/$($infFiles.Count)): $infFullName"
    
        # Filter to known device classes
        # Some INFs include trailing comments after the value (e.g. "{GUID} ; TODO: ..."), so normalize to the GUID token only.
        $classGuidRaw = Get-PrivateProfileString -FileName $longInfFullName -SectionName "version" -KeyName "ClassGUID"
        $classGuid = $classGuidRaw
        if (-not [string]::IsNullOrWhiteSpace($classGuid)) {
            # Remove any trailing ';' comment and trim whitespace
            $classGuid = ($classGuid -split ';', 2)[0].Trim()

            # Extract the GUID token if the value contains other text
            if ($classGuid -match '\{[0-9A-Fa-f\-]{36}\}') {
                $classGuid = $matches[0]
            }
        }

        # WriteLog "ClassGUID: $classGuid"
        if ($classGuid -notin $filterGUIDs) {
            # WriteLog "Skipping PE driver INF due to GUID: $infFullName"
            $skippedInfCount++
            continue
        }
    
        # Avoid drivers that reference keywords from the exclusion list to keep the total size small
        if (((Get-Content -Path $infFullName) -match $exclusionList).Length -ne 0) {
            WriteLog "Skipping PE driver INF due to exclusion match: $infFullName"
            $skippedInfCount++
            continue
        }
    
        $matchedInfCount++
    
        # Log the INF being processed
        $providerName = (Get-PrivateProfileString -FileName $longInfFullName -SectionName "Version" -KeyName "Provider").Trim("%")
        if ([string]::IsNullOrWhiteSpace($providerName)) {
            $providerName = "Unknown Provider"
        }
    
        WriteLog "Processing PE driver INF: $infFullName"
        WriteLog "Provider: $providerName | ClassGUID: $classGuid"
        WriteLog "Target folder: $targetPath"
    
        [void](New-Item -Path $targetPath -ItemType Directory -Force)
    
        # Copy the INF itself
        try {
            Copy-Item -LiteralPath "$infFullName" -Destination "$targetPath" -Force -ErrorAction Stop
            $copiedFileCount++
            WriteLog "Copied: $infFullName -> $targetPath"
        }
        catch {
            $errorCount++
            WriteLog "ERROR: Failed to copy INF '$infFullName' to '$targetPath': $($_.Exception.Message)"
        }
    
        # Copy the catalog file (if specified)
        $CatalogFileName = Get-PrivateProfileString -FileName $longInfFullName -SectionName "version" -KeyName "Catalogfile"
        if (-not [string]::IsNullOrWhiteSpace($CatalogFileName)) {
            $catalogSource = Join-Path -Path $infPath -ChildPath $CatalogFileName
            if (Test-Path -Path $catalogSource) {
                try {
                    Copy-Item -LiteralPath "$catalogSource" -Destination "$targetPath" -Force -ErrorAction Stop
                    $copiedFileCount++
                    WriteLog "Copied: $catalogSource -> $targetPath"
                }
                catch {
                    $errorCount++
                    WriteLog "ERROR: Failed to copy catalog '$catalogSource' to '$targetPath': $($_.Exception.Message)"
                }
            }
            else {
                $errorCount++
                WriteLog "ERROR: Catalog file not found: $catalogSource (INF: $infFullName)"
            }
        }
        else {
            WriteLog "WARNING: No CatalogFile entry found in INF: $infFullName"
        }
    
        # Copy all files referenced by SourceDisksFiles sections
        foreach ($sectionName in $sourceDisksFileSections) {
            $sourceDiskFiles = Get-PrivateProfileSection -FileName $longInfFullName -SectionName $sectionName
            if ($sourceDiskFiles.Count -eq 0) {
                continue
            }
    
            WriteLog "Copying files from INF section [$sectionName] ($($sourceDiskFiles.Count) entries)"
    
            foreach ($sourceDiskFile in $sourceDiskFiles.Keys) {
                # Determine if the file lives in a subfolder relative to the INF path
                $rawValue = $sourceDiskFiles[$sourceDiskFile]
                $subdir = ""
    
                if (($null -ne $rawValue) -and ($rawValue.Contains(","))) {
                    $splitParts = $rawValue -split ","
                    if ($splitParts.Count -ge 2) {
                        $subdir = $splitParts[1]
                    }
                }
    
                if ([string]::IsNullOrWhiteSpace($subdir)) {
                    $subdir = ""
                }
    
                # Build source and destination paths
                if ([string]::IsNullOrEmpty($subdir)) {
                    $sourceFilePath = Join-Path -Path $infPath -ChildPath $sourceDiskFile
                    $destinationFolder = $targetPath
                }
                else {
                    $sourceFolder = Join-Path -Path $infPath -ChildPath $subdir
                    $sourceFilePath = Join-Path -Path $sourceFolder -ChildPath $sourceDiskFile
                    $destinationFolder = Join-Path -Path $targetPath -ChildPath $subdir
                    [void](New-Item -Path $destinationFolder -ItemType Directory -Force)
                }
    
                # Copy with logging and error handling
                if (Test-Path -Path $sourceFilePath) {
                    try {
                        Copy-Item -LiteralPath "$sourceFilePath" -Destination "$destinationFolder" -Force -ErrorAction Stop
                        $copiedFileCount++
                        WriteLog "Copied: $sourceFilePath -> $destinationFolder"
                    }
                    catch {
                        $errorCount++
                        WriteLog "ERROR: Failed to copy '$sourceFilePath' to '$destinationFolder' (INF: $infFullName): $($_.Exception.Message)"
                    }
                }
                else {
                    $errorCount++
                    WriteLog "ERROR: Source file not found for [$sectionName] entry '$sourceDiskFile': $sourceFilePath (INF: $infFullName)"
                }
            }
        }
    }
    
    # Final summary
    WriteLog "PE driver copy summary: INF total=$($infFiles.Count) matched=$matchedInfCount skipped=$skippedInfCount filesCopied=$copiedFileCount errors=$errorCount"
}

function New-PEMedia {
    param ()
    #Need to use the Demployment and Imaging tools environment to create winPE media
    $DandIEnv = "$adkPath`Assessment and Deployment Kit\Deployment Tools\DandISetEnv.bat"
    $WinPEFFUPath = "$FFUDevelopmentPath\WinPE"

    If (Test-path -Path "$WinPEFFUPath") {
        WriteLog "Removing old WinPE path at $WinPEFFUPath"
        Remove-Item -Path "$WinPEFFUPath" -Recurse -Force | out-null
    }

    WriteLog "Copying WinPE files to $WinPEFFUPath"
    # Use & cmd over invoke-process as Invoke-process has issues with the winpe media folder copying all contents
    if ($WindowsArch -eq 'x64') {
        & cmd /c """$DandIEnv"" && copype amd64 $WinPEFFUPath" | Out-Null
        # Invoke-Process cmd "/c ""$DandIEnv"" && copype amd64 $WinPEFFUPath" | Out-Null
    }
    elseif ($WindowsArch -eq 'arm64') {
        & cmd /c """$DandIEnv"" && copype arm64 $WinPEFFUPath" | Out-Null
        # Invoke-Process cmd "/c ""$DandIEnv"" && copype arm64 $WinPEFFUPath" | Out-Null
    }
    WriteLog 'Files copied successfully'

    WriteLog 'Mounting WinPE media to add WinPE optional components'
    Mount-WindowsImage -ImagePath "$WinPEFFUPath\media\sources\boot.wim" -Index 1 -Path "$WinPEFFUPath\mount" | Out-Null
    WriteLog 'Mounting complete'

    $Packages = @(
        "WinPE-WMI.cab",
        "en-us\WinPE-WMI_en-us.cab",
        "WinPE-NetFX.cab",
        "en-us\WinPE-NetFX_en-us.cab",
        "WinPE-Scripting.cab",
        "en-us\WinPE-Scripting_en-us.cab",
        "WinPE-PowerShell.cab",
        "en-us\WinPE-PowerShell_en-us.cab",
        "WinPE-StorageWMI.cab",
        "en-us\WinPE-StorageWMI_en-us.cab",
        "WinPE-DismCmdlets.cab",
        "en-us\WinPE-DismCmdlets_en-us.cab"
    )

    if ($WindowsArch -eq 'x64') {
        $PackagePathBase = "$adkPath`Assessment and Deployment Kit\Windows Preinstallation Environment\amd64\WinPE_OCs\"
    }
    elseif ($WindowsArch -eq 'arm64') {
        $PackagePathBase = "$adkPath`Assessment and Deployment Kit\Windows Preinstallation Environment\arm64\WinPE_OCs\"
    }
    

    foreach ($Package in $Packages) {
        $PackagePath = Join-Path $PackagePathBase $Package
        WriteLog "Adding Package $Package"
        Add-WindowsPackage -Path "$WinPEFFUPath\mount" -PackagePath $PackagePath | Out-Null
        WriteLog "Adding package complete"
    }
    WriteLog "Copying $FFUDevelopmentPath\WinPEDeployFFUFiles\* to WinPE deploy media"
    Copy-Item -Path "$FFUDevelopmentPath\WinPEDeployFFUFiles\*" -Destination "$WinPEFFUPath\mount" -Recurse -Force | Out-Null
    WriteLog 'Copy complete'
    #If $CopyPEDrivers = $true, add drivers to WinPE media using dism
    if ($CopyPEDrivers) {
        if ($UseDriversAsPEDrivers) {
            WriteLog "UseDriversAsPEDrivers is set. Building WinPE driver set from Drivers folder (bypassing PEDrivers folder contents)."
            if (Test-Path -Path $PEDriversFolder) {
                try {
                    Remove-Item -Path (Join-Path $PEDriversFolder '*') -Recurse -Force -ErrorAction SilentlyContinue | Out-Null
                }
                catch {
                    WriteLog "Warning: Failed clearing existing PEDriversFolder contents: $($_.Exception.Message)"
                }
            }
            else {
                try {
                    New-Item -Path $PEDriversFolder -ItemType Directory -Force | Out-Null
                }
                catch {
                    WriteLog "Error: Failed to create PEDriversFolder at $PEDriversFolder - continuing may fail when adding drivers."
                }
            }
            WriteLog "Copying required WinPE drivers from Drivers folder"
            Copy-Drivers -Path $DriversFolder -Output $PEDriversFolder
        }
        else {
            WriteLog "Copying PE drivers from PEDrivers folder"
        }
        
        WriteLog "Adding drivers to WinPE media"
        try {
            $WinPEMount = "$WinPEFFUPath\Mount"

            # Inject drivers using deep SUBST mapping (reuse one drive letter and loop each INF folder)
            Invoke-DismDriverInjectionWithSubstLoop -ImagePath $WinPEMount -DriverRoot $PEDriversFolder
        }
        catch {
            WriteLog 'Some drivers failed to be added. This can be expected. Continuing.'
        }
        WriteLog "Adding drivers complete"
    }
    $WinPEISOFile = $DeployISO
    WriteLog 'Dismounting WinPE media' 
    Dismount-WindowsImage -Path "$WinPEFFUPath\mount" -Save | Out-Null
    WriteLog 'Dismount complete'
    #Make ISO
    if ($WindowsArch -eq 'x64') {
        $OSCDIMGPath = "$adkPath`Assessment and Deployment Kit\Deployment Tools\amd64\Oscdimg"
    }
    elseif ($WindowsArch -eq 'arm64') {
        $OSCDIMGPath = "$adkPath`Assessment and Deployment Kit\Deployment Tools\arm64\Oscdimg"
    }
    $OSCDIMG = "$OSCDIMGPath\oscdimg.exe"
    WriteLog "Creating WinPE ISO at $WinPEISOFile"
    # & "$OSCDIMG" -m -o -u2 -udfver102 -bootdata:2`#p0,e,b$OSCDIMGPath\etfsboot.com`#pEF,e,b$OSCDIMGPath\Efisys_noprompt.bin $WinPEFFUPath\media $FFUDevelopmentPath\$WinPEISOName | Out-null
    if ($WindowsArch -eq 'x64') {
        $OSCDIMGArgs = "-m -o -u2 -udfver102 -bootdata:2`#p0,e,b`"$OSCDIMGPath\etfsboot.com`"`#pEF,e,b`"$OSCDIMGPath\Efisys.bin`" `"$WinPEFFUPath\media`" `"$WinPEISOFile`""
    }
    elseif ($WindowsArch -eq 'arm64') {
        $OSCDIMGArgs = "-m -o -u2 -udfver102 -bootdata:1`#pEF,e,b`"$OSCDIMGPath\Efisys.bin`" `"$WinPEFFUPath\media`" `"$WinPEISOFile`""
    }
    Invoke-Process $OSCDIMG $OSCDIMGArgs | Out-Null
    WriteLog "ISO created successfully"
    WriteLog "Cleaning up $WinPEFFUPath"
    Remove-Item -Path "$WinPEFFUPath" -Recurse -Force
    WriteLog 'Cleanup complete'
    # Deferred cleanup of preserved driver model folders (only after WinPE Deploy media is created)
    if ($UseDriversAsPEDrivers -and $CompressDownloadedDriversToWim -and $CopyPEDrivers) {
        WriteLog "Beginning deferred cleanup of preserved driver model folders (UseDriversAsPEDrivers + compression scenario)."
        $removedCount = 0
        $skippedCount = 0
        if (Test-Path -Path $DriversFolder) {
            Get-ChildItem -Path $DriversFolder -Directory -ErrorAction SilentlyContinue | ForEach-Object {
                $makeDir = $_.FullName
                Get-ChildItem -Path $makeDir -Directory -ErrorAction SilentlyContinue | ForEach-Object {
                    $modelDir = $_.FullName
                    $markerFile = Join-Path -Path $modelDir -ChildPath '__PreservedForPEDrivers.txt'
                    $leaf = Split-Path -Path $modelDir -Leaf
                    $wimPath = Join-Path -Path $makeDir -ChildPath ($leaf + '.wim')
                    if ((Test-Path -Path $markerFile -PathType Leaf) -and (Test-Path -Path $wimPath -PathType Leaf)) {
                        try {
                            WriteLog "Removing preserved driver folder: $modelDir (WIM located at $wimPath)"
                            Remove-Item -Path $modelDir -Recurse -Force -ErrorAction Stop
                            $removedCount++
                        }
                        catch {
                            WriteLog "Warning: Failed to remove preserved folder $modelDir : $($_.Exception.Message)"
                            $skippedCount++
                        }
                    }
                    else {
                        $skippedCount++
                    }
                }
            }
            WriteLog "Deferred driver cleanup complete. Removed: $removedCount; Skipped: $skippedCount"
        }
        else {
            WriteLog "Drivers folder $DriversFolder not found during deferred cleanup."
        }
    }
}

function Optimize-FFUCaptureDrive {
    param (
        [string]$VhdxPath,
        [bool]$EnableVolumeRetrim = $false,
        [string]$SystemPartitionDriveLetter,
        [string]$WindowsPartitionDriveLetter,
        [string]$RecoveryPartitionDriveLetter,
        [bool]$CreateRecoveryPartition = $true,
        [object[]]$AdditionalDataPartitions = @()
    )
    try {
        # Resolve whether the VHDX is already attached and get the disk reference
        $vhdInfo = Get-VHD -Path $VhdxPath
        if ($vhdInfo.Attached) {
            WriteLog 'VHDX is already mounted; using existing attachment for volume optimization'
            $mountedDisk = Get-Disk -Number $vhdInfo.DiskNumber
        }
        else {
            WriteLog 'Mounting VHDX for volume optimization'
            $mountedDisk = Mount-VHD -Path $VhdxPath -Passthru | Get-Disk
        }

        $partitionLayout = Resolve-VhdxPartitionLayout -Disk $mountedDisk -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $AdditionalDataPartitions
        $partitionLayout = Set-VhdxBuildPartitionDriveLetters -Layout $partitionLayout -SystemPartitionDriveLetter $SystemPartitionDriveLetter -WindowsPartitionDriveLetter $WindowsPartitionDriveLetter -RecoveryPartitionDriveLetter $RecoveryPartitionDriveLetter -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $AdditionalDataPartitions
        $osPartition = $partitionLayout.WindowsPartition
        if ($null -eq $osPartition -or [string]::IsNullOrWhiteSpace($osPartition.DriveLetter)) {
            throw 'Unable to resolve Windows partition drive letter for VHDX optimization.'
        }

        # Optimize filesystem layout before compacting the VHDX file
        WriteLog 'Defragmenting Windows partition...'
        Optimize-Volume -DriveLetter $osPartition.DriveLetter -Defrag -NormalPriority
        WriteLog 'Performing slab consolidation on Windows partition...'
        Optimize-Volume -DriveLetter $osPartition.DriveLetter -SlabConsolidate -NormalPriority

        # Optionally issue TRIM to improve reclaimable blocks before full VHDX compact
        if ($EnableVolumeRetrim) {
            try {
                WriteLog 'Running retrim on Windows partition before VHD compaction...'
                Optimize-Volume -DriveLetter $osPartition.DriveLetter -ReTrim -NormalPriority
            }
            catch {
                WriteLog "Retrim skipped with error: $($_.Exception.Message)"
            }
        }

        # Dismount and perform full-mode VHD optimization offline
        WriteLog 'Dismounting VHDX'
        Dismount-ScratchVhdx -VhdxPath $VhdxPath
        WriteLog 'Mounting VHDX as read-only for optimization'
        Mount-VHD -Path $VhdxPath -NoDriveLetter -ReadOnly | Out-Null
        WriteLog 'Optimizing VHDX in full mode...'
        Optimize-VHD -Path $VhdxPath -Mode Full
        WriteLog 'Dismounting VHDX'
        Dismount-ScratchVhdx -VhdxPath $VhdxPath
    }
    catch {
        throw $_
    }
}

function Get-CaptureVhdContext {
    param(
        [Parameter(Mandatory = $true)]
        [string]$VhdxPath,
        [string]$SystemPartitionDriveLetter,
        [string]$WindowsPartitionDriveLetter,
        [string]$RecoveryPartitionDriveLetter,
        [bool]$CreateRecoveryPartition = $true,
        [object[]]$AdditionalDataPartitions = @()
    )

    WriteLog 'Resolving VHDX context for host-side FFU capture'

    $vhdInfo = Get-VHD -Path $VhdxPath
    if ($vhdInfo.Attached) {
        WriteLog 'VHDX is already mounted for capture'
        $captureDisk = Get-Disk -Number $vhdInfo.DiskNumber
    }
    else {
        WriteLog 'Mounting VHDX for capture'
        $captureDisk = Mount-VHD -Path $VhdxPath -Passthru | Get-Disk
    }

    $partitionLayout = Resolve-VhdxPartitionLayout -Disk $captureDisk -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $AdditionalDataPartitions
    $partitionLayout = Set-VhdxBuildPartitionDriveLetters -Layout $partitionLayout -SystemPartitionDriveLetter $SystemPartitionDriveLetter -WindowsPartitionDriveLetter $WindowsPartitionDriveLetter -RecoveryPartitionDriveLetter $RecoveryPartitionDriveLetter -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $AdditionalDataPartitions
    $captureOsPartition = $partitionLayout.WindowsPartition
    if ($null -eq $captureOsPartition) {
        throw 'Unable to resolve Windows partition for FFU capture.'
    }
    if ([string]::IsNullOrWhiteSpace($captureOsPartition.DriveLetter)) {
        throw 'Unable to resolve Windows partition drive letter for FFU capture.'
    }

    return [pscustomobject]@{
        Disk                   = $captureDisk
        OsPartition            = $captureOsPartition
        OsPartitionDriveLetter = $captureOsPartition.DriveLetter
        WindowsPartition       = "$($captureOsPartition.DriveLetter):\"
        Layout                 = $partitionLayout
    }
}

function Get-ShortenedWindowsSKU {
    param (
        [string]$WindowsSKU
    )
    $shortenedWindowsSKU = switch ($WindowsSKU) {
        'Core' { 'Home' }
        'Home' { 'Home' }
        'CoreN' { 'Home_N' }
        'Home N' { 'Home_N' }
        'CoreSingleLanguage' { 'Home_SL' }
        'Home Single Language' { 'Home_SL' }
        'Education' { 'Edu' }
        'EducationN' { 'Edu_N' }
        'Education N' { 'Edu_N' }
        'Professional' { 'Pro' }
        'Pro' { 'Pro' }
        'ProfessionalN' { 'Pro_N' }
        'Pro N' { 'Pro_N' }
        'ProfessionalEducation' { 'Pro_Edu' }
        'Pro Education' { 'Pro_Edu' }
        'ProfessionalEducationN' { 'Pro_Edu_N' }
        'Pro Education N' { 'Pro_Edu_N' }
        'ProfessionalWorkstation' { 'Pro_WKS' }
        'Pro for Workstations' { 'Pro_WKS' }
        'ProfessionalWorkstationN' { 'Pro_WKS_N' }
        'Pro N for Workstations' { 'Pro_WKS_N' }
        'Enterprise' { 'Ent' }
        'EnterpriseN' { 'Ent_N' }
        'Enterprise N' { 'Ent_N' }
        'Enterprise N LTSC' { 'Ent_N_LTSC' }
        'EnterpriseS' { 'Ent_LTSC' }
        'EnterpriseSN' { 'Ent_N_LTSC' }
        'Enterprise LTSC' { 'Ent_LTSC' }
        'Enterprise 2016 LTSB' { 'Ent_LTSC' }
        'Enterprise N 2016 LTSB' { 'Ent_N_LTSC' }
        'IoT Enterprise LTSC' { 'IoT_Ent_LTSC' }
        'IoTEnterpriseS' { 'IoT_Ent_LTSC' }
        'IoT Enterprise N LTSC' { 'IoT_Ent_N_LTSC' }
        'ServerStandard' { 'Srv_Std' }
        'Standard' { 'Srv_Std' }
        'ServerDatacenter' { 'Srv_Dtc' }
        'Datacenter' { 'Srv_Dtc' }
        'Standard (Desktop Experience)' { 'Srv_Std_DE' }
        'Datacenter (Desktop Experience)' { 'Srv_Dtc_DE' }  
    }
    return $shortenedWindowsSKU

}
function Get-FFUCaptureNamingInfo {
    param(
        [Parameter(Mandatory = $true)]
        [string]$ShortenedWindowsSKU,

        [Parameter(Mandatory = $true)]
        [int]$WindowsRelease,

        [Parameter(Mandatory = $true)]
        [string]$WindowsVersion,

        [Parameter(Mandatory = $true)]
        [string]$InstallationType,

        [bool]$IsWindows10LtscClient = $false
    )

    $windowsReleaseToken = [string]$WindowsRelease
    if ($InstallationType -eq 'Client') {
        if (($WindowsRelease -eq 10) -or $IsWindows10LtscClient) {
            $windowsReleaseToken = 'Win10'
        }
        else {
            $windowsReleaseToken = 'Win11'
        }
    }

    $defaultFilePrefix = if ($InstallationType -eq 'Client') { $windowsReleaseToken } else { "Server$WindowsRelease" }
    $buildDate = Get-Date -uformat %b%Y

    return [pscustomobject]@{
        WindowsReleaseToken = $windowsReleaseToken
        WindowsVersion      = $WindowsVersion
        BuildDate           = $buildDate
        DefaultFileName     = "$defaultFilePrefix`_$WindowsVersion`_$ShortenedWindowsSKU`_$buildDate.ffu"
        CaptureName         = "$windowsReleaseToken$WindowsVersion$ShortenedWindowsSKU"
    }
}

function New-FFUFileName {

    $ffuCaptureNamingInfo = Get-FFUCaptureNamingInfo -ShortenedWindowsSKU $shortenedWindowsSKU -WindowsRelease $WindowsRelease -WindowsVersion $WindowsVersion -InstallationType $installationType -IsWindows10LtscClient:$isWindows10LtscClient
    $resolvedFFUNameTemplate = $CustomFFUNameTemplate

    # Replace '{WindowsRelease}' with the Windows release (e.g., 10, 11, 2016, 2019, 2022, 2025)
    $resolvedFFUNameTemplate = $resolvedFFUNameTemplate -replace '{WindowsRelease}', $ffuCaptureNamingInfo.WindowsReleaseToken
    # Replace '{WindowsVersion}' with the Windows version (e.g., 1607, 1809, 21h2, 22h2, 23h2, 24h2, etc)
    $resolvedFFUNameTemplate = $resolvedFFUNameTemplate -replace '{WindowsVersion}', $ffuCaptureNamingInfo.WindowsVersion
    # Replace '{SKU}' with the SKU of the Windows image (e.g., Pro, Enterprise, etc.)
    $resolvedFFUNameTemplate = $resolvedFFUNameTemplate -replace '{SKU}', $shortenedWindowsSKU
    # Replace '{BuildDate}' with the current month and year (e.g., Jan2023)
    $resolvedFFUNameTemplate = $resolvedFFUNameTemplate -replace '{BuildDate}', $ffuCaptureNamingInfo.BuildDate
    # Replace '{yyyy}' with the current year in 4-digit format (e.g., 2023)
    $resolvedFFUNameTemplate = $resolvedFFUNameTemplate -replace '{yyyy}', (Get-Date -UFormat '%Y')
    # Replace '{MM}' with the current month in 2-digit format (e.g., 01 for January)
    $resolvedFFUNameTemplate = $resolvedFFUNameTemplate -creplace '{MM}', (Get-Date -UFormat '%m')
    # Replace '{dd}' with the current day of the month in 2-digit format (e.g., 05)
    $resolvedFFUNameTemplate = $resolvedFFUNameTemplate -replace '{dd}', (Get-Date -UFormat '%d')
    # Replace '{HH}' with the current hour in 24-hour format (e.g., 14 for 2 PM)
    $resolvedFFUNameTemplate = $resolvedFFUNameTemplate -creplace '{HH}', (Get-Date -UFormat '%H')
    # Replace '{hh}' with the current hour in 12-hour format (e.g., 02 for 2 PM)
    $resolvedFFUNameTemplate = $resolvedFFUNameTemplate -creplace '{hh}', (Get-Date -UFormat '%I')
    # Replace '{mm}' with the current minute in 2-digit format (e.g., 09)
    $resolvedFFUNameTemplate = $resolvedFFUNameTemplate -creplace '{mm}', (Get-Date -UFormat '%M')
    # Replace '{tt}' with the current AM/PM designator (e.g., AM or PM)
    $resolvedFFUNameTemplate = $resolvedFFUNameTemplate -replace '{tt}', (Get-Date -UFormat '%p')
    if ($resolvedFFUNameTemplate -notlike '*.ffu') {
        $resolvedFFUNameTemplate += '.ffu'
    }
    return $resolvedFFUNameTemplate
}

function New-FFU {
    $captureContext = Get-CaptureVhdContext -VhdxPath $VHDXPath -SystemPartitionDriveLetter $SystemPartitionDriveLetter -WindowsPartitionDriveLetter $WindowsPartitionDriveLetter -RecoveryPartitionDriveLetter $RecoveryPartitionDriveLetter -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions
    $captureDisk = $captureContext.Disk
    $resolvedFFUOptimizePartitionNumber = 0
    if ($Optimize -eq $true) {
        $resolvedFFUOptimizePartitionNumber = Get-FFUOptimizePartitionNumber -Layout $captureContext.Layout -RequestedPartitionNumber $OptimizeFFUPartitionNumber -AdditionalDataPartitions $normalizedAdditionalDataPartitions
    }
    $ffuCaptureNamingInfo = Get-FFUCaptureNamingInfo -ShortenedWindowsSKU $shortenedWindowsSKU -WindowsRelease $WindowsRelease -WindowsVersion $WindowsVersion -InstallationType $installationType -IsWindows10LtscClient:$isWindows10LtscClient

    try {
        Set-Progress -Percentage 68 -Message "Capturing FFU from VHDX..."

        WriteLog 'Creating FFU File Name'
        if ($CustomFFUNameTemplate) {
            $FFUFileName = New-FFUFileName
        }
        else {
            $FFUFileName = $ffuCaptureNamingInfo.DefaultFileName
        }
        WriteLog "FFU file name: $FFUFileName"
        $FFUFile = "$FFUCaptureLocation\$FFUFileName"
        WriteLog 'Capturing FFU from mounted VHDX on host'
        Invoke-Process cmd "/c ""$DandIEnv"" && dism /Capture-FFU /ImageFile:$FFUFile /CaptureDrive:\\.\PhysicalDrive$($captureDisk.DiskNumber) /Name:$($ffuCaptureNamingInfo.CaptureName) /Compress:Default" | Out-Null

        WriteLog 'FFU Capture complete'
    }
    finally {
        Dismount-ScratchVhdx -VhdxPath $VHDXPath
    }

    #Without this 120 second sleep, we sometimes see an error when mounting the FFU due to a file handle lock. Needed for both driver and optimize steps.
    
    If ($InstallDrivers -or $Optimize) {
        WriteLog 'Sleeping 2 minutes to prevent file handle lock'
        Start-Sleep 120
    }

    #Add drivers
    If ($InstallDrivers) {
        Set-Progress -Percentage 75 -Message "Injecting drivers into FFU..."
        WriteLog 'Adding drivers'
        WriteLog "Creating $FFUDevelopmentPath\Mount directory"
        New-Item -Path "$FFUDevelopmentPath\Mount" -ItemType Directory -Force | Out-Null
        WriteLog "Created $FFUDevelopmentPath\Mount directory"
        WriteLog "Mounting $FFUFile to $FFUDevelopmentPath\Mount"
        Mount-WindowsImage -ImagePath $FFUFile -Index 1 -Path "$FFUDevelopmentPath\Mount" | Out-null
        WriteLog 'Mounting complete'
        WriteLog 'Adding drivers - This will take a few minutes, please be patient'
        try {
            # Inject drivers using deep SUBST mapping (reuse one drive letter and loop each INF folder)
            Invoke-DismDriverInjectionWithSubstLoop -ImagePath "$FFUDevelopmentPath\Mount" -DriverRoot "$DriversFolder"
        }
        catch {
            WriteLog 'Some drivers failed to be added to the FFU. This can be expected. Continuing.'
        }
        WriteLog 'Adding drivers complete'
        WriteLog "Dismount $FFUDevelopmentPath\Mount"
        Dismount-WindowsImage -Path "$FFUDevelopmentPath\Mount" -Save | Out-Null
        WriteLog 'Dismount complete'
        WriteLog "Remove $FFUDevelopmentPath\Mount folder"
        Remove-Item -Path "$FFUDevelopmentPath\Mount" -Recurse -Force | Out-null
        WriteLog 'Folder removed'
    }
    #Optimize FFU
    if ($Optimize -eq $true) {
        Set-Progress -Percentage 85 -Message "Optimizing FFU..."
        WriteLog 'Optimizing FFU - This will take a few minutes, please be patient'
        #Need to use ADK version of DISM to address bug in DISM - perhaps Windows 11 24H2 will fix this
        $optimizePartitionArgument = if ($resolvedFFUOptimizePartitionNumber -gt 0) { " /PartitionNumber:$resolvedFFUOptimizePartitionNumber" } else { '' }
        if (-not [string]::IsNullOrWhiteSpace($optimizePartitionArgument)) {
            WriteLog "Optimizing FFU with DISM partition number $resolvedFFUOptimizePartitionNumber."
        }
        Invoke-Process cmd "/c ""$DandIEnv"" && dism /optimize-ffu /imagefile:$FFUFile$optimizePartitionArgument" | Out-Null
        #Invoke-Process cmd "/c dism /optimize-ffu /imagefile:$FFUFile" | Out-Null
        WriteLog 'Optimizing FFU complete'
        Set-Progress -Percentage 90 -Message "FFU post-processing complete."
    }
    

}
function Remove-FFUVM {
    param (
        [Parameter(Mandatory = $false)]
        [string]$VMName
    )
    #Get the VM object and remove the VM, the HGSGuardian, and the certs
    If ($VMName) {
        $FFUVM = get-vm $VMName | Where-Object { $_.state -ne 'running' }
    }   
    If ($null -ne $FFUVM) {
        WriteLog 'Cleaning up VM'
        $certPath = 'Cert:\LocalMachine\Shielded VM Local Certificates\'
        $VMName = $FFUVM.Name
        WriteLog "Removing VM: $VMName"
        Remove-VM -Name $VMName -Force
        WriteLog 'Removal complete'
        WriteLog "Removing $VMPath"
        Remove-Item -Path $VMPath -Force -Recurse
        WriteLog 'Removal complete'
        WriteLog "Removing HGSGuardian for $VMName" 
        Remove-HgsGuardian -Name $VMName -WarningAction SilentlyContinue
        WriteLog 'Removal complete'
        WriteLog 'Cleaning up HGS Guardian certs'
        $certs = Get-ChildItem -Path $certPath -Recurse | Where-Object { $_.Subject -like "*$VMName*" }
        foreach ($cert in $Certs) {
            Remove-item -Path $cert.PSPath -force | Out-Null
        }
        WriteLog 'Cert removal complete'
    }
    #If just building the FFU from vhdx, remove the vhdx path
    If (-not $InstallApps -and $vhdxDisk) {
        WriteLog 'Cleaning up VHDX'
        WriteLog "Removing $VMPath"
        Remove-Item -Path $VMPath -Force -Recurse | Out-Null
        WriteLog 'Removal complete'
    }

    #Remove orphaned mounted images
    $mountedImages = Get-WindowsImage -Mounted
    if ($mountedImages) {
        foreach ($image in $mountedImages) {
            $mountPath = $image.Path
            WriteLog "Dismounting image at $mountPath"
            Dismount-WindowsImage -Path $mountPath -discard
            WriteLog "Successfully dismounted image at $mountPath"
        }
    } 
    #Remove Mount folder if it exists
    If (Test-Path -Path $FFUDevelopmentPath\Mount) {
        WriteLog "Remove $FFUDevelopmentPath\Mount folder"
        Remove-Item -Path "$FFUDevelopmentPath\Mount" -Recurse -Force
        WriteLog 'Folder removed'
    }
    #Remove unused mountpoints
    WriteLog 'Remove unused mountpoints'
    Invoke-Process cmd "/c mountvol /r" | Out-Null
    WriteLog 'Removal complete'
}
Function Get-USBDrive {
    # Log the start of the USB drive check
    WriteLog 'Checking for USB drives'
    
    # Check if external hard disk media is allowed and user has not specified USB drives
    If ($AllowExternalHardDiskMedia -and (-not($USBDriveList))) {
        # Get all removable and external hard disk media drives
        [array]$USBDrives = (Get-CimInstance -ClassName Win32_DiskDrive -Filter "MediaType='Removable Media' OR MediaType='External hard disk media'")
        [array]$ExternalHardDiskDrives = $USBDrives | Where-Object { $_.MediaType -eq 'External hard disk media' }
        $ExternalCount = $ExternalHardDiskDrives.Count
        $USBDrivesCount = $USBDrives.Count
        
        # Check if user should be prompted for external hard disk media
        if ($PromptExternalHardDiskMedia) {
            if ($ExternalHardDiskDrives) {
                # Log and warn about found external hard disk media drives
                if ($VerbosePreference -ne 'Continue') {
                    Write-Warning 'Found external hard disk media drives'
                    Write-Warning 'Will prompt for user input to select the drive to use to prevent accidental data loss'
                    Write-Warning 'If you do not want to be prompted for this in the future, set -PromptExternalHardDiskMedia to $false'
                }
                WriteLog 'Found external hard disk media drives'
                WriteLog 'Will prompt for user input to select the drive to use to prevent accidental data loss'
                WriteLog 'If you do not want to be prompted for this in the future, set -PromptExternalHardDiskMedia to $false'
                
                # Prepare output for user selection
                $Output = @()
                for ($i = 0; $i -lt $ExternalHardDiskDrives.Count; $i++) {
                    $ExternalDiskNumber = $ExternalHardDiskDrives[$i].Index
                    $ExternalDisk = Get-Disk -Number $ExternalDiskNumber
                    $Index = $i + 1
                    $Name = $ExternalDisk.FriendlyName
                    $SerialNumber = $ExternalHardDiskDrives[$i].serialnumber
                    $PartitionStyle = $ExternalDisk.PartitionStyle
                    $Status = $ExternalDisk.OperationalStatus
                    $Properties = [ordered]@{
                        'Drive Number'    = $Index
                        'Drive Name'      = $Name
                        'Serial Number'   = $SerialNumber
                        'Partition Style' = $PartitionStyle
                        'Status'          = $Status
                    }
                    $Output += New-Object PSObject -Property $Properties
                }
                
                # Format and display the output
                $FormattedOutput = $Output | Format-Table -AutoSize -Property 'Drive Number', 'Drive Name', 'Serial Number', 'Partition Style', 'Status' | Out-String
                if ($VerbosePreference -ne 'Continue') {
                    $FormattedOutput | Out-Host
                }
                WriteLog $FormattedOutput
                
                # Prompt user to select a drive
                do {
                    $inputChoice = Read-Host "Enter the number corresponding to the external hard disk media drive you want to use"
                    if ($inputChoice -match '^\d+$') {
                        $inputChoice = [int]$inputChoice
                        if ($inputChoice -ge 1 -and $inputChoice -le $ExternalCount) {
                            $SelectedIndex = $inputChoice - 1
                            $ExternalDiskNumber = $ExternalHardDiskDrives[$SelectedIndex].Index
                            $ExternalDisk = Get-Disk -Number $ExternalDiskNumber
                            $USBDrives = $ExternalHardDiskDrives[$SelectedIndex]
                            $USBDrivesCount = $USBDrives.Count
                            if ($VerbosePreference -ne 'Continue') {
                                Write-Host "Drive $inputChoice was selected"
                            }
                            WriteLog "Drive $inputChoice was selected"
                        }
                        else {
                            # Handle invalid selection
                            if ($VerbosePreference -ne 'Continue') {
                                Write-Host "Invalid selection. Please try again."
                            }
                            WriteLog "Invalid selection. Please try again."
                        }
                        
                        # Check if the selected drive is offline
                        if ($ExternalDisk.OperationalStatus -eq 'Offline') {
                            if ($VerbosePreference -ne 'Continue') {
                                Write-Error "Selected Drive is in an Offline State. Please check the drive status in Disk Manager and try again."
                            }
                            WriteLog "Selected Drive is in an Offline State. Please check the drive status in Disk Manager and try again."
                            exit 1
                        }
                    }
                    else {
                        # Handle invalid input
                        if ($VerbosePreference -ne 'Continue') {
                            Write-Host "Invalid selection. Please try again."
                        }
                        WriteLog "Invalid selection. Please try again."
                    }
                } while ($null -eq $selectedIndex)
            }
        }
        else {
            # Log the count of found USB drives
            if ($VerbosePreference -ne 'Continue') {
                Write-Host "Found $USBDrivesCount total USB drives"
                If ($ExternalCount -gt 0) {
                    Write-Host "$ExternalCount external drives"
                }
            }
            WriteLog "Found $USBDrivesCount total USB drives"
            If ($ExternalCount -gt 0) {
                WriteLog "$ExternalCount external drives"
            }
        }
    }
    elseif ($USBDriveList) {
        # Log the count of specified USB drives
        # USBDriveList values can be a single UniqueId string, or an array of UniqueIds (multiple same-model drives)
        $USBDriveListCount = 0
        foreach ($model in $USBDriveList.Keys) {
            $USBDriveListCount += @($USBDriveList[$model]).Count
        }
        WriteLog "Looking for $USBDriveListCount USB drives from USB Drive List"

        # Get only the specified USB drives based on model and UniqueId
        $USBDrives = @()
        foreach ($model in $USBDriveList.Keys) {
            $configUniqueIds = @($USBDriveList[$model])

            foreach ($configUniqueId in $configUniqueIds) {
                if ([string]::IsNullOrWhiteSpace([string]$configUniqueId)) {
                    continue
                }

                WriteLog "Looking for USB drive model $model with UniqueId $configUniqueId"
                # First get candidate drives by model and media type
                $candidateDrives = Get-CimInstance -ClassName Win32_DiskDrive -Filter "Model LIKE '%$model%' AND (MediaType='Removable Media' OR MediaType='External hard disk media')"
                $foundDrive = $null
                foreach ($candidate in $candidateDrives) {
                    # Get the disk to retrieve UniqueId
                    $disk = Get-Disk -Number $candidate.Index -ErrorAction SilentlyContinue
                    if ($disk -and $disk.UniqueId) {
                        # Trim the machine name suffix (everything after the colon) from UniqueId
                        $diskUniqueId = if ($disk.UniqueId -match ':') {
                            $disk.UniqueId.Split(':')[0]
                        }
                        else {
                            $disk.UniqueId
                        }
                        # Match on the trimmed UniqueId
                        if ($diskUniqueId -eq $configUniqueId) {
                            $foundDrive = $candidate
                            break
                        }
                    }
                }
                if ($foundDrive) {
                    if ($USBDrives.Index -notcontains $foundDrive.Index) {
                        WriteLog "Found USB drive model $($foundDrive.Model) with UniqueId $configUniqueId"
                        $USBDrives += $foundDrive
                    }
                    else {
                        WriteLog "USB drive model $($foundDrive.Model) with UniqueId $configUniqueId was already added. Skipping duplicate."
                    }
                }
                else {
                    WriteLog "USB drive model $model with UniqueId $configUniqueId not found"
                }
            }
        }
        $USBDrivesCount = $USBDrives.Count
        WriteLog "Found $USBDrivesCount of $USBDriveListCount USB drives from USB Drive List"
    }
    else {
        # Get only removable media drives
        [array]$USBDrives = (Get-CimInstance -ClassName Win32_DiskDrive -Filter "MediaType='Removable Media'")
        $USBDrivesCount = $USBDrives.Count
        WriteLog "Found $USBDrivesCount Removable USB drives"
    }
    
    # Check if any USB drives were found
    if ($USBDrives.Count -eq 0) {
        WriteLog "No USB drive found. Exiting"
        Write-Error "No USB drive found. Exiting"
        exit 1
    }
    
    # Return the found USB drives and their count
    return $USBDrives, $USBDrivesCount
}
Function New-DeploymentUSB {
    param(
        [switch]$CopyFFU,
        [string[]]$FFUFilesToCopy
    )
    WriteLog "CopyFFU is set to $CopyFFU"
    $BuildUSBPath = $PSScriptRoot
    WriteLog "BuildUSBPath is $BuildUSBPath"

    $SelectedFFUFile = $null

    # 1. Get FFU File(s) - This happens once before parallel processing
    if ($CopyFFU.IsPresent) {
        if ($null -ne $FFUFilesToCopy -and $FFUFilesToCopy.Count -gt 0) {
            $SelectedFFUFile = $FFUFilesToCopy
            WriteLog "Using preselected FFU file list. Count: $($FFUFilesToCopy.Count)"
            WriteLog "FFU files to copy:"
            foreach ($f in $FFUFilesToCopy) {
                WriteLog ("- {0}" -f (Split-Path $f -Leaf))
            }
        }
        else {
            $FFUFiles = Get-ChildItem -Path "$BuildUSBPath\FFU" -Filter "*.ffu"
            $FFUCount = $FFUFiles.count
    
            switch ($FFUCount) {
                0 {
                    Write-Error "No FFU files found in $BuildUSBPath\FFU. Cannot copy FFU to USB drive."
                    return
                }
                1 {
                    $SelectedFFUFile = $FFUFiles.FullName
                    WriteLog "One FFU file found, will use: $SelectedFFUFile"
                }
                default {
                    WriteLog "Found $FFUCount FFU files"
                    if ($VerbosePreference -ne 'Continue') {
                        Write-Host "Found $FFUCount FFU files"
                    }
                    $output = @()
                    for ($i = 0; $i -lt $FFUCount; $i++) {
                        $output += [PSCustomObject]@{
                            'FFU Number'    = $i + 1
                            'FFU Name'      = $FFUFiles[$i].Name
                            'Last Modified' = $FFUFiles[$i].LastWriteTime
                        }
                    }
                    $output | Format-Table -AutoSize | Out-String | Write-Host
    
                    do {
                        $inputChoice = Read-Host "Enter the number for the FFU to copy, or 'A' for all"
                        if ($inputChoice -eq 'A') {
                            $SelectedFFUFile = $FFUFiles.FullName
                            WriteLog 'Will copy all FFU Files'
                        }
                        elseif ($inputChoice -match '^\d+$' -and [int]$inputChoice -ge 1 -and [int]$inputChoice -le $FFUCount) {
                            $SelectedFFUFile = $FFUFiles[[int]$inputChoice - 1].FullName
                            WriteLog "$SelectedFFUFile was selected"
                        }
                        else {
                            Write-Host "Invalid selection. Please try again."
                        }
                    } while ($null -eq $SelectedFFUFile)
                }
            }
        }
    }

    # Mount ISO once before the loop
    WriteLog "Mounting deployment ISO: $DeployISO"
    $ISOMountPoint = (Mount-DiskImage -ImagePath $DeployISO -PassThru | Get-Volume).DriveLetter + ":\"
    WriteLog "ISO mounted at $ISOMountPoint"

    # 2. Partition and format USB drives in parallel
    WriteLog "Starting parallel creation for $USBDrivesCount USB drive(s)."

    $resolvedUSBThrottle = if ($MaxUSBDrives -gt 0) { [math]::Min($MaxUSBDrives, $USBDrivesCount) } else { $USBDrivesCount }
    WriteLog "Using USB drive throttle limit: $resolvedUSBThrottle (MaxUSBDrives param: $MaxUSBDrives; Drives to process: $USBDrivesCount)"

    $USBDrives | ForEach-Object -Parallel {
        $USBDrive = $_
        
        # Import common module for logging in this thread
        Import-Module "$($using:PSScriptRoot)\FFU.Common" -Force
        Set-CommonCoreLogPath -Path $using:LogFile

        function Get-LocalUnattendSourcePath {
            param(
                [string]$UnattendFolder,
                [string]$WindowsArch,
                [string]$UnattendX64FilePath,
                [string]$UnattendArm64FilePath
            )

            $resolvedArch = if ($WindowsArch -ieq 'arm64') { 'arm64' } else { 'x64' }
            $resolvedSourcePath = if ($resolvedArch -eq 'arm64') {
                if ([string]::IsNullOrWhiteSpace($UnattendArm64FilePath)) {
                    Join-Path $UnattendFolder 'unattend_arm64.xml'
                }
                else {
                    $UnattendArm64FilePath
                }
            }
            else {
                if ([string]::IsNullOrWhiteSpace($UnattendX64FilePath)) {
                    Join-Path $UnattendFolder 'unattend_x64.xml'
                }
                else {
                    $UnattendX64FilePath
                }
            }

            WriteLog "Resolved unattend source path for ${resolvedArch}: $resolvedSourcePath"
            return $resolvedSourcePath
        }

        function Initialize-UnattendComputerNamePath {
            param(
                [xml]$UnattendXml,
                [string]$WindowsArch
            )

            $unattendRoot = $UnattendXml.DocumentElement
            if (($null -eq $unattendRoot) -or ($unattendRoot.LocalName -ne 'unattend')) {
                throw 'Unattend XML is missing the unattend root element.'
            }

            $unattendNamespace = $unattendRoot.NamespaceURI
            if ([string]::IsNullOrWhiteSpace($unattendNamespace)) {
                throw 'Unattend XML is missing the default unattend namespace.'
            }

            $namespaceManager = New-Object System.Xml.XmlNamespaceManager($UnattendXml.NameTable)
            $namespaceManager.AddNamespace('un', $unattendNamespace)

            $specializeSettings = $unattendRoot.SelectSingleNode("un:settings[@pass='specialize']", $namespaceManager)
            $createdSpecializeSettings = $false
            if ($null -eq $specializeSettings) {
                $specializeSettings = $UnattendXml.CreateElement('settings', $unattendNamespace)
                $null = $specializeSettings.SetAttribute('pass', 'specialize')
                $firstSettingsNode = $unattendRoot.SelectSingleNode('un:settings', $namespaceManager)
                if ($null -ne $firstSettingsNode) {
                    $null = $unattendRoot.InsertBefore($specializeSettings, $firstSettingsNode)
                }
                else {
                    $null = $unattendRoot.AppendChild($specializeSettings)
                }
                $createdSpecializeSettings = $true
            }

            $shellSetupComponent = $specializeSettings.SelectSingleNode("un:component[@name='Microsoft-Windows-Shell-Setup']", $namespaceManager)
            $createdShellSetupComponent = $false
            if ($null -eq $shellSetupComponent) {
                $processorArchitecture = if ($WindowsArch -ieq 'arm64') { 'arm64' } else { 'amd64' }
                $shellSetupComponent = $UnattendXml.CreateElement('component', $unattendNamespace)
                $null = $shellSetupComponent.SetAttribute('name', 'Microsoft-Windows-Shell-Setup')
                $null = $shellSetupComponent.SetAttribute('processorArchitecture', $processorArchitecture)
                $null = $shellSetupComponent.SetAttribute('publicKeyToken', '31bf3856ad364e35')
                $null = $shellSetupComponent.SetAttribute('language', 'neutral')
                $null = $shellSetupComponent.SetAttribute('versionScope', 'nonSxS')
                $null = $shellSetupComponent.SetAttribute('xmlns:wcm', 'http://www.w3.org/2000/xmlns/', 'http://schemas.microsoft.com/WMIConfig/2002/State')
                $null = $shellSetupComponent.SetAttribute('xmlns:xsi', 'http://www.w3.org/2000/xmlns/', 'http://www.w3.org/2001/XMLSchema-instance')

                $firstComponentNode = $specializeSettings.SelectSingleNode('un:component', $namespaceManager)
                if ($null -ne $firstComponentNode) {
                    $null = $specializeSettings.InsertBefore($shellSetupComponent, $firstComponentNode)
                }
                else {
                    $null = $specializeSettings.AppendChild($shellSetupComponent)
                }
                $createdShellSetupComponent = $true
            }

            $computerNameElement = $shellSetupComponent.SelectSingleNode('un:ComputerName', $namespaceManager)
            $createdComputerNameElement = $false
            if ($null -eq $computerNameElement) {
                $computerNameElement = $UnattendXml.CreateElement('ComputerName', $unattendNamespace)
                $null = $shellSetupComponent.AppendChild($computerNameElement)
                $createdComputerNameElement = $true
            }

            return [PSCustomObject]@{
                ComputerNameElement        = $computerNameElement
                CreatedSpecializeSettings  = $createdSpecializeSettings
                CreatedShellSetupComponent = $createdShellSetupComponent
                CreatedComputerNameElement = $createdComputerNameElement
            }
        }

        function Save-LocalStagedUnattendFile {
            param(
                [string]$SourcePath,
                [string]$DestinationPath,
                [string]$DeviceNamingMode,
                [string]$DeviceNameTemplate,
                [string]$WindowsArch,
                [bool]$LegacyPrefixesWillBeStaged = $false
            )

            if ($DeviceNamingMode -eq 'None') {
                Copy-Item -Path $SourcePath -Destination $DestinationPath -Force | Out-Null
                return
            }

            [xml]$unattendXml = Get-Content -Path $SourcePath
            $computerNamePath = Initialize-UnattendComputerNamePath -UnattendXml $unattendXml -WindowsArch $WindowsArch

            if ($computerNamePath.CreatedSpecializeSettings -or $computerNamePath.CreatedShellSetupComponent -or $computerNamePath.CreatedComputerNameElement) {
                $createdParts = @()
                if ($computerNamePath.CreatedSpecializeSettings) {
                    $createdParts += 'specialize settings'
                }
                if ($computerNamePath.CreatedShellSetupComponent) {
                    $createdParts += 'Microsoft-Windows-Shell-Setup component'
                }
                if ($computerNamePath.CreatedComputerNameElement) {
                    $createdParts += 'ComputerName element'
                }
                WriteLog "Created $($createdParts -join ', ') while staging unattend file $DestinationPath"
            }

            if ($DeviceNamingMode -eq 'Prompt') {
                $computerNamePath.ComputerNameElement.InnerText = 'MyComputer'
            }
            elseif ($DeviceNamingMode -eq 'Template') {
                $computerNamePath.ComputerNameElement.InnerText = $DeviceNameTemplate
            }
            elseif ($DeviceNamingMode -eq 'Prefixes') {
                if ($computerNamePath.CreatedComputerNameElement) {
                    $computerNamePath.ComputerNameElement.InnerText = '*'
                }
            }
            elseif ($DeviceNamingMode -eq 'SerialComputerNames') {
                if ($computerNamePath.CreatedComputerNameElement) {
                    $computerNamePath.ComputerNameElement.InnerText = '*'
                }
            }
            elseif (($DeviceNamingMode -eq 'Legacy') -and $computerNamePath.CreatedComputerNameElement) {
                $computerNamePath.ComputerNameElement.InnerText = if ($LegacyPrefixesWillBeStaged) { '*' } else { 'MyComputer' }
            }

            $unattendXml.Save($DestinationPath)
        }

        $DiskNumber = $USBDrive.DeviceID.Replace("\\.\PHYSICALDRIVE", "")
        WriteLog "Thread $([System.Threading.Thread]::CurrentThread.ManagedThreadId) processing DiskNumber $DiskNumber ($($USBDrive.Model))"

        # Partitioning
        $Disk = Get-Disk -Number $DiskNumber
        if ($Disk.PartitionStyle -ne "RAW") {
            $Disk | Clear-Disk -RemoveData -RemoveOEM -Confirm:$false
            $Disk = Get-Disk -Number $DiskNumber
        }
        if ($Disk.PartitionStyle -eq "RAW") {
            $Disk | Initialize-Disk -PartitionStyle MBR -Confirm:$false
        }
        else {
            $Disk | Get-Partition | Remove-Partition -Confirm:$false
            $Disk | Set-Disk -PartitionStyle MBR
        }
        
        $BootPartition = $Disk | New-Partition -Size 2GB -IsActive -AssignDriveLetter
        $DeployPartition = $Disk | New-Partition -UseMaximumSize -AssignDriveLetter
        Format-Volume -Partition $BootPartition -FileSystem FAT32 -NewFileSystemLabel "TempBoot" -Confirm:$false | Out-Null
        Format-Volume -Partition $DeployPartition -FileSystem NTFS -NewFileSystemLabel "TempDeploy" -Confirm:$false | Out-Null
        
        $BootPartitionDriveLetter = "$($BootPartition.DriveLetter):\"
        $DeployPartitionDriveLetter = "$($DeployPartition.DriveLetter):\"
        WriteLog "Disk $DiskNumber partitioned. Boot: $BootPartitionDriveLetter, Deploy: $DeployPartitionDriveLetter"

        # Copy WinPE files
        WriteLog "Copying WinPE files from $($using:ISOMountPoint) to $BootPartitionDriveLetter"
        robocopy $using:ISOMountPoint $BootPartitionDriveLetter /E /COPYALL /R:5 /W:5 /J /NFL /NDL /NJH /NJS /nc /ns /np | Out-Null

        # Copy other files
        if ($using:CopyFFU.IsPresent -and $null -ne $using:SelectedFFUFile) {
            $selectedFFUFiles = @($using:SelectedFFUFile)
            if ($selectedFFUFiles.Count -gt 1) {
                WriteLog "Copying multiple FFU files to $DeployPartitionDriveLetter"
            }
            elseif ($selectedFFUFiles.Count -eq 1) {
                WriteLog "Copying $(Split-Path $selectedFFUFiles[0] -Leaf) to $DeployPartitionDriveLetter"
            }
            foreach ($ffuFile in $selectedFFUFiles) {
                robocopy (Split-Path $ffuFile -Parent) $DeployPartitionDriveLetter (Split-Path $ffuFile -Leaf) /J /NFL /NDL /NJH /NJS /nc /ns /np | Out-Null
            }
        }

        if ($using:CopyDrivers) {
            $DriversPathOnUSB = Join-Path $DeployPartitionDriveLetter "Drivers"
            WriteLog "Copying drivers to $DriversPathOnUSB"
            robocopy $using:DriversFolder $DriversPathOnUSB /E /COPYALL /R:5 /W:5 /J /XF .gitkeep /NFL /NDL /NJH /NJS /nc /ns /np | Out-Null
        }

        if ($using:CopyPPKG) {
            $PPKGPathOnUSB = Join-Path $DeployPartitionDriveLetter "PPKG"
            WriteLog "Copying PPKGs to $PPKGPathOnUSB"
            robocopy $using:PPKGFolder $PPKGPathOnUSB /E /COPYALL /R:5 /W:5 /J /NFL /NDL /NJH /NJS /nc /ns /np | Out-Null
        }

        if ($using:CopyUnattend) {
            $UnattendPathOnUSB = Join-Path $DeployPartitionDriveLetter "Unattend"
            WriteLog "Copying unattend file to $UnattendPathOnUSB"
            New-Item -Path $UnattendPathOnUSB -ItemType Directory -ErrorAction SilentlyContinue | Out-Null
            $unattendSource = Get-LocalUnattendSourcePath -UnattendFolder $using:UnattendFolder -WindowsArch $using:WindowsArch -UnattendX64FilePath $using:UnattendX64FilePath -UnattendArm64FilePath $using:UnattendArm64FilePath
            $legacyPrefixesWillBeStaged = ($using:DeviceNamingMode -eq 'Legacy') -and (Test-Path -Path $using:resolvedDeviceNamePrefixesPath -PathType Leaf)
            Save-LocalStagedUnattendFile -SourcePath $unattendSource -DestinationPath (Join-Path $UnattendPathOnUSB 'Unattend.xml') -DeviceNamingMode $using:DeviceNamingMode -DeviceNameTemplate $using:normalizedDeviceNameTemplate -WindowsArch $using:WindowsArch -LegacyPrefixesWillBeStaged $legacyPrefixesWillBeStaged
            if ($using:DeviceNamingMode -eq 'Prefixes') {
                WriteLog "Writing prefixes.txt file to $UnattendPathOnUSB"
                $using:effectiveDeviceNamePrefixes | Set-Content -Path (Join-Path $UnattendPathOnUSB 'prefixes.txt') -Encoding UTF8
            }
            elseif ($using:DeviceNamingMode -eq 'SerialComputerNames') {
                WriteLog "Writing SerialComputerNames.csv file to $UnattendPathOnUSB"
                $using:effectiveDeviceNameSerialComputerNames | Set-Content -Path (Join-Path $UnattendPathOnUSB 'SerialComputerNames.csv') -Encoding UTF8
            }
            elseif ($legacyPrefixesWillBeStaged) {
                WriteLog "Copying prefixes.txt file to $UnattendPathOnUSB"
                Copy-Item -Path $using:resolvedDeviceNamePrefixesPath -Destination (Join-Path $UnattendPathOnUSB 'prefixes.txt') -Force | Out-Null
            }
            WriteLog 'Copy completed'
        }

        if ($using:CopyAutopilot) {
            $AutopilotPathOnUSB = Join-Path $DeployPartitionDriveLetter "Autopilot"
            WriteLog "Copying Autopilot files to $AutopilotPathOnUSB"
            robocopy $using:AutopilotFolder $AutopilotPathOnUSB /E /COPYALL /R:5 /W:5 /J /NFL /NDL /NJH /NJS /nc /ns /np | Out-Null
        }

        # Rename volumes
        WriteLog "Renaming volumes for disk $DiskNumber"
        Set-Volume -DriveLetter $BootPartition.DriveLetter -NewFileSystemLabel "Boot"
        Set-Volume -DriveLetter $DeployPartition.DriveLetter -NewFileSystemLabel "Deploy"
        WriteLog "Finished processing disk $DiskNumber"

    } -ThrottleLimit $resolvedUSBThrottle

    # Dismount ISO after all parallel jobs are complete
    WriteLog "Dismounting deployment ISO."
    Dismount-DiskImage -ImagePath $DeployISO | Out-Null

    WriteLog "USB Drives completed"
}


function Get-FFUEnvironment {
    WriteLog 'Dirty.txt file detected. Last run did not complete succesfully. Will clean environment'
    try {
        Remove-InProgressItems -FFUDevelopmentPath $FFUDevelopmentPath
    }
    catch {
        WriteLog "Remove-InProgressItems failed: $($_.Exception.Message)"
    }
    if ($CleanupCurrentRunDownloads) {
        try {
            Cleanup-CurrentRunDownloads -FFUDevelopmentPath $FFUDevelopmentPath
        }
        catch {
            WriteLog "Cleanup-CurrentRunDownloads failed: $($_.Exception.Message)"
        }
        try {
            Restore-RunBackups -FFUDevelopmentPath $FFUDevelopmentPath
        }
        catch {
            WriteLog "Restore-RunBackups failed: $($_.Exception.Message)"
        }
    }
    # Check for running VMs that start with '_FFU-' and are in the 'Off' state
    $vms = Get-VM

    # Loop through each VM
    foreach ($vm in $vms) {
        if ($vm.Name.StartsWith("_FFU-")) {
            if ($vm.State -eq 'Running') {
                Stop-VM -Name $vm.Name -TurnOff -Force
            }
            # If conditions are met, delete the VM
            Remove-FFUVM -VMName $vm.Name
        }
    }
    # Check for MSFT Virtual disks where location contains FFUDevelopment in the path
    $disks = Get-Disk -FriendlyName *virtual*
    foreach ($disk in $disks) {
        $diskNumber = $disk.Number
        $vhdLocation = $disk.Location
        if ($vhdLocation -like "*FFUDevelopment*") {
            WriteLog "Dismounting Virtual Disk $diskNumber with Location $vhdLocation"
            Dismount-ScratchVhdx -VhdxPath $vhdLocation
            $parentFolder = Split-Path -Parent $vhdLocation
            WriteLog "Removing folder $parentFolder"
            Remove-Item -Path $parentFolder -Recurse -Force
        }
    }

    # Check for mounted DiskImages
    $volumes = Get-Volume | Where-Object { $_.DriveType -eq 'CD-ROM' }
    foreach ($volume in $volumes) {
        $letter = $volume.DriveLetter
        WriteLog "Dismounting DiskImage for volume $letter"
        Get-Volume $letter | Get-DiskImage | Dismount-DiskImage | Out-Null
        WriteLog "Dismounting complete"
    }

    # Remove unused mountpoints
    WriteLog 'Remove unused mountpoints'
    Invoke-Process cmd "/c mountvol /r" | Out-Null
    WriteLog 'Removal complete'

    # Check for content in the VM folder and delete any folders that start with _FFU-
    if ([string]::IsNullOrWhiteSpace($VMLocation)) {
        $VMLocation = Join-Path $FFUDevelopmentPath 'VM'
        WriteLog "VMLocation not set; defaulting to $VMLocation"
    }
    if (Test-Path -Path $VMLocation) {
        $folders = Get-ChildItem -Path $VMLocation -Directory
        foreach ($folder in $folders) {
            if ($folder.Name -like '_FFU-*') {
                WriteLog "Removing folder $($folder.FullName)"
                Remove-Item -Path $folder.FullName -Recurse -Force
            }
        }
    }
    else {
        WriteLog "VMLocation path $VMLocation not found; skipping VM folder cleanup"
    }

    # Remove orphaned mounted images
    $mountedImages = Get-WindowsImage -Mounted
    if ($mountedImages) {
        foreach ($image in $mountedImages) {
            $mountPath = $image.Path
            WriteLog "Dismounting image at $mountPath"
            try {
                Dismount-WindowsImage -Path $mountPath -discard | Out-null
                WriteLog "Successfully dismounted image at $mountPath"
            }
            catch {
                WriteLog "Failed to dismount image at $mountPath with error: $_"
            }
        }
    }

    # Remove Mount folder if it exists
    if (Test-Path -Path "$FFUDevelopmentPath\Mount") {
        WriteLog "Remove $FFUDevelopmentPath\Mount folder"
        Remove-Item -Path "$FFUDevelopmentPath\Mount" -Recurse -Force
        WriteLog 'Folder removed'
    }

    #Clear any corrupt Windows mount points
    WriteLog 'Clearing any corrupt Windows mount points'
    Clear-WindowsCorruptMountPoint | Out-null
    WriteLog 'Complete'

    #Clean up registry
    if (Test-Path -Path 'HKLM:\FFU') {
        Writelog 'Found HKLM:\FFU, removing it' 
        Invoke-Process reg "unload HKLM\FFU" | Out-Null
    }

    #Run shared cleanup to avoid duplicated logic
    Invoke-FFUPostBuildCleanup -RootPath $FFUDevelopmentPath -AppsPath $AppsPath -DriversPath $DriversFolder -FFUCapturePath $FFUCaptureLocation -DeployISOPath $DeployISO -AppsISOPath $AppsISO -RemoveDeployISO:$CleanupDeployISO -RemoveAppsISO:$CleanupAppsISO -RemoveDrivers:$CleanupDrivers -RemoveFFU:$RemoveFFU -RemoveApps:$RemoveApps -RemoveUpdates:$RemoveUpdates -RemoveDownloadedESD:$RemoveDownloadedESD -KBPath:$KBPath

    # Remove existing Apps.iso
    if (Test-Path -Path $AppsISO) {
        WriteLog "Removing $AppsISO"
        Remove-Item -Path $AppsISO -Force -ErrorAction SilentlyContinue
        WriteLog 'Removal complete'
    }
    # Remove per-run session folder if present (Cancel/-Cleanup scenario)
    $sessionDir = Join-Path $FFUDevelopmentPath '.session'
    if (Test-Path -Path $sessionDir) {
        WriteLog 'Removing .session folder'
        Remove-Item -Path $sessionDir -Recurse -Force -ErrorAction SilentlyContinue
        WriteLog 'Removal complete'
    }
    WriteLog 'Removing dirty.txt file'
    Remove-Item -Path "$FFUDevelopmentPath\dirty.txt" -Force
    WriteLog "Cleanup complete"
}
Function Remove-DisabledArtifacts {
    # Remove Office artifacts if Install Office is disabled
    if (-not $InstallOffice) {
        $removed = $false
        if (Test-Path -Path $installOfficePath) {
            WriteLog "Install Office disabled - removing $installOfficePath"
            Remove-Item -Path $installOfficePath -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if (Test-Path -Path $OfficePath) {
            WriteLog 'Removing Office and ODT download'
            $OfficeDownloadPath = "$OfficePath\Office"
            Remove-Item -Path $OfficeDownloadPath -Recurse -Force -ErrorAction SilentlyContinue
            Remove-Item -Path "$OfficePath\setup.exe" -Recurse -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if ($removed) { WriteLog 'Removal complete' }
    }


    # Remove Defender artifacts if Defender update is disabled
    if (-not $UpdateLatestDefender) {
        $removed = $false
        if (Test-Path -Path $installDefenderPath) {
            WriteLog "Update Defender disabled - removing $installDefenderPath"
            Remove-Item -Path $installDefenderPath -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if (Test-Path -Path $DefenderPath) {
            WriteLog "Update Defender disabled - removing $DefenderPath"
            Remove-Item -Path $DefenderPath -Recurse -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if ($removed) { WriteLog 'Removal complete' }
    }

    # Remove MSRT artifacts if MSRT update is disabled
    if (-not $UpdateLatestMSRT) {
        $removed = $false
        if (Test-Path -Path $installMSRTPath) {
            WriteLog "Update MSRT disabled - removing $installMSRTPath"
            Remove-Item -Path $installMSRTPath -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if (Test-Path -Path $MSRTPath) {
            WriteLog "Update MSRT disabled - removing $MSRTPath"
            Remove-Item -Path $MSRTPath -Recurse -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if ($removed) { WriteLog 'Removal complete' }
    }

    # Remove OneDrive artifacts if OneDrive update is disabled
    if (-not $UpdateOneDrive) {
        $removed = $false
        if (Test-Path -Path $installODPath) {
            WriteLog "Update OneDrive disabled - removing $installODPath"
            Remove-Item -Path $installODPath -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if (Test-Path -Path $OneDrivePath) {
            WriteLog "Update OneDrive disabled - removing $OneDrivePath"
            Remove-Item -Path $OneDrivePath -Recurse -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if ($removed) { WriteLog 'Removal complete' }
    }

    # Remove Edge artifacts if Edge update is disabled
    if (-not $UpdateEdge) {
        $removed = $false
        if (Test-Path -Path $installEdgePath) {
            WriteLog "Update Edge disabled - removing $installEdgePath"
            Remove-Item -Path $installEdgePath -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if (Test-Path -Path $EdgePath) {
            WriteLog "Update Edge disabled - removing $EdgePath"
            Remove-Item -Path $EdgePath -Recurse -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if ($removed) { WriteLog 'Removal complete' }
    }
    
    # Remove LTSC CU in-VM artifacts when this scenario is not selected
    if (-not ($UpdateLatestCU -and $installationType -eq 'Client' -and $WindowsRelease -in 2016, 2019, 2021 -and $WindowsSKU -like '*LTS*')) {
        $removed = $false
        if (Test-Path -Path $InstallLTSCUpdatePath) {
            WriteLog "Windows 10 LTSB/LTSC latest CU in-VM install not selected - removing $InstallLTSCUpdatePath"
            Remove-Item -Path $InstallLTSCUpdatePath -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if (Test-Path -Path $LtscCUStagePath) {
            WriteLog "Windows 10 LTSB/LTSC latest CU in-VM install not selected - removing $LtscCUStagePath"
            Remove-Item -Path $LtscCUStagePath -Recurse -Force -ErrorAction SilentlyContinue
            $removed = $true
        }
        if ($removed) { WriteLog 'Removal complete' }
    }
}

function Export-ConfigFile {
    [CmdletBinding()]
    param (
        [Parameter()]
        $paramNames
    )
    $filteredParamNames = Get-Parameters -ParamNames $paramNames
    
    # Retrieve their values
    $paramsToExport = @{}
    foreach ($paramName in $filteredParamNames) {
        $paramsToExport[$paramName] = Get-Variable -Name $paramName -ValueOnly
    }
    
    # Sort the keys alphabetically
    $orderedParams = [ordered]@{}
    foreach ($key in ($paramsToExport.Keys | Sort-Object)) {
        $orderedParams[$key] = $paramsToExport[$key]
    }
    
    # Convert to JSON and save
    $orderedParams | ConvertTo-Json -Depth 10 | Set-Content -Path $ExportConfigFile -Encoding UTF8
}
function Get-PEArchitecture {
    param(
        [string]$FilePath
    )
    
    # Read the entire file as bytes.1
    $bytes = [System.IO.File]::ReadAllBytes($FilePath)
    
    # Check for the 'MZ' signature.
    if ($bytes[0] -ne 0x4D -or $bytes[1] -ne 0x5A) {
        throw "The file is not a valid PE file."
    }
    
    # The PE header offset is stored at offset 0x3C.
    $peHeaderOffset = [System.BitConverter]::ToInt32($bytes, 0x3C)
    
    # Verify the PE signature "PE\0\0".
    if ($bytes[$peHeaderOffset] -ne 0x50 -or $bytes[$peHeaderOffset + 1] -ne 0x45) {
        throw "Invalid PE header."
    }
    
    # The Machine field is located immediately after the PE signature.
    $machine = [System.BitConverter]::ToUInt16($bytes, $peHeaderOffset + 4)
    
    switch ($machine) {
        0x014c { return "x86" }
        0x8664 { return "x64" }
        0xAA64 { return "ARM64" }
        default { return ("Unknown architecture: 0x{0:X}" -f $machine) }
    }
}

function New-RunSession {
    param(
        [string]$FFUDevelopmentPath,
        [string]$DriversFolder,
        [string]$OrchestrationPath
    )
    try {
        $sessionDir = Join-Path $FFUDevelopmentPath '.session'
        $backupDir = Join-Path $sessionDir 'backups'
        $inprogDir = Join-Path $sessionDir 'inprogress'
        if (-not (Test-Path $sessionDir)) { New-Item -ItemType Directory -Path $sessionDir -Force | Out-Null }
        if (-not (Test-Path $backupDir)) { New-Item -ItemType Directory -Path $backupDir -Force | Out-Null }
        if (-not (Test-Path $inprogDir)) { New-Item -ItemType Directory -Path $inprogDir -Force | Out-Null }

        $manifest = [ordered]@{
            SchemaVersion    = 2
            RunStartUtc      = (Get-Date).ToUniversalTime().ToString('o')
            JsonBackups      = @()
            OfficeXmlBackups = @()
            FileBackups      = @()
            DownloadTargets  = @()
        }

        if ($DriversFolder) {
            $driverMapPath = Join-Path $DriversFolder 'DriverMapping.json'
            if (Test-Path $driverMapPath) {
                $backup = Join-Path $backupDir 'DriverMapping.json'
                Copy-Item -Path $driverMapPath -Destination $backup -Force
                $manifest.JsonBackups += @{ Path = $driverMapPath; Backup = $backup }
                WriteLog "Backed up DriverMapping.json to $backup"
            }
        }
        # Backup Office XMLs (DeployFFU.xml, DownloadFFU.xml) if present so we can restore them after cleanup
        if ($OfficePath) {
            foreach ($n in @('DeployFFU.xml', 'DownloadFFU.xml')) {
                $src = Join-Path $OfficePath $n
                if (Test-Path $src) {
                    $dst = Join-Path $backupDir $n
                    try {
                        Copy-Item -Path $src -Destination $dst -Force
                        $manifest.OfficeXmlBackups += @{ Path = $src; Backup = $dst }
                        WriteLog "Backed up $n to $dst"
                    }
                    catch { WriteLog "Failed backing up $($n): $($_.Exception.Message)" }
                }
            }
        }

        $manifestPath = Join-Path $sessionDir 'currentRun.json'
        $manifest | ConvertTo-Json -Depth 5 | Set-Content -Path $manifestPath -Encoding UTF8
        WriteLog "Run session initialized at $sessionDir"
    }
    catch {
        WriteLog "New-RunSession failed: $($_.Exception.Message)"
    }
}
function Get-CurrentRunManifest {
    param([string]$FFUDevelopmentPath)
    $manifestPath = Join-Path $FFUDevelopmentPath '.session\currentRun.json'
    if (Test-Path $manifestPath) { return (Get-Content $manifestPath -Raw | ConvertFrom-Json) }
    return $null
}
function Save-RunManifest {
    param([string]$FFUDevelopmentPath, [object]$Manifest)
    if ($null -eq $Manifest) { return }
    $manifestPath = Join-Path $FFUDevelopmentPath '.session\currentRun.json'
    $Manifest | ConvertTo-Json -Depth 5 | Set-Content -Path $manifestPath -Encoding UTF8
}
function Backup-RunFile {
    param(
        [string]$FFUDevelopmentPath,
        [string]$Path
    )
    if ([string]::IsNullOrWhiteSpace($FFUDevelopmentPath) -or [string]::IsNullOrWhiteSpace($Path)) { return }
    if (-not (Test-Path -LiteralPath $Path -PathType Leaf)) { return }

    $manifest = Get-CurrentRunManifest -FFUDevelopmentPath $FFUDevelopmentPath
    if ($null -eq $manifest) { return }

    # Ensure FileBackups exists for newer schema while remaining compatible with older manifests.
    if ($null -eq $manifest.PSObject.Properties['FileBackups']) {
        Add-Member -InputObject $manifest -MemberType NoteProperty -Name FileBackups -Value @()
    }

    # Skip when this path is already backed up (legacy or current schema).
    $alreadyBackedUp = $false
    foreach ($entry in @($manifest.FileBackups) + @($manifest.JsonBackups) + @($manifest.OfficeXmlBackups)) {
        if ($entry.Path -eq $Path) {
            $alreadyBackedUp = $true
            break
        }
    }
    if ($alreadyBackedUp) { return }

    try {
        # Store a uniquely named backup file so restore can safely overwrite or recreate the original path.
        $backupDir = Join-Path (Join-Path $FFUDevelopmentPath '.session') 'backups'
        if (-not (Test-Path -Path $backupDir)) {
            New-Item -ItemType Directory -Path $backupDir -Force | Out-Null
        }

        $safeLeaf = Split-Path -Path $Path -Leaf
        $backupName = "{0}_{1}" -f ([guid]::NewGuid().ToString('N')), $safeLeaf
        $backupPath = Join-Path $backupDir $backupName

        Copy-Item -LiteralPath $Path -Destination $backupPath -Force

        $fileBackups = @($manifest.FileBackups)
        $fileBackups += @{ Path = $Path; Backup = $backupPath }
        $manifest.FileBackups = $fileBackups
        Save-RunManifest -FFUDevelopmentPath $FFUDevelopmentPath -Manifest $manifest

        WriteLog "Backed up existing file before modification: $Path"
    }
    catch {
        WriteLog "Backup-RunFile failed for $($Path): $($_.Exception.Message)"
    }
}
function Mark-DownloadInProgress {
    param([string]$FFUDevelopmentPath, [string]$TargetPath)
    if ([string]::IsNullOrWhiteSpace($FFUDevelopmentPath) -or [string]::IsNullOrWhiteSpace($TargetPath)) { return }
    $sessionInprog = Join-Path (Join-Path $FFUDevelopmentPath '.session') 'inprogress'
    if (-not (Test-Path $sessionInprog)) { New-Item -ItemType Directory -Path $sessionInprog -Force | Out-Null }
    $marker = Join-Path $sessionInprog ("{0}.marker" -f ([guid]::NewGuid()))
    $payload = @{ TargetPath = $TargetPath; CreatedUtc = (Get-Date).ToUniversalTime().ToString('o') }
    $payload | ConvertTo-Json -Depth 3 | Set-Content -Path $marker -Encoding UTF8

    # Track file download targets so cleanup can remove them even when BITS preserves old source timestamps.
    try {
        if ([System.IO.Path]::HasExtension($TargetPath)) {
            $manifest = Get-CurrentRunManifest -FFUDevelopmentPath $FFUDevelopmentPath
            if ($manifest) {
                if ($null -eq $manifest.PSObject.Properties['DownloadTargets']) {
                    Add-Member -InputObject $manifest -MemberType NoteProperty -Name DownloadTargets -Value @()
                }

                $downloadTargets = @($manifest.DownloadTargets)
                if ($TargetPath -notin $downloadTargets) {
                    $downloadTargets += $TargetPath
                    $manifest.DownloadTargets = $downloadTargets
                    Save-RunManifest -FFUDevelopmentPath $FFUDevelopmentPath -Manifest $manifest
                    WriteLog "Registered current-run download target: $TargetPath"
                }
            }
        }
    }
    catch {
        WriteLog "Failed registering current-run download target $($TargetPath): $($_.Exception.Message)"
    }

    WriteLog "Marked in-progress: $TargetPath"
}
function Clear-DownloadInProgress {
    param([string]$FFUDevelopmentPath, [string]$TargetPath)
    $sessionInprog = Join-Path (Join-Path $FFUDevelopmentPath '.session') 'inprogress'
    if (-not (Test-Path $sessionInprog)) { return }
    Get-ChildItem -Path $sessionInprog -Filter *.marker -ErrorAction SilentlyContinue | ForEach-Object {
        try {
            $data = Get-Content $_.FullName -Raw | ConvertFrom-Json
            if ($data.TargetPath -eq $TargetPath) { Remove-Item -Path $_.FullName -Force }
        }
        catch {}
    }
    WriteLog "Cleared in-progress: $TargetPath"
}
function Remove-InProgressItems {
    param([string]$FFUDevelopmentPath)
    $sessionInprog = Join-Path (Join-Path $FFUDevelopmentPath '.session') 'inprogress'
    if (-not (Test-Path $sessionInprog)) { return }

    function Remove-PathWithRetry {
        param(
            [string]$path,
            [bool]$isDirectory
        )
        for ($i = 0; $i -lt 3; $i++) {
            try {
                if ($isDirectory) {
                    Remove-Item -Path $path -Recurse -Force -ErrorAction Stop
                }
                else {
                    # clear readonly if set
                    try { (Get-Item -LiteralPath $path -ErrorAction SilentlyContinue).Attributes = 'Normal' } catch {}
                    Remove-Item -Path $path -Force -ErrorAction Stop
                }
                return $true
            }
            catch {
                Start-Sleep -Milliseconds 350
            }
        }
        return -not (Test-Path -LiteralPath $path)
    }

    Get-ChildItem -Path $sessionInprog -Filter *.marker -ErrorAction SilentlyContinue | ForEach-Object {
        try {
            $data = Get-Content $_.FullName -Raw | ConvertFrom-Json
            $target = $data.TargetPath
            try {
                if ($DriversFolder -and $target) {
                    $fullTarget = [System.IO.Path]::GetFullPath($target).TrimEnd('\')
                    $driversRoot = [System.IO.Path]::GetFullPath($DriversFolder).TrimEnd('\')
                    if ($fullTarget.StartsWith($driversRoot, [System.StringComparison]::OrdinalIgnoreCase)) {
                        $remainder = $fullTarget.Substring($driversRoot.Length).TrimStart('\')
                        $parts = $remainder -split '\\'
                        if ($parts.Length -ge 1) {
                            $knownMakes = @('Dell', 'HP', 'Lenovo', 'Microsoft')
                            if ($parts.Length -ge 2 -and $knownMakes -contains $parts[0]) {
                                # Drivers\<Make>\<Model>\...
                                $modelFolder = Join-Path (Join-Path $driversRoot $parts[0]) $parts[1]
                            }
                            else {
                                # Drivers\<Model>\... (when DriversFolder already includes Make)
                                $modelFolder = Join-Path $driversRoot $parts[0]
                            }
                            if ($modelFolder) {
                                WriteLog "Promoting in-progress driver target to model folder: $modelFolder (from $target)"
                                $target = $modelFolder
                            }
                        }
                    }
                }
            }
            catch {}

            if (Test-Path $target) {
                # Special-case Office: preserve DeployFFU.xml and DownloadFFU.xml; remove everything else with retries.
                $targetFull = [System.IO.Path]::GetFullPath($target).TrimEnd('\')
                $officeFull = $null
                if ($OfficePath) { $officeFull = [System.IO.Path]::GetFullPath($OfficePath).TrimEnd('\') }

                if ($officeFull -and ($targetFull -ieq $officeFull) -and (Test-Path $OfficePath -PathType Container)) {
                    $preserve = @('DeployFFU.xml', 'DownloadFFU.xml')
                    WriteLog "Cleaning in-progress Office folder: preserving $($preserve -join ', ') and removing other content."
                    Get-ChildItem -Path $OfficePath -Force | ForEach-Object {
                        if ($preserve -notcontains $_.Name) {
                            $itemPath = $_.FullName
                            $isDir = $_.PSIsContainer
                            WriteLog "Removing Office item: $itemPath"
                            $removed = $false
                            try { $removed = Remove-PathWithRetry -path $itemPath -isDirectory:$isDir } catch {}
                            if (-not $removed) {
                                # If setup.exe (or ODT stub) is locked, try to stop the exact owning process by path and retry.
                                try {
                                    $basename = [System.IO.Path]::GetFileName($itemPath)
                                    if (-not $isDir -and $basename -in @('setup.exe', 'odtsetup.exe')) {
                                        Get-Process -ErrorAction SilentlyContinue | Where-Object { $_.Path -eq $itemPath } | Stop-Process -Force -ErrorAction SilentlyContinue
                                        Start-Sleep -Milliseconds 500
                                        $removed = Remove-PathWithRetry -path $itemPath -isDirectory:$false
                                    }
                                }
                                catch {
                                    WriteLog "Process stop attempt for $itemPath failed: $($_.Exception.Message)"
                                }
                            }
                            if (-not $removed) {
                                WriteLog "Failed removing Office item $itemPath after retries."
                            }
                        }
                    }
                }
                else {
                    WriteLog "Removing in-progress target: $target"
                    $isDir = Test-Path $target -PathType Container
                    [void](Remove-PathWithRetry -path $target -isDirectory:$isDir)
                }
            }

            Remove-Item -Path $_.FullName -Force
        }
        catch {
            WriteLog "Failed Remove-InProgressItems marker '$($_.FullName)': $($_.Exception.Message)"
        }
    }
    # Also clean up any driver content created this run (model folders, temp folders, and files)
    # only when broader current-run cleanup is requested.
    try {
        if ($DriversFolder -and (Test-Path $DriversFolder)) {
            $manifest = Get-CurrentRunManifest -FFUDevelopmentPath $FFUDevelopmentPath
            $driverProcessingMarker = Join-Path (Join-Path $FFUDevelopmentPath '.session') 'driverProcessing.active'
            $shouldCleanupDriverCurrentRun = ($CleanupCurrentRunDownloads -or (Test-Path -LiteralPath $driverProcessingMarker))

            if ($manifest -and $manifest.RunStartUtc -and $shouldCleanupDriverCurrentRun) {
                if (-not $CleanupCurrentRunDownloads) {
                    WriteLog 'Driver processing marker detected during cancel; running in-progress driver cleanup for current run.'
                }
                $runStart = [datetime]::Parse($manifest.RunStartUtc)
                $affectedMakeRoots = @{}

                # Remove tracked driver download targets for this run.
                # This handles BITS files whose timestamps can be inherited from source metadata.
                # Preserve HP platform list artifacts because they are model-list cache assets, not per-model download artifacts.
                if ($manifest.PSObject.Properties['DownloadTargets']) {
                    foreach ($downloadTarget in @($manifest.DownloadTargets)) {
                        if ([string]::IsNullOrWhiteSpace($downloadTarget)) { continue }

                        try {
                            $fullTarget = [System.IO.Path]::GetFullPath($downloadTarget).TrimEnd('\')
                            $driversRoot = [System.IO.Path]::GetFullPath($DriversFolder).TrimEnd('\')
                            if (-not $fullTarget.StartsWith($driversRoot, [System.StringComparison]::OrdinalIgnoreCase)) { continue }

                            $leafName = [System.IO.Path]::GetFileName($fullTarget)
                            if ($leafName -in @('platformList.cab', 'platformList.xml')) { continue }

                            if (Test-Path -LiteralPath $downloadTarget -PathType Leaf) {
                                WriteLog "Removing tracked driver file from current run: $downloadTarget"
                                $removed = $false
                                try { $removed = Remove-PathWithRetry -path $downloadTarget -isDirectory:$false } catch {}

                                if ($removed) {
                                    $relative = $fullTarget.Substring($driversRoot.Length).TrimStart('\')
                                    $parts = $relative -split '\\'
                                    if ($parts.Length -ge 1 -and -not [string]::IsNullOrWhiteSpace($parts[0])) {
                                        $makeRoot = Join-Path $DriversFolder $parts[0]
                                        $affectedMakeRoots[$makeRoot] = $true
                                    }
                                }
                                else {
                                    WriteLog "Failed removing tracked driver file from current run after retries: $downloadTarget"
                                }
                            }
                        }
                        catch {
                            WriteLog "Failed removing tracked driver target '$($downloadTarget)': $($_.Exception.Message)"
                        }
                    }
                }

                # Remove OEM temp folders like _TEMP_* (safe to always remove)
                Get-ChildItem -Path $DriversFolder -Directory -Recurse -ErrorAction SilentlyContinue |
                Where-Object { $_.Name -like '_TEMP_*' } |
                ForEach-Object {
                    WriteLog "Removing driver temp folder: $($_.FullName)"
                    Remove-Item -Path $_.FullName -Recurse -Force -ErrorAction SilentlyContinue
                }

                # Remove residual BITS temp files under Drivers regardless of timestamp.
                Get-ChildItem -Path $DriversFolder -Directory -ErrorAction SilentlyContinue | ForEach-Object {
                    $makeRoot = $_.FullName
                    Get-ChildItem -Path $makeRoot -File -Recurse -Force -Filter 'BIT*.tmp' -ErrorAction SilentlyContinue | ForEach-Object {
                        WriteLog "Removing residual BITS temp file: $($_.FullName)"
                        $removed = $false
                        try { $removed = Remove-PathWithRetry -path $_.FullName -isDirectory:$false } catch {}
                        if ($removed) {
                            $affectedMakeRoots[$makeRoot] = $true
                        }
                        else {
                            WriteLog "Failed removing residual BITS temp file after retries: $($_.FullName)"
                        }
                    }
                }

                # Remove model folders created/modified this run; never remove top-level make roots
                Get-ChildItem -Path $DriversFolder -Directory -ErrorAction SilentlyContinue | ForEach-Object {
                    $makeRoot = $_.FullName
                    # Model-level folders are immediate children under a make root (e.g. Drivers\Lenovo\<Model>)
                    Get-ChildItem -Path $makeRoot -Directory -ErrorAction SilentlyContinue |
                    Where-Object { $_.CreationTimeUtc -ge $runStart -or $_.LastWriteTimeUtc -ge $runStart } |
                    ForEach-Object {
                        WriteLog "Removing driver model folder from current run: $($_.FullName)"
                        $affectedMakeRoots[$makeRoot] = $true
                        Remove-Item -Path $_.FullName -Recurse -Force -ErrorAction SilentlyContinue
                    }
                }

                # Remove run-created driver files that can live outside model folders
                # (for example BITS temp files and vendor index JSON files under make roots).
                Get-ChildItem -Path $DriversFolder -Directory -ErrorAction SilentlyContinue | ForEach-Object {
                    $makeRoot = $_.FullName
                    Get-ChildItem -Path $makeRoot -File -Recurse -Force -ErrorAction SilentlyContinue |
                    Where-Object { $_.CreationTimeUtc -ge $runStart } |
                    ForEach-Object {
                        WriteLog "Removing driver file from current run: $($_.FullName)"
                        $affectedMakeRoots[$makeRoot] = $true
                        $removed = $false
                        try { $removed = Remove-PathWithRetry -path $_.FullName -isDirectory:$false } catch {}
                        if (-not $removed) {
                            WriteLog "Failed removing driver file from current run after retries: $($_.FullName)"
                        }
                    }
                }

                # Remove empty make root folders that were affected during this run.
                Get-ChildItem -Path $DriversFolder -Directory -ErrorAction SilentlyContinue | ForEach-Object {
                    $makeRoot = $_.FullName
                    if ($affectedMakeRoots.ContainsKey($makeRoot)) {
                        $any = Get-ChildItem -Path $makeRoot -Force -Recurse -ErrorAction SilentlyContinue | Select-Object -First 1
                        if ($null -eq $any) {
                            WriteLog "Removing empty make root folder after current-run cleanup: $makeRoot"
                            Remove-Item -Path $makeRoot -Recurse -Force -ErrorAction SilentlyContinue
                        }
                        else {
                            WriteLog "Skipping non-empty affected make root folder: $makeRoot"
                        }
                    }
                }
            }
        }
    }
    catch {
        WriteLog "Driver in-progress cleanup step failed: $($_.Exception.Message)"
    }
}
function Cleanup-CurrentRunDownloads {
    param([string]$FFUDevelopmentPath)
    $manifest = Get-CurrentRunManifest -FFUDevelopmentPath $FFUDevelopmentPath
    if ($null -eq $manifest) { WriteLog "No current run manifest; skipping current-run cleanup."; return }
    $runStart = [datetime]::Parse($manifest.RunStartUtc)

    # Remove tracked file download targets for this run.
    # This handles BITS files that may keep old source timestamps and bypass time-based cleanup.
    $affectedDriverMakeRoots = @{}
    if ($manifest.PSObject.Properties['DownloadTargets']) {
        foreach ($downloadTarget in @($manifest.DownloadTargets)) {
            if ([string]::IsNullOrWhiteSpace($downloadTarget)) { continue }

            try {
                if (Test-Path -LiteralPath $downloadTarget -PathType Leaf) {
                    WriteLog "Removing tracked current-run download file: $downloadTarget"

                    $removed = $false
                    for ($i = 0; $i -lt 5; $i++) {
                        try {
                            try { (Get-Item -LiteralPath $downloadTarget -ErrorAction SilentlyContinue).Attributes = 'Normal' } catch {}
                            Remove-Item -LiteralPath $downloadTarget -Force -ErrorAction Stop
                            $removed = $true
                            break
                        }
                        catch {
                            Start-Sleep -Milliseconds 400
                        }
                    }

                    if ($removed) {
                        if ($DriversFolder) {
                            try {
                                $fullTarget = [System.IO.Path]::GetFullPath($downloadTarget).TrimEnd('\')
                                $driversRoot = [System.IO.Path]::GetFullPath($DriversFolder).TrimEnd('\')
                                if ($fullTarget.StartsWith($driversRoot, [System.StringComparison]::OrdinalIgnoreCase)) {
                                    $parent = Split-Path -Path $fullTarget -Parent
                                    $parentParent = Split-Path -Path $parent -Parent
                                    if ($parent -and $parentParent) {
                                        $fullParentParent = [System.IO.Path]::GetFullPath($parentParent).TrimEnd('\')
                                        if ($fullParentParent -ieq $driversRoot) {
                                            $affectedDriverMakeRoots[$parent] = $true
                                        }
                                    }
                                }
                            }
                            catch {}
                        }
                    }
                    else {
                        WriteLog "Failed removing tracked current-run download file after retries: $downloadTarget"
                    }
                }
            }
            catch {
                WriteLog "Failed removing tracked current-run download file $($downloadTarget): $($_.Exception.Message)"
            }
        }
    }

    # Remove empty make roots affected by tracked download file removal.
    foreach ($makeRoot in $affectedDriverMakeRoots.Keys) {
        try {
            if (Test-Path -LiteralPath $makeRoot -PathType Container) {
                $any = Get-ChildItem -Path $makeRoot -Force -Recurse -ErrorAction SilentlyContinue | Select-Object -First 1
                if ($null -eq $any) {
                    WriteLog "Removing empty make root folder after tracked download cleanup: $makeRoot"
                    Remove-Item -LiteralPath $makeRoot -Recurse -Force -ErrorAction SilentlyContinue
                }
            }
        }
        catch {
            WriteLog "Failed removing empty make root folder '$($makeRoot)': $($_.Exception.Message)"
        }
    }

    # 1) Generic current-run scrub across known download roots.
    $roots = @()
    if ($AppsPath) { $roots += (Join-Path $AppsPath 'Win32'); $roots += (Join-Path $AppsPath 'MSStore') }
    if ($DefenderPath) { $roots += $DefenderPath }
    if ($MSRTPath) { $roots += $MSRTPath }
    if ($OneDrivePath) { $roots += $OneDrivePath }
    if ($EdgePath) { $roots += $EdgePath }
    if ($KBPath) { $roots += $KBPath }
    if ($DriversFolder) { $roots += $DriversFolder }

    foreach ($root in $roots | Where-Object { $_ -and (Test-Path $_) }) {
        $isDriversRoot = $false
        try {
            if ($DriversFolder) {
                $isDriversRoot = ([System.IO.Path]::GetFullPath($root).TrimEnd('\') -ieq [System.IO.Path]::GetFullPath($DriversFolder).TrimEnd('\'))
            }
        }
        catch {}

        if ($isDriversRoot) {
            WriteLog "Scanning Drivers folder (creation-time filter) in $root"

            # Remove driver folders created this run (skip non-empty make roots)
            Get-ChildItem -Path $root -Directory -Recurse -ErrorAction SilentlyContinue |
            Where-Object { $_.CreationTimeUtc -ge $runStart } |
            Sort-Object FullName -Descending | ForEach-Object {
                try {
                    $parent = Split-Path -Path $_.FullName -Parent
                    $parentIsDriversRoot = ([System.IO.Path]::GetFullPath($parent).TrimEnd('\') -ieq [System.IO.Path]::GetFullPath($root).TrimEnd('\'))
                    if ($parentIsDriversRoot) {
                        # Only remove top-level make folders if created this run AND empty (avoid deleting existing Lenovo/HP/Dell/Microsoft trees)
                        $any = Get-ChildItem -Path $_.FullName -Force -Recurse -ErrorAction SilentlyContinue | Select-Object -First 1
                        if ($null -eq $any) {
                            WriteLog "Removing empty make folder created this run: $($_.FullName)"
                            Remove-Item -Path $_.FullName -Recurse -Force -ErrorAction SilentlyContinue
                        }
                    }
                    else {
                        WriteLog "Removing current-run driver folder: $($_.FullName)"
                        Remove-Item -Path $_.FullName -Recurse -Force -ErrorAction SilentlyContinue
                    }
                }
                catch { WriteLog "Failed removing driver folder $($_.FullName): $($_.Exception.Message)" }
            }

            # Remove driver files created this run
            Get-ChildItem -Path $root -File -Recurse -Force -ErrorAction SilentlyContinue |
            Where-Object { $_.CreationTimeUtc -ge $runStart } |
            ForEach-Object {
                try {
                    WriteLog "Removing current-run driver file: $($_.FullName)"
                    Remove-Item -Path $_.FullName -Force -ErrorAction SilentlyContinue
                }
                catch { WriteLog "Failed removing driver file $($_.FullName): $($_.Exception.Message)" }
            }

            # Remove residual BITS temp files under Drivers regardless of timestamp.
            Get-ChildItem -Path $root -File -Recurse -Force -Filter 'BIT*.tmp' -ErrorAction SilentlyContinue |
            ForEach-Object {
                try {
                    WriteLog "Removing residual BITS temp file: $($_.FullName)"
                    Remove-Item -Path $_.FullName -Force -ErrorAction SilentlyContinue
                }
                catch { WriteLog "Failed removing residual BITS temp file $($_.FullName): $($_.Exception.Message)" }
            }

            # Prune empty driver folders (skip existing make roots)
            Get-ChildItem -Path $root -Directory -Recurse -ErrorAction SilentlyContinue |
            Sort-Object FullName -Descending | ForEach-Object {
                try {
                    $any = Get-ChildItem -Path $_.FullName -Force -Recurse -ErrorAction SilentlyContinue | Select-Object -First 1
                    if ($null -eq $any) {
                        $parent = Split-Path -Path $_.FullName -Parent
                        $parentIsDriversRoot = ([System.IO.Path]::GetFullPath($parent).TrimEnd('\') -ieq [System.IO.Path]::GetFullPath($root).TrimEnd('\'))
                        if ($parentIsDriversRoot) {
                            # Only remove empty make roots if they were created this run
                            if ($_.CreationTimeUtc -ge $runStart) {
                                WriteLog "Removing empty make folder created this run: $($_.FullName)"
                                Remove-Item -Path $_.FullName -Recurse -Force -ErrorAction SilentlyContinue
                            }
                        }
                        else {
                            WriteLog "Removing empty driver subfolder: $($_.FullName)"
                            Remove-Item -Path $_.FullName -Recurse -Force -ErrorAction SilentlyContinue
                        }
                    }
                }
                catch { WriteLog "Failed pruning empty driver folder $($_.FullName): $($_.Exception.Message)" }
            }
        }
        else {
            WriteLog "Scanning for current-run items in $root"
            # Remove only folders created this run so pre-existing modified files are not deleted.
            Get-ChildItem -Path $root -Directory -Recurse -ErrorAction SilentlyContinue |
            Where-Object { $_.CreationTimeUtc -ge $runStart } |
            Sort-Object FullName -Descending | ForEach-Object {
                try {
                    WriteLog "Removing current-run folder: $($_.FullName)"
                    Remove-Item -Path $_.FullName -Recurse -Force -ErrorAction SilentlyContinue
                }
                catch { WriteLog "Failed removing folder $($_.FullName): $($_.Exception.Message)" }
            }
            # Remove only files created this run so pre-existing modified files are restored instead of removed.
            Get-ChildItem -Path $root -File -Recurse -ErrorAction SilentlyContinue |
            Where-Object { $_.CreationTimeUtc -ge $runStart -and $_.Name -notin @('DeployFFU.xml', 'DownloadFFU.xml') } |
            ForEach-Object {
                try {
                    WriteLog "Removing current-run file: $($_.FullName)"
                    Remove-Item -Path $_.FullName -Force -ErrorAction SilentlyContinue
                }
                catch { WriteLog "Failed removing file $($_.FullName): $($_.Exception.Message)" }
            }
        }
    }

    # 2) Office folder policy: keep XML configs, remove only current-run Office content
    if ($OfficePath -and (Test-Path $OfficePath)) {
        $preserve = @('DeployFFU.xml', 'DownloadFFU.xml')
        $officeItemsToRemove = Get-ChildItem -Path $OfficePath -Force -ErrorAction SilentlyContinue | Where-Object {
            ($preserve -notcontains $_.Name) -and ($_.CreationTimeUtc -ge $runStart)
        }

        if ($officeItemsToRemove) {
            WriteLog "Cleaning Office folder: preserving $($preserve -join ', ') and removing current-run content only."
            $officeItemsToRemove | ForEach-Object {
                try {
                    WriteLog "Removing Office item: $($_.FullName)"
                    if ($_.PSIsContainer) {
                        Remove-Item -Path $_.FullName -Recurse -Force -ErrorAction SilentlyContinue
                    }
                    else {
                        Remove-Item -Path $_.FullName -Force -ErrorAction SilentlyContinue
                    }
                }
                catch { WriteLog "Failed removing Office item $($_.FullName): $($_.Exception.Message)" }
            }
        }
        else {
            WriteLog 'Skipping Office cleanup: no current-run Office content detected.'
        }
    }

    # 3) Remove generated orchestration artifacts created this run when there was no pre-run backup.
    if ($orchestrationPath -and (Test-Path $orchestrationPath)) {
        $generatedArtifacts = @(
            'Update-Defender.ps1',
            'Update-MSRT.ps1',
            'Update-OneDrive.ps1',
            'Update-Edge.ps1',
            'Install-Office.ps1',
            'Install-LTSCUpdate.ps1',
            'WinGetWin32Apps.json',
            'AppsScriptVariables.json'
        )

        foreach ($artifactName in $generatedArtifacts) {
            try {
                $artifactPath = Join-Path $orchestrationPath $artifactName
                if (Test-Path -LiteralPath $artifactPath) {
                    $hasBackup = (@($manifest.FileBackups) + @($manifest.JsonBackups) + @($manifest.OfficeXmlBackups)) | Where-Object { $_.Path -eq $artifactPath } | Select-Object -First 1
                    if ($null -eq $hasBackup) {
                        $fi = Get-Item -LiteralPath $artifactPath
                        if ($fi.LastWriteTimeUtc -ge $runStart) {
                            WriteLog "Removing current-run artifact: $artifactPath"
                            Remove-Item -LiteralPath $artifactPath -Force -ErrorAction SilentlyContinue
                        }
                    }
                }
            }
            catch { WriteLog "Failed removing generated artifact $($artifactName): $($_.Exception.Message)" }
        }
    }

    # 4) If Defender/OneDrive/Edge/MSRT folders exist, remove them entirely (they're session downloads)
    foreach ($p in @($DefenderPath, $OneDrivePath, $EdgePath, $MSRTPath)) {
        if ($p -and (Test-Path $p)) {
            try {
                WriteLog "Removing current-run folder (entire): $p"
                Remove-Item -Path $p -Recurse -Force -ErrorAction SilentlyContinue
            }
            catch { WriteLog "Failed removing folder $($p): $($_.Exception.Message)" }
        }
    }

    # 5) Remove any ESDs downloaded this run
    Get-ChildItem -Path $PSScriptRoot -Filter *.esd -File -ErrorAction SilentlyContinue | Where-Object { $_.LastWriteTimeUtc -ge $runStart } | ForEach-Object {
        try {
            WriteLog "Removing current-run ESD: $($_.FullName)"
            Remove-Item -Path $_.FullName -Force -ErrorAction SilentlyContinue
        }
        catch { WriteLog "Failed removing ESD $($_.FullName): $($_.Exception.Message)" }
    }

    # 6) Remove empty top-level subfolders under Apps (cosmetic)
    if ($AppsPath -and (Test-Path $AppsPath)) {
        Get-ChildItem -Path $AppsPath -Directory -ErrorAction SilentlyContinue | ForEach-Object {
            try {
                $any = Get-ChildItem -Path $_.FullName -Force -Recurse -ErrorAction SilentlyContinue | Select-Object -First 1
                if ($null -eq $any) {
                    WriteLog "Removing empty folder: $($_.FullName)"
                    Remove-Item -Path $_.FullName -Recurse -Force -ErrorAction SilentlyContinue
                }
            }
            catch { WriteLog "Failed removing empty folder $($_.FullName): $($_.Exception.Message)" }
        }
    }
}
function Restore-RunBackups {
    param([string]$FFUDevelopmentPath)
    $manifest = Get-CurrentRunManifest -FFUDevelopmentPath $FFUDevelopmentPath
    if ($null -eq $manifest) { return }
    $runStart = [datetime]::Parse($manifest.RunStartUtc)

    # Build a unified backup list for compatibility across manifest versions.
    $allBackups = @()
    if ($manifest.PSObject.Properties['FileBackups']) { $allBackups += @($manifest.FileBackups) }
    if ($manifest.PSObject.Properties['JsonBackups']) { $allBackups += @($manifest.JsonBackups) }
    if ($manifest.PSObject.Properties['OfficeXmlBackups']) { $allBackups += @($manifest.OfficeXmlBackups) }

    # Restore any pre-existing files that were backed up before in-run modification.
    foreach ($entry in $allBackups) {
        if ($null -eq $entry -or [string]::IsNullOrWhiteSpace($entry.Path) -or [string]::IsNullOrWhiteSpace($entry.Backup)) { continue }
        $path = $entry.Path
        $backup = $entry.Backup
        try {
            if (Test-Path -LiteralPath $backup) {
                WriteLog "Restoring file from backup: $path"
                Copy-Item -LiteralPath $backup -Destination $path -Force
            }
        }
        catch { WriteLog "Failed restoring backup for $($path): $($_.Exception.Message)" }
    }

    # Remove current-run JSON files that were generated this run and had no pre-run backup.
    $candidateJsons = @()
    if ($DriversFolder) { $candidateJsons += (Join-Path $DriversFolder 'DriverMapping.json') }
    if ($orchestrationPath) { $candidateJsons += (Join-Path $orchestrationPath 'WinGetWin32Apps.json') }

    foreach ($jp in $candidateJsons) {
        if (Test-Path -LiteralPath $jp) {
            $hasBackup = $allBackups | Where-Object { $_.Path -eq $jp } | Select-Object -First 1
            if ($null -eq $hasBackup) {
                $fi = Get-Item -LiteralPath $jp
                if ($fi.LastWriteTimeUtc -ge $runStart) {
                    WriteLog "Removing current-run JSON: $jp"
                    Remove-Item -LiteralPath $jp -Force -ErrorAction SilentlyContinue
                }
            }
        }
    }

    # Restore Office XML backups if present; ensure Office folder exists and only XMLs remain.
    if ($manifest.OfficeXmlBackups -and $OfficePath) {
        if (-not (Test-Path $OfficePath)) {
            try { New-Item -ItemType Directory -Path $OfficePath -Force | Out-Null } catch {}
        }
        foreach ($ox in $manifest.OfficeXmlBackups) {
            try {
                WriteLog "Restoring Office XML from backup: $($ox.Path)"
                Copy-Item -Path $ox.Backup -Destination $ox.Path -Force
            }
            catch { WriteLog "Failed restoring Office XML $($ox.Path): $($_.Exception.Message)" }
        }
        # Ensure only DeployFFU.xml and DownloadFFU.xml remain.
        $preserve = @('DeployFFU.xml', 'DownloadFFU.xml')
        Get-ChildItem -Path $OfficePath -Force -ErrorAction SilentlyContinue | ForEach-Object {
            if ($preserve -notcontains $_.Name) {
                try {
                    if ($_.PSIsContainer) { Remove-Item -Path $_.FullName -Recurse -Force -ErrorAction SilentlyContinue }
                    else { Remove-Item -Path $_.FullName -Force -ErrorAction SilentlyContinue }
                }
                catch { WriteLog "Failed removing extra Office item $($_.FullName): $($_.Exception.Message)" }
            }
        }
    }
}
function Restore-RunJsonBackups {
    param([string]$FFUDevelopmentPath)
    Restore-RunBackups -FFUDevelopmentPath $FFUDevelopmentPath
}

###END FUNCTIONS


if (-not $Cleanup) {
    #Remove old log file if found
    if (Test-Path -Path $Logfile) {
        Remove-item -Path $LogFile -Force
    }

    $startTime = Get-Date
    Write-Host "FFU build process started at" $startTime
    Write-Host "This process can take 20 minutes or more. Please do not close this window or any additional windows that pop up"
    Write-Host "To track progress, please open the log file $Logfile or use the -Verbose parameter next time"
}


if ($Cleanup) {
    WriteLog 'User cancelled, starting cleanup process'
    WriteLog 'Cleanup requested via -Cleanup. Running Get-FFUEnvironment...'
    Get-FFUEnvironment
    return
}

WriteLog 'Begin Logging'
New-RunSession -FFUDevelopmentPath $FFUDevelopmentPath -DriversFolder $DriversFolder -OrchestrationPath $orchestrationPath
Set-Progress -Percentage 1 -Message "FFU build process started..."

####### Generate Config File #######

if ($ExportConfigFile) {
    WriteLog 'Exporting Config File'
    # Get the parameter names from the script and exclude ExportConfigFile
    $paramNames = $MyInvocation.MyCommand.Parameters.Keys | Where-Object { $_ -ne 'ExportConfigFile' }
    try {
        Export-ConfigFile($paramNames)
        WriteLog "Config file exported to $ExportConfigFile"
    }
    catch {
        WriteLog 'Failed to export config file'
        throw $_
    }
}

####### End Generate Config File #######

#Setting long path support - this prevents issues where some applications have deep directory structures
#and oscdimg fails to create the Apps ISO
try {
    $LongPathsEnabled = Get-ItemPropertyValue -Path 'HKLM:\SYSTEM\CurrentControlSet\Control\FileSystem' -Name 'LongPathsEnabled' -ErrorAction Stop
}
catch {
    $LongPathsEnabled = $null
}
if ($LongPathsEnabled -ne 1) {
    WriteLog 'LongPathsEnabled is not set to 1. Setting it to 1'
    Set-ItemProperty -Path 'HKLM:\SYSTEM\CurrentControlSet\Control\FileSystem' -Name 'LongPathsEnabled' -Value 1
    WriteLog 'LongPathsEnabled set to 1'
}


Set-Progress -Percentage 2 -Message "Validating parameters..."
###PARAMETER VALIDATION

#Set build partition drive letters and validate they are available for use; this is required before any build steps that require drive access to ensure the expected drive letters are reserved and to fail fast if there are conflicts.
try {
    $normalizedAdditionalDataPartitions = @(ConvertTo-NormalizedDataPartitions -DataPartitions $AdditionalDataPartitions)
    $configureDeployedDataPartitionDriveLetters = $normalizedAdditionalDataPartitions.Count -gt 0
    $enforceBuildVmDataPartitionDriveLetters = $InstallApps -and $configureDeployedDataPartitionDriveLetters
    $fillRemainingPartitionCount = if ($OSPartitionSize -le 0) { 1 } else { 0 }
    $fillRemainingPartitionCount += @($normalizedAdditionalDataPartitions | Where-Object { $_.FillRemaining }).Count
    if ($fillRemainingPartitionCount -gt 1) {
		throw 'Only one Windows or data partition can fill remaining disk space. Set a fixed Windows partition size before using FillRemaining on a data partition.'
    }
    $partitionLayoutSignature = Get-PartitionLayoutSignature -OSPartitionSize $OSPartitionSize -RecoveryPartitionSize $RecoveryPartitionSize -CreateRecoveryPartition $CreateRecoveryPartition -DataPartitions $normalizedAdditionalDataPartitions

    $partitionDriveLetters = Get-NormalizedPartitionDriveLetters -SystemPartitionDriveLetter $SystemPartitionDriveLetter -WindowsPartitionDriveLetter $WindowsPartitionDriveLetter -RecoveryPartitionDriveLetter $RecoveryPartitionDriveLetter -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions -ValidateAvailable
    $SystemPartitionDriveLetter = $partitionDriveLetters.SystemPartitionDriveLetter
    $WindowsPartitionDriveLetter = $partitionDriveLetters.WindowsPartitionDriveLetter
    if ($CreateRecoveryPartition) {
        $RecoveryPartitionDriveLetter = $partitionDriveLetters.RecoveryPartitionDriveLetter
    }
    $dataPartitionDriveLetterLog = if ($normalizedAdditionalDataPartitions.Count -gt 0) { ', Data=' + (($normalizedAdditionalDataPartitions | ForEach-Object { "$($_.Name):$($_.DriveLetter)" }) -join ', ') } else { '' }
    $recoveryPartitionDriveLetterLog = if ($CreateRecoveryPartition) { ", Recovery=$RecoveryPartitionDriveLetter" } else { ', Recovery=disabled' }
    WriteLog "Using build partition drive letters: System=$SystemPartitionDriveLetter, Windows=$WindowsPartitionDriveLetter$recoveryPartitionDriveLetterLog$dataPartitionDriveLetterLog"
}
catch {
    $partitionDriveLetterValidationError = "Build validation failed: $($_.Exception.Message)"
    Set-Progress -Percentage 2 -Message $partitionDriveLetterValidationError
    WriteLog $partitionDriveLetterValidationError
    throw $partitionDriveLetterValidationError
}

#Validate CopyDrivers dependency on BuildUSBDrive
if ($CopyDrivers -and (-not $BuildUSBDrive)) {
    WriteLog "-CopyDrivers is set to `$true, but -BuildUSBDrive is not set to `$true"
    throw "-CopyDrivers is set to `$true, but -BuildUSBDrive is not set to `$true. Please set -BuildUSBDrive to `$true and try again."
}

#Validate drivers folder
if ($InstallDrivers -or $CopyDrivers) {
    WriteLog 'Doing driver validation'
    if ($DriversFolder -match '\s') {
        WriteLog "Driver folder path $DriversFolder contains spaces. Please remove spaces from the path and try again."
        throw "Driver folder path $DriversFolder contains spaces. Please remove spaces from the path and try again."
    }
    if ($Make -and $Model) {
        WriteLog "Make and Model are set to $Make and $Model, will attempt to download drivers"
    }
    elseif ($DriversJsonPath -and (Test-Path -Path $DriversJsonPath)) {
        WriteLog "Drivers JSON path is set to $DriversJsonPath, will attempt to download drivers"
    }
    else {
        if (!(Test-Path -Path $DriversFolder)) {
            WriteLog "-InstallDrivers or -CopyDrivers is set to `$true, but the $DriversFolder folder is missing"
            throw "-InstallDrivers or -CopyDrivers is set to `$true, but the $DriversFolder folder is missing"
        }
        if ((Get-ChildItem -Path $DriversFolder -Recurse | Measure-Object -Property Length -Sum).Sum -lt 1MB) {
            WriteLog "-InstallDrivers or -CopyDrivers is set to `$true, but the $DriversFolder folder is empty, and no drivers are specified for download."
            throw "-InstallDrivers or -CopyDrivers is set to `$true, but the $DriversFolder folder is empty, and no drivers are specified for download."
        }
        WriteLog "Drivers folder found with content. Will use existing drivers."
    }   
}
#Validate PEDrivers folder
if ($CopyPEDrivers) {
    WriteLog 'Doing PEDriver validation'
    # Check if $PEdriversFolder has spaces in the path, if it does, throw an error
    if ($PEDriversFolder -match '\s') {
        WriteLog "Driver folder path $PEDriversFolder contains spaces. Please remove spaces from the path and try again."
        throw "Driver folder path $PEDriversFolder contains spaces. Please remove spaces from the path and try again."
    }
    if ($UseDriversAsPEDrivers) {
        # When using Drivers as PE drivers, skip strict PEDrivers folder existence/content checks.
        $driverSourceAvailable = $false
        if ($DriversJsonPath -and (Test-Path -Path $DriversJsonPath)) {
            $driverSourceAvailable = $true
            WriteLog "Drivers JSON path is set to $DriversJsonPath; drivers will be downloaded for WinPE."
        }
        elseif ($Make -and $Model) {
            $driverSourceAvailable = $true
            WriteLog "Make/Model ($Make / $Model) specified; drivers will be downloaded for WinPE."
        }
        elseif ((Test-Path -Path $DriversFolder) -and ((Get-ChildItem -Path $DriversFolder -Recurse | Measure-Object -Property Length -Sum).Sum -ge 1MB)) {
            $driverSourceAvailable = $true
            WriteLog "Drivers folder contains existing content; will reuse for WinPE."
        }
        if (-not $driverSourceAvailable) {
            WriteLog "-UseDriversAsPEDrivers is set, but no driver sources are available (Drivers folder missing/empty and no download instructions)."
            throw "-UseDriversAsPEDrivers is set, but no driver sources are available (Drivers folder missing/empty and no download instructions)."
        }
        WriteLog "UseDriversAsPEDrivers is set. Skipping PEDrivers folder existence/content checks; drivers will be sourced from Drivers folder (or downloaded)."
        WriteLog 'PEDriver validation complete'
    }
    else {
        if (!(Test-Path -Path $PEDriversFolder)) {
            WriteLog "-CopyPEDrivers is set to `$true, but the $PEDriversFolder folder is missing"
            throw "-CopyPEDrivers is set to `$true, but the $PEDriversFolder folder is missing"
        }
        if ((Get-ChildItem -Path $PEDriversFolder -Recurse | Measure-Object -Property Length -Sum).Sum -lt 1MB) {
            WriteLog "-CopyPEDrivers is set to `$true, but the $PEDriversFolder folder is empty"
            throw "-CopyPEDrivers is set to `$true, but the $PEDriversFolder folder is empty"
        }
        WriteLog 'PEDriver validation complete'
    }
}

#Validate PPKG folder
if ($CopyPPKG) {
    WriteLog 'Doing PPKG validation'
    if (!(Test-Path -Path $PPKGFolder)) {
        WriteLog "-CopyPPKG is set to `$true, but the $PPKGFolder folder is missing"
        throw "-CopyPPKG is set to `$true, but the $PPKGFolder folder is missing"
    }
    #Check for at least one .PPKG file
    if (!(Get-ChildItem -Path $PPKGFolder -Filter *.ppkg)) {
        WriteLog "-CopyPPKG is set to `$true, but the $PPKGFolder folder is missing a .PPKG file"
        throw "-CopyPPKG is set to `$true, but the $PPKGFolder folder is missing a .PPKG file"
    }
    WriteLog 'PPKG validation complete'
}

#Validate Autopilot folder
if ($CopyAutopilot) {
    WriteLog 'Doing Autopilot validation'
    if (!(Test-Path -Path $AutopilotFolder)) {
        WriteLog "-CopyAutopilot is set to `$true, but the $AutopilotFolder folder is missing"
        throw "-CopyAutopilot is set to `$true, but the $AutopilotFolder folder is missing"
    }
    #Check for .JSON file
    if (!(Get-ChildItem -Path $AutopilotFolder -Filter *.json)) {
        WriteLog "-CopyAutopilot is set to `$true, but the $AutopilotFolder folder is missing a .JSON file"
        throw "-CopyAutopilot is set to `$true, but the $AutopilotFolder folder is missing a .JSON file"
    }
    WriteLog 'Autopilot validation complete'
}

# Validate unattend source file
if ($CopyUnattend -or $InjectUnattend) {
    WriteLog 'Doing Unattend validation'
    $selectedUnattendMode = if ($CopyUnattend) { 'CopyUnattend' } else { 'InjectUnattend' }
    $unattendSourcePath = Get-UnattendSourcePath -UnattendFolder $UnattendFolder -WindowsArch $WindowsArch -UnattendX64FilePath $UnattendX64FilePath -UnattendArm64FilePath $UnattendArm64FilePath
    if (!(Test-Path -Path $unattendSourcePath -PathType Leaf)) {
        WriteLog "-$selectedUnattendMode is set to `$true, but the selected unattend XML file is missing: $unattendSourcePath"
        throw "-$selectedUnattendMode is set to `$true, but the selected unattend XML file is missing: $unattendSourcePath"
    }

    $selectedUnattendFile = Get-Item -Path $unattendSourcePath -ErrorAction SilentlyContinue
    if (($null -eq $selectedUnattendFile) -or ($selectedUnattendFile.Length -le 0)) {
        WriteLog "-$selectedUnattendMode is set to `$true, but the selected unattend XML file is empty: $unattendSourcePath"
        throw "-$selectedUnattendMode is set to `$true, but the selected unattend XML file is empty: $unattendSourcePath"
    }

    if ($DeviceNamingMode -ne 'None') {
        try {
            [xml]$validationUnattendXml = Get-Content -Path $unattendSourcePath
            $null = Initialize-UnattendComputerNamePath -UnattendXml $validationUnattendXml -WindowsArch $WindowsArch
        }
        catch {
            throw "DeviceNamingMode $DeviceNamingMode requires a valid specialize/Microsoft-Windows-Shell-Setup/ComputerName path in $unattendSourcePath. $($_.Exception.Message)"
        }
    }

    WriteLog 'Unattend validation complete'
}

#Override $InstallApps value if using ESD to build FFU. This is due to a strange issue where building the FFU
#from vhdx doesn't work (you get an older style OOBE screen and get stuck in an OOBE reboot loop when hitting next).
#This behavior doesn't happen with WIM files.
# If (-not ($ISOPath) -and (-not ($InstallApps))) {
#     $InstallApps = $true
#     WriteLog "Script will download Windows media. Setting `$InstallApps to `$true to build VM to capture FFU. Must do this when using MCT ESD."
# }

if (($InstallOffice -eq $true) -and ($InstallApps -eq $false)) {
    throw "If variable InstallOffice is set to `$true, InstallApps must also be set to `$true."
}
if ($EnableVMNetworking) {
    if ($InstallApps -eq $false) {
        WriteLog 'EnableVMNetworking is set to true, but InstallApps is false. No VM will be created, so VM networking will be ignored.'
    }
    else {
        if ([string]::IsNullOrWhiteSpace($VMSwitchName)) {
            throw '-EnableVMNetworking requires -VMSwitchName. Select or enter a Hyper-V switch and try again.'
        }

        WriteLog "Experimental VM networking enabled. Validating -VMSwitchName $VMSwitchName"
        #Check $VMSwitchName by using Get-VMSwitch
        $VMSwitch = Get-VMSwitch -Name $VMSwitchName -ErrorAction SilentlyContinue
        if (-not $VMSwitch) {
            throw "-VMSwitchName $VMSwitchName not found. Please check the -VMSwitchName parameter and try again."
        }
        WriteLog '-EnableVMNetworking validation complete'
    }
}
elseif ($VMSwitchName) {
    WriteLog "VM networking is disabled. Stored -VMSwitchName $VMSwitchName will not be used unless -EnableVMNetworking is `$true."
}

if (-not ($ISOPath) -and ($OptionalFeatures -like '*netfx3*')) {
    throw "netfx3 specified as an optional feature, however Windows ISO isn't defined. Unable to get netfx3 source files from downloaded ESD media. Please specify a Windows ISO in the ISOPath parameter."
}
if (($LogicalSectorSizeBytes -eq 4096) -and ($installdrivers -eq $true)) {
    $installdrivers = $false
    $CopyDrivers = $true
    WriteLog 'LogicalSectorSizeBytes is set to 4096, which is not supported for driver injection. Setting $installdrivers to $false'
    WriteLog 'As a workaround, setting -copydrivers $true to copy drivers to the deploy partition drivers folder'
    WriteLog 'We are investigating this issue and will update the script if/when we have a fix'
}
if ($BuildUSBDrive -eq $true) {
    $USBDrives, $USBDrivesCount = Get-USBDrive
}
if (($InstallApps -eq $false) -and (($UpdateLatestDefender -eq $true) -or ($UpdateOneDrive -eq $true) -or ($UpdateEdge -eq $true) -or ($UpdateLatestMSRT -eq $true) -or ($installLatestCuInVm -eq $true))) {
    WriteLog 'You have selected to update Defender, Malicious Software Removal Tool, OneDrive, Edge, or the latest Windows 10 LTSB/LTSC cumulative update, however you are setting InstallApps to false. These updates require the InstallApps variable to be set to true. Please set InstallApps to true and try again.'
    throw "InstallApps variable must be set to `$true to update Defender, OneDrive, Edge, MSRT, or the latest Windows 10 LTSB/LTSC cumulative update"
}
# if (($WindowsArch -eq 'ARM64') -and ($InstallOffice -eq $true)) {
#     $InstallOffice = $false
#     WriteLog 'M365 Apps/Office currently fails to install on ARM64 VMs without an internet connection. Setting InstallOffice to false'
# }

if (($WindowsArch -eq 'ARM64') -and ($UpdateOneDrive -eq $true)) {
    $UpdateOneDrive = $false
    WriteLog 'OneDrive currently fails to install on ARM64 VMs (even with the OneDrive ARM setup files). Setting UpdateOneDrive to false'
}

if (($WindowsArch -eq 'ARM64') -and ($UpdateLatestMSRT -eq $true)) {
    $UpdateLatestMSRT = $false
    WriteLog 'Windows Malicious Software Removal Tool is not available for the ARM64 architecture.'
}
#If downloading ESD from MCT, hardcode WindowsVersion to 22H2 for Windows 10 and 25H2 for Windows 11
#MCT media only provides 22H2 and 25H2 media
#This prevents issues with VHDX Caching unecessarily and with searching for CUs
if ($ISOPath -eq '') {
    if ($WindowsRelease -eq '10') {
        $WindowsVersion = '22H2'
    }
    if ($WindowsRelease -eq '11') {
        $WindowsVersion = '25H2'
    }
}

###END PARAMETER VALIDATION

#Get script variable values
LogVariableValues

#Check if environment is dirty
If (Test-Path -Path "$FFUDevelopmentPath\dirty.txt") {
    Get-FFUEnvironment
}
WriteLog 'Creating dirty.txt file'
$dirtyFilePath = Join-Path -Path $FFUDevelopmentPath -ChildPath 'dirty.txt'
New-Item -Path $dirtyFilePath -ItemType "file" | Out-Null

# Early CLI prompt for additional FFUs (only if enabled and not provided)
if ($BuildUSBDrive -and $CopyAdditionalFFUFiles -and ((-not $AdditionalFFUFiles) -or ($AdditionalFFUFiles.Count -eq 0))) {
    try {
        $ffuFolder = Join-Path $FFUDevelopmentPath 'FFU'
        if (Test-Path -Path $ffuFolder) {
            $cand = Get-ChildItem -Path $ffuFolder -Filter '*.ffu' -File | Sort-Object LastWriteTime -Descending
            if ($cand.Count -gt 0) {
                Write-Host ""
                Write-Host "Additional FFU files available in $($ffuFolder):"
                $i = 1
                foreach ($c in $cand) {
                    Write-Host ("{0,3}. {1}  [{2}]" -f $i, $c.Name, $c.LastWriteTime)
                    $i++
                }
                Write-Host ""
                $resp = Read-Host "Select additional FFUs to copy (e.g. 1,3,5) or 'A' for all, or press Enter to skip"
                if ($resp -match '^[Aa]$') {
                    $AdditionalFFUFiles = @($cand.FullName)
                }
                elseif ($resp -match '^\s*\d+(\s*,\s*\d+)*\s*$') {
                    $indices = $resp.Split(',') | ForEach-Object { [int]($_.Trim()) }
                    $sel = @()
                    foreach ($idx in $indices) {
                        if ($idx -ge 1 -and $idx -le $cand.Count) {
                            $sel += $cand[$idx - 1].FullName
                        }
                    }
                    $AdditionalFFUFiles = @($sel | Select-Object -Unique)
                }
                else {
                    # Skip if blank or invalid
                    if (-not [string]::IsNullOrWhiteSpace($resp)) {
                        WriteLog "Invalid additional FFU selection input. Skipping."
                    }
                }
            }
        }
    }
    catch {
        WriteLog "Early additional FFU selection prompt failed: $($_.Exception.Message)"
    }
}

#Get drivers first since user could be prompted for additional info
Set-Progress -Percentage 3 -Message "Processing drivers..."
if ($driversJsonPath -and (Test-Path $driversJsonPath) -and ($InstallDrivers -or $CopyDrivers)) {
    WriteLog "Processing drivers from JSON file: $driversJsonPath"
    Import-Module "$PSScriptRoot\FFUUI.Core\FFUUI.Core.psm1"
    # FFU.Common.Drivers.psm1 is imported by FFUUI.Core.psm1

    $driversToProcess = @()
    $jsonData = Get-Content -Path $driversJsonPath -Raw | ConvertFrom-Json

    foreach ($makeEntry in $jsonData.PSObject.Properties) {
        $makeName = $makeEntry.Name
        if ($makeEntry.Value.PSObject.Properties['Models']) {
            foreach ($modelEntry in $makeEntry.Value.Models) {
                if ($null -eq $modelEntry -or -not $modelEntry.PSObject.Properties['Name']) {
                    WriteLog "Skipping model entry for Make '$makeName' due to missing Name."
                    continue
                }

                $modelName = $modelEntry.Name
                if ([string]::IsNullOrWhiteSpace($modelName)) {
                    WriteLog "Skipping model entry for Make '$makeName' because Name is empty."
                    continue
                }
                $modelName = $modelName.Trim()

                $driverItem = $null
                switch ($makeName) {
                    'Microsoft' {
                        $driverItem = [PSCustomObject]@{
                            Make  = $makeName
                            Model = $modelName
                            Link  = if ($modelEntry.PSObject.Properties['Link']) { $modelEntry.Link } else { $null }
                        }
                    }
                    'HP' {
                        $systemId = if ($modelEntry.PSObject.Properties['SystemId'] -and -not [string]::IsNullOrWhiteSpace($modelEntry.SystemId)) { $modelEntry.SystemId.Trim() } else { $null }
                        $baseName = if ($modelEntry.PSObject.Properties['ProductName'] -and -not [string]::IsNullOrWhiteSpace($modelEntry.ProductName)) { $modelEntry.ProductName } else { $modelName }
                        if ($modelName -match '(.+?)\s*\((.+?)\)$') {
                            if ([string]::IsNullOrWhiteSpace($baseName)) { $baseName = $matches[1].Trim() }
                            if ([string]::IsNullOrWhiteSpace($systemId)) { $systemId = $matches[2].Trim() }
                        }
                        if ($baseName -match '(.+?)\s*\((.+?)\)$') {
                            $baseName = $matches[1].Trim()
                            if ([string]::IsNullOrWhiteSpace($systemId)) { $systemId = $matches[2].Trim() }
                        }
                        if ([string]::IsNullOrWhiteSpace($baseName)) { $baseName = $modelName }
                        $displayModel = if ([string]::IsNullOrWhiteSpace($systemId)) { $baseName.Trim() } else { "$($baseName.Trim()) ($($systemId.Trim()))" }
                        $driverItem = [PSCustomObject]@{
                            Make  = $makeName
                            Model = $displayModel
                        }
                        if (-not [string]::IsNullOrWhiteSpace($baseName)) {
                            $driverItem | Add-Member -NotePropertyName ProductName -NotePropertyValue $baseName.Trim()
                        }
                        if (-not [string]::IsNullOrWhiteSpace($systemId)) {
                            $driverItem | Add-Member -NotePropertyName SystemId -NotePropertyValue $systemId
                        }
                    }
                    'Lenovo' {
                        $machineType = if ($modelEntry.PSObject.Properties['MachineType']) { $modelEntry.MachineType } else { $null }
                        $productName = $modelName
                        if ([string]::IsNullOrWhiteSpace($machineType) -and $modelName -match '(.+?)\s*\((.+?)\)$') {
                            $productName = $matches[1].Trim()
                            $machineType = $matches[2].Trim()
                        }
                        if ([string]::IsNullOrWhiteSpace($machineType)) {
                            WriteLog "Skipping Lenovo model '$modelName' because MachineType is missing."
                            continue
                        }
                        $displayModel = "$productName ($machineType)"
                        $driverItem = [PSCustomObject]@{
                            Make        = $makeName
                            Model       = $displayModel
                            ProductName = $productName
                            MachineType = $machineType
                        }
                    }
                    'Dell' {
                        $systemId = if ($modelEntry.PSObject.Properties['SystemId']) { $modelEntry.SystemId } else { $null }
                        $baseName = $modelName
                        if ([string]::IsNullOrWhiteSpace($systemId) -and $modelName -match '(.+?)\s*\((.+?)\)$') {
                            $baseName = $matches[1].Trim()
                            $systemId = $matches[2].Trim()
                        }
                        $displayModel = if ([string]::IsNullOrWhiteSpace($systemId)) { $baseName } else { "$baseName ($systemId)" }
                        $driverItem = [PSCustomObject]@{
                            Make     = $makeName
                            Model    = $displayModel
                            SystemId = $systemId
                            CabUrl   = if ($modelEntry.PSObject.Properties['CabUrl']) { $modelEntry.CabUrl } else { $null }
                        }
                    }
                    default {
                        WriteLog "Skipping unsupported Make '$makeName' in Drivers.json."
                    }
                }

                if ($null -ne $driverItem) {
                    $driversToProcess += $driverItem
                }
            }
        }
    }

    if ($driversToProcess.Count -eq 0) {
        WriteLog "No drivers found to process in $driversJsonPath."
    }
    else {
        WriteLog "Found $($driversToProcess.Count) driver entries to process from $driversJsonPath."

        # Resolve effective release used for OEM driver operations.
        $localDriverWindowsRelease = Get-EffectiveDriverWindowsRelease -WindowsRelease $WindowsRelease -WindowsSKU $WindowsSKU
        if ($localDriverWindowsRelease -ne $WindowsRelease) {
            WriteLog "Normalized WindowsRelease for Drivers.json processing from $WindowsRelease to $localDriverWindowsRelease (SKU='$WindowsSKU')."
        }

        $preserveSourceOnCompress = ($UseDriversAsPEDrivers -and $CompressDownloadedDriversToWim)
        $taskArguments = @{
            DriversFolder            = $DriversFolder
            WindowsRelease           = $localDriverWindowsRelease
            WindowsArch              = $WindowsArch
            WindowsVersion           = $WindowsVersion 
            Headers                  = $Headers
            UserAgent                = $UserAgent
            CompressToWim            = $CompressDownloadedDriversToWim
            PreserveSourceOnCompress = $preserveSourceOnCompress
        }
        
        WriteLog "Starting parallel driver processing using Invoke-ParallelProcessing..."
        # Mark driver processing as active so Cancel+No can still clean partial driver artifacts.
        $driverProcessingMarker = Join-Path (Join-Path $FFUDevelopmentPath '.session') 'driverProcessing.active'
        try {
            Set-Content -Path $driverProcessingMarker -Value ((Get-Date).ToUniversalTime().ToString('o')) -Encoding UTF8
        }
        catch {
            WriteLog "Failed to create driver processing marker at $($driverProcessingMarker): $($_.Exception.Message)"
        }

        try {
            # Use the configured Threads value to control driver download concurrency
            $parallelResults = Invoke-ParallelProcessing -ItemsToProcess $driversToProcess `
                -TaskType 'DownloadDriverByMake' `
                -TaskArguments $taskArguments `
                -IdentifierProperty 'Model' `
                -WindowObject $null `
                -ListViewControl $null `
                -MainThreadLogPath $LogFile `
                -ThrottleLimit $Threads
        }
        finally {
            if (Test-Path -Path $driverProcessingMarker) {
                try {
                    Remove-Item -Path $driverProcessingMarker -Force -ErrorAction SilentlyContinue
                }
                catch {
                    WriteLog "Failed to remove driver processing marker at $($driverProcessingMarker): $($_.Exception.Message)"
                }
            }
        }

        # After processing, update the driver mapping file and detect failures
        $successfullyDownloaded = [System.Collections.Generic.List[PSCustomObject]]::new()
        $failedDownloads = [System.Collections.Generic.List[PSCustomObject]]::new()

        if ($null -ne $parallelResults) {
            # Create a lookup table from the original items to retain full metadata for mapping.
            $driverLookup = @{}
            foreach ($driver in $driversToProcess) {
                if (-not [string]::IsNullOrWhiteSpace($driver.Model)) {
                    $driverLookup[$driver.Model] = $driver
                }
            }

            foreach ($result in $parallelResults) {
                if ($null -eq $result) { continue }

                $lookupModelName = $null
                $resultStatus = $null
                $resultDriverPath = $null
                $resultSuccess = $false

                if ($result -is [hashtable]) {
                    $lookupModelName = $result['Identifier']
                    $resultStatus = $result['Status']
                    if ($result.ContainsKey('DriverPath')) { $resultDriverPath = $result['DriverPath'] }
                    if ($result.ContainsKey('Success')) {
                        $resultSuccess = [bool]$result['Success']
                    }
                    elseif ($result.ContainsKey('ResultCode')) {
                        $resultSuccess = ($result['ResultCode'] -eq 0)
                    }
                }
                elseif ($result -is [pscustomobject]) {
                    if ($result.PSObject.Properties.Name -contains 'Identifier' -and -not [string]::IsNullOrWhiteSpace($result.Identifier)) {
                        $lookupModelName = $result.Identifier
                    }
                    elseif ($result.PSObject.Properties.Name -contains 'Model' -and -not [string]::IsNullOrWhiteSpace($result.Model)) {
                        $lookupModelName = $result.Model
                    }

                    if ($result.PSObject.Properties.Name -contains 'Status') { $resultStatus = $result.Status }
                    if ($result.PSObject.Properties.Name -contains 'DriverPath') { $resultDriverPath = $result.DriverPath }
                    if ($result.PSObject.Properties.Name -contains 'Success') {
                        $resultSuccess = [bool]$result.Success
                    }
                    elseif ($result.PSObject.Properties.Name -contains 'ResultCode') {
                        $resultSuccess = ($result.ResultCode -eq 0)
                    }
                }

                if ($resultSuccess -and -not [string]::IsNullOrWhiteSpace($resultDriverPath)) {
                    $modelKey = if (-not [string]::IsNullOrWhiteSpace($lookupModelName)) { $lookupModelName } else { 'Unknown model' }
                    $driverMetadata = $null
                    if (-not [string]::IsNullOrWhiteSpace($lookupModelName) -and $driverLookup.ContainsKey($lookupModelName)) {
                        $driverMetadata = $driverLookup[$lookupModelName]
                    }

                    if ($driverMetadata) {
                        $driverRecord = [PSCustomObject]@{
                            Make       = $driverMetadata.Make
                            Model      = $modelKey
                            DriverPath = $resultDriverPath
                        }
                        if ($driverMetadata.PSObject.Properties['SystemId'] -and -not [string]::IsNullOrWhiteSpace($driverMetadata.SystemId)) {
                            $driverRecord | Add-Member -NotePropertyName SystemId -NotePropertyValue $driverMetadata.SystemId
                        }
                        if ($driverMetadata.PSObject.Properties['MachineType'] -and -not [string]::IsNullOrWhiteSpace($driverMetadata.MachineType)) {
                            $driverRecord | Add-Member -NotePropertyName MachineType -NotePropertyValue $driverMetadata.MachineType
                        }
                        if ($driverMetadata.PSObject.Properties['ProductName'] -and -not [string]::IsNullOrWhiteSpace($driverMetadata.ProductName)) {
                            $driverRecord | Add-Member -NotePropertyName ProductName -NotePropertyValue $driverMetadata.ProductName
                        }
                        $successfullyDownloaded.Add($driverRecord)
                    }
                    else {
                        WriteLog "Warning: Could not find driver metadata for successful download of model '$modelKey'. Skipping from DriverMapping.json."
                    }
                }
                else {
                    $failureModel = if (-not [string]::IsNullOrWhiteSpace($lookupModelName)) { $lookupModelName } else { 'Unknown model' }
                    $failureStatus = if (-not [string]::IsNullOrWhiteSpace($resultStatus)) { $resultStatus } else { 'Driver download failed without a status message. Check the log for details.' }
                    $failedDownloads.Add([PSCustomObject]@{
                            Model  = $failureModel
                            Status = $failureStatus
                        })
                    WriteLog "Driver download failed for '$failureModel'. Status: $failureStatus"
                }
            }
        }
        else {
            WriteLog "Invoke-ParallelProcessing returned null or no results."
        }

        if ($failedDownloads.Count -gt 0) {
            $firstFailure = $failedDownloads[0]
            $errorMessage = "Driver download failed for model '$($firstFailure.Model)': $($firstFailure.Status)"
            WriteLog $errorMessage
            throw $errorMessage
        }

        if ($successfullyDownloaded.Count -gt 0) {
            try {
                WriteLog "Updating DriverMapping.json with $($successfullyDownloaded.Count) successfully downloaded drivers."
                Update-DriverMappingJson -DownloadedDrivers $successfullyDownloaded -DriversFolder $DriversFolder
            }
            catch {
                WriteLog "Warning: Failed to update DriverMapping.json: $($_.Exception.Message)"
                # This is not a fatal error for the build process itself, so just show a warning.
                Write-Warning "The driver download process completed, but failed to update the DriverMapping.json file. Please check the log for details."
            }
        }

        WriteLog "Finished processing drivers from $driversJsonPath."
    }
}
# Existing single-model driver download logic
elseif (($Make -and $Model) -and ($InstallDrivers -or $CopyDrivers)) {
    # Resolve effective release used for OEM driver operations.
    $localDriverWindowsRelease = Get-EffectiveDriverWindowsRelease -WindowsRelease $WindowsRelease -WindowsSKU $WindowsSKU
    if ($localDriverWindowsRelease -ne $WindowsRelease) {
        WriteLog "Normalized WindowsRelease for single-model driver processing from $WindowsRelease to $localDriverWindowsRelease (SKU='$WindowsSKU')."
    }

    try {
        if ($Make -eq 'HP') {
            WriteLog 'Getting HP drivers'
            Get-HPDrivers -Make $Make -Model $Model -WindowsArch $WindowsArch -WindowsRelease $localDriverWindowsRelease -WindowsVersion $WindowsVersion
            WriteLog 'Getting HP drivers completed successfully'
        }
        if ($make -eq 'Microsoft') {
            WriteLog 'Getting Microsoft drivers'
            Get-MicrosoftDrivers -Make $Make -Model $Model -WindowsArch $WindowsArch -WindowsRelease $localDriverWindowsRelease
            WriteLog 'Getting Microsoft drivers completed successfully'
        }
        if ($make -eq 'Lenovo') {
            WriteLog 'Getting Lenovo drivers'
            Get-LenovoDrivers -Model $Model -WindowsArch $WindowsArch -WindowsRelease $localDriverWindowsRelease
            WriteLog 'Getting Lenovo drivers completed successfully'
        }
        if ($make -eq 'Dell') {
            WriteLog 'Getting Dell drivers'
            Get-DellDrivers -Model $Model -WindowsArch $WindowsArch -WindowsRelease $localDriverWindowsRelease
            WriteLog 'Getting Dell drivers completed successfully'
        }
    }
    catch {
        Writelog "Getting drivers failed with error $_"
        throw $_
    }
}
            
#Get Windows ADK
Set-Progress -Percentage 5 -Message "Validating ADK installation..."
try {
    $adkPath = Get-ADK
    #Need to use the Deployment and Imaging tools environment to use dism from the Sept 2023 ADK to optimize FFU 
    $DandIEnv = Join-Path $adkPath "Assessment and Deployment Kit\Deployment Tools\DandISetEnv.bat"

}
catch {
    WriteLog 'ADK not found'
    throw $_
}
            
#Create apps ISO for Office and/or 3rd party apps
if ($InstallApps) {
    Set-Progress -Percentage 6 -Message "Downloading and preparing applications..."
    $appsIsoRefreshRequired = -not (Test-Path -Path $AppsISO -PathType Leaf)
    if (-not $appsIsoRefreshRequired) {
        WriteLog "Apps ISO exists at: $AppsISO"

        $appIsoInputPaths = @(
            $AppListPath,
            $UserAppListPath,
            $appInstallConfigPath,
            $wingetWin32jsonFile,
            $appsScriptVarsJsonPath,
            $OfficeDownloadXML,
            $OfficeConfigXMLFile
        )

        if ($InjectUnattend) {
            $appIsoInputPaths += Get-UnattendSourcePath -UnattendFolder $UnattendFolder -WindowsArch $WindowsArch -UnattendX64FilePath $UnattendX64FilePath -UnattendArm64FilePath $UnattendArm64FilePath
        }

        $appsIsoRefreshRequired = Test-AppsIsoRefreshRequired -AppsISOPath $AppsISO -AppsPath $AppsPath -InputPaths $appIsoInputPaths
    }

    if (Test-Path -Path $wingetWin32jsonFile -PathType Leaf) {
        try {
            WriteLog "Removing generated Winget Win32 app manifest before app preparation: $wingetWin32jsonFile"
            Remove-Item -Path $wingetWin32jsonFile -Force -ErrorAction Stop
            WriteLog 'Removal complete'
        }
        catch {
            WriteLog "Failed removing generated Winget Win32 app manifest: $($_.Exception.Message)"
            if ($appsIsoRefreshRequired) { throw $_ }
        }
    }

    if ($appsIsoRefreshRequired) {
        try {
            if (Test-Path -Path $AppsISO -PathType Leaf) {
                WriteLog "Refreshing Apps ISO because app inputs changed: $AppsISO"
                Remove-Item -Path $AppsISO -Force -ErrorAction SilentlyContinue
                WriteLog 'Removal complete'
            }

            #Check for and download WinGet applications
            if (Test-Path -Path $AppListPath) {
                $appList = Get-Content -Path $AppListPath -Raw | ConvertFrom-Json
                
                WriteLog 'Checking for previously downloaded Winget applications'
            
                # Initialize variables
                $missingApps = @()
                $existingMSStoreApps = @()
                $hasExistingApps = $false
            
                # Check for WinGetWin32Apps.json
                $wingetAppsJson = $null
                if (Test-Path -Path $wingetWin32jsonFile) {
                    WriteLog "$wingetWin32jsonFile found"
                    $wingetAppsJson = Get-Content -Path $wingetWin32jsonFile -Raw | ConvertFrom-Json
                    $hasExistingApps = $true
                }
            
                # Check MSStore folder for existing apps
                if (Test-Path -Path "$AppsPath\MSStore") {
                    WriteLog "$AppsPath\MSStore folder found"
                    
                    # Get root folder names in MSStore directory
                    $MSStoreFolder = Get-ChildItem -Path "$AppsPath\MSStore" -Directory
                    
                    # Check content size of each folder
                    foreach ($folder in $MSStoreFolder) {
                        $folderSize = (Get-ChildItem -Path $folder.FullName -Recurse | Measure-Object -Property Length -Sum).Sum
                        if ($folderSize -gt 1MB) {
                            $existingMSStoreApps += $folder.Name
                            $hasExistingApps = $true
                        }
                    }
                }
            
                # If there are no existing apps, use the original AppList.json directly
                if (-not $hasExistingApps) {
                    WriteLog "No existing applications found. Using original AppList.json for all apps."
                    Get-Apps -AppList $AppListPath -AppsPath $AppsPath -UserAppListPath $UserAppListPath -WindowsArch $WindowsArch -OrchestrationPath $OrchestrationPath -LogFilePath $LogFile -ThrottleLimit $Threads
                }
                else {
                    # Compare apps in AppList.json with existing installations
                    foreach ($app in $appList.apps) {
                        $appFound = $false
                        
                        # Check Win32 apps regardless of source
                        if ($wingetAppsJson) {
                            # Check for exact match or architecture-specific entries
                            $wingetApp = $wingetAppsJson | Where-Object { 
                                $_.Name -eq $app.name -or 
                                $_.Name -eq "$($app.name) (x86)" -or 
                                $_.Name -eq "$($app.name) (x64)" -or
                                $_.Name -eq "$($app.name) (arm64)"
                            }
                            
                            if ($wingetApp) {
                                # Verify content exists in Win32 folder
                                $appFolder = Join-Path -Path "$AppsPath\Win32" -ChildPath $app.name
                                if (Test-Path -Path $appFolder) {
                                    # Check for actual files in the folder or its subdirectories
                                    $allFiles = Get-ChildItem -Path $appFolder -Recurse -File -ErrorAction SilentlyContinue
                                    
                                    if ($allFiles) {
                                        # Verify actual content size
                                        $folderSize = ($allFiles | Measure-Object -Property Length -Sum).Sum
                                        if ($folderSize -gt 1MB) {
                                            $appFound = $true
                                            WriteLog "Found existing Win32 app: $($app.name) (Size: $([math]::Round($folderSize/1MB, 2)) MB)"
                                        }
                                        else {
                                            WriteLog "Win32 app folder exists but content is too small: $($app.name) (Size: $([math]::Round($folderSize/1MB, 2)) MB)"
                                        }
                                    }
                                    else {
                                        WriteLog "Win32 app folder exists but contains no files: $($app.name)"
                                    }
                                }
                                else {
                                    WriteLog "Win32 app folder does not exist: $($app.name)"
                                }
                            }
                        }
            
                        # If not found in Win32, check MSStore folder regardless of source
                        if (-not $appFound) {
                            if ($existingMSStoreApps -contains $app.name) {
                                $appFound = $true
                                WriteLog "Found existing MSStore app: $($app.name)"
                            }
                        }
            
                        # If app not found in either location, add to missing apps list
                        if (-not $appFound) {
                            $missingApps += $app
                            WriteLog "App not found, will download: $($app.name)"
                        }
                    }
            
                    # If missing apps found, create modified AppList.json
                    if ($missingApps.Count -gt 0) {
                        $modifiedAppList = @{
                            apps = $missingApps
                        }
                        
                        $modifiedAppListPath = Join-Path -Path $AppsPath -ChildPath "ModifiedAppList.json"
                        $modifiedAppList | ConvertTo-Json | Set-Content -Path $modifiedAppListPath
                        WriteLog "Created ModifiedAppList.json with $($missingApps.Count) apps to download"
            
                        # Download missing apps
                        WriteLog "Downloading missing applications"
                        Get-Apps -AppList $modifiedAppListPath -AppsPath $AppsPath -UserAppListPath $UserAppListPath -WindowsArch $WindowsArch -OrchestrationPath $OrchestrationPath -LogFilePath $LogFile -ThrottleLimit $Threads
                        
                        # Cleanup modified app list
                        Remove-Item -Path $modifiedAppListPath -Force
                    }
                    else {
                        WriteLog "All applications already downloaded, skipping downloads"
                    }
                }
            }
            # Check if the configured BYO app list exists and output which apps will be installed.
            # It is expected that the user will have already copied the applications and created the BYO app list file.
            if (Test-Path -Path $UserAppListPath) {
                $userAppList = Get-Content -Path $UserAppListPath -Raw | ConvertFrom-Json
                WriteLog "$(Split-Path -Path $UserAppListPath -Leaf) found, the following apps will be installed:"
                foreach ($app in $userAppList) {
                    WriteLog "$($app.name)"
                }
            }

            # Remove residual update artifacts for any updates disabled via flags
            Remove-DisabledArtifacts

            #Install Office
            if ($InstallOffice) {
                #Check if Office has already been downloaded, if so, skip download
                WriteLog 'Checking for M365 Apps/Office download'
                $officeDataFolder = "$AppsPath\Office\Office\Data"
                if (Test-Path -Path $officeDataFolder) {
                    # Check the size of the $officeDataFolder folder
                    $OfficeSize = (Get-ChildItem -Path $officeDataFolder -Recurse | Measure-Object -Property Length -Sum).Sum
                    if ($OfficeSize -gt 1MB) {
                        WriteLog "Found Office download in $officeDataFolder, skipping download"
                    }
                    else {
                        WriteLog 'Downloading M365 Apps/Office'
                        Get-Office
                        WriteLog 'Downloading M365 Apps/Office completed successfully'
                    }

                }
                else {
                    WriteLog 'Downloading M365 Apps/Office'
                    Get-Office
                    WriteLog 'Downloading M365 Apps/Office completed successfully'
                }
                
            }

            #Update Latest Defender Platform and Definitions - these can't be serviced into the VHDX, will be saved to AppsPath
            if ($UpdateLatestDefender) {
                # Check if Defender has already been downloaded, if so, skip download
                WriteLog "`$UpdateLatestDefender is set to true, checking for latest Defender Platform and Security updates"
                if (Test-Path -Path $DefenderPath) {
                    # Check the size of the $DefenderPath folder
                    $DefenderSize = (Get-ChildItem -Path $DefenderPath -Recurse | Measure-Object -Property Length -Sum).Sum
                    if ($DefenderSize -gt 1MB) {
                        WriteLog "Found Defender download in $DefenderPath, skipping download"
                        $DefenderDownloaded = $true
                    }
                }
                if (-not $DefenderDownloaded) {
                    WriteLog "Creating $DefenderPath"
                    New-Item -Path $DefenderPath -ItemType Directory -Force | Out-Null

                    # Define array of updates to download
                    $defenderUpdates = @(
                        @{
                            Name        = "Update for Microsoft Defender Antivirus antimalware platform"
                            Description = "Defender Platform"
                        },
                        @{
                            Name        = "Windows Security Platform"
                            Description = "Windows Security Platform"
                        }
                    )

                    # Add 30 second delay to allow for Windows Security Platform to install
                    # I suspect this is related to AppxSVC not being immediately ready when booting to audit mode
                    # Long-term solution would be the check for AppxSVC being started, but for now the 30 second sleep seems to work consistently
                    $installDefenderCommand = "Start-Sleep -Seconds 30`r`n"

                    # Download each update
                    foreach ($update in $defenderUpdates) {
                        WriteLog "Searching for $($update.Name) from Microsoft Update Catalog and saving to $DefenderPath"
                        $KBFilePath = Save-KB -Name $update.Name -Path $DefenderPath
                        WriteLog "Latest $($update.Description) saved to $DefenderPath\$KBFilePath"
                        # Add the KB file path to the installDefenderCommand
                        $installDefenderCommand += "& `"`$env:FFUAppsRoot\Defender\$KBFilePath`"`r`n"
                    }
                
                    # Download latest Defender Definitions
                    WriteLog "Downloading latest Defender Definitions"
                    # Defender def updates can be found https://www.microsoft.com/en-us/wdsi/defenderupdates
                    if ($WindowsArch -eq 'x64') {
                        $DefenderDefURL = 'https://go.microsoft.com/fwlink/?LinkID=121721&arch=x64'
                    }
                    if ($WindowsArch -eq 'ARM64') {
                        $DefenderDefURL = 'https://go.microsoft.com/fwlink/?LinkID=121721&arch=arm64'
                    }
                    try {
                        WriteLog "Defender definitions URL is $DefenderDefURL"
                        Start-BitsTransferWithRetry -Source $DefenderDefURL -Destination "$DefenderPath\mpam-fe.exe"
                        WriteLog "Defender Definitions downloaded to $DefenderPath\mpam-fe.exe"
                        $installDefenderCommand += "& `"`$env:FFUAppsRoot\Defender\mpam-fe.exe`""
                    }
                    catch {
                        Write-Host "Downloading Defender Definitions Failed"
                        WriteLog "Downloading Defender Definitions Failed with error $_"
                        throw $_
                    }

                    # Create Update-Defender.ps1
                    WriteLog "Creating $installDefenderPath"
                    # Back up any pre-existing script with the same name before overwrite.
                    Backup-RunFile -FFUDevelopmentPath $FFUDevelopmentPath -Path $installDefenderPath
                    Set-Content -Path $installDefenderPath -Value $installDefenderCommand -Force
                    if (Test-Path -Path $installDefenderPath) {
                        WriteLog "$installDefenderPath created successfully"
                    }
                    else {
                        WriteLog "$installDefenderPath failed to create"
                        throw "$installDefenderPath failed to create"
                    }
                }
            }
            # Download latest MSRT
            if ($UpdateLatestMSRT) {
                WriteLog "`$UpdateLatestMSRT is set to true, checking for latest Windows Malicious Software Removal Tool"
                # Check if MSRT has already been downloaded, if so, skip download
                if (Test-Path -Path $MSRTPath) {
                    # Check the size of the $MSRTPath folder
                    $MSRTSize = (Get-ChildItem -Path $MSRTPath -Recurse | Measure-Object -Property Length -Sum).Sum
                    if ($MSRTSize -gt 1MB) {
                        WriteLog "Found MSRT download in $MSRTPath, skipping download"
                        $MSRTDownloaded = $true
                    }
                }
                if (-Not $MSRTDownloaded) {
                    # Create the search string for MSRT based on Windows architecture and release
                    if ($WindowsArch -eq 'x64') {
                        if ($installationType -eq 'client' -and (-not $isLTSC)) {
                            $Name = """Windows Malicious Software Removal Tool x64""" + " " + """Windows $WindowsRelease""" 
                        }
                        # Handle LTSB/LTSC
                        elseif ($installationType -eq 'client' -and $isLTSC -eq $true) {
                            $Name = """Windows Malicious Software Removal Tool x64""" + " " + """LTSB"""    
                        }
                        #Windows Server 2025 isn't listed as a product in the Microsoft Update Catalog, so we'll use the 2019 version
                        elseif ($installationType -eq 'server' -and $WindowsRelease -eq '24H2') {
                            $Name = """Windows Malicious Software Removal Tool x64""" + " " + """Windows Server 2019"""
                        }
                        else {
                            $Name = """Windows Malicious Software Removal Tool x64""" + " " + """Windows Server $WindowsRelease""" 
                        }
                    }
                    if ($WindowsArch -eq 'x86') {
                        $Name = """Windows Malicious Software Removal Tool""" + " " + """Windows $WindowsRelease""" 
                    }
                    #Check if $MSRTPath exists, if not, create it
                    if (-not (Test-Path -Path $MSRTPath)) {
                        WriteLog "Creating $MSRTPath"
                        New-Item -Path $MSRTPath -ItemType Directory -Force | Out-Null
                    }
                    WriteLog "Searching for $Name from Microsoft Update Catalog and saving to $MSRTPath"
                    WriteLog "Getting Windows Malicious Software Removal Tool URL"
                    $MSRTFileName = Save-KB -Name $Name -Path $MSRTPath
                    WriteLog "Latest Windows Malicious Software Removal Tool saved to $MSRTPath\$MSRTFileName"
                
                    # Create Update-MSRT.ps1
                    $installMSRTPath = Join-Path -Path $orchestrationPath -ChildPath "Update-MSRT.ps1"
                    WriteLog "Creating $installMSRTPath"
                    $installMSRTCommand = "& `"`$env:FFUAppsRoot\MSRT\$MSRTFileName`" /quiet"
                    # Back up any pre-existing script with the same name before overwrite.
                    Backup-RunFile -FFUDevelopmentPath $FFUDevelopmentPath -Path $installMSRTPath
                    Set-Content -Path $installMSRTPath -Value $installMSRTCommand -Force
                    # Validate that the file created successfully
                    if (Test-Path -Path $installMSRTPath) {
                        WriteLog "$installMSRTPath created successfully"
                    }
                    else {
                        WriteLog "$installMSRTPath failed to create"
                        throw "$installMSRTPath failed to create"
                    }
                }   
            }

            #Download and Install OneDrive Per Machine
            if ($UpdateOneDrive) {
                WriteLog "`$UpdateOneDrive is set to true, checking for latest OneDrive client"
                # Check if OneDrive has already been downloaded, if so, skip download
                if (Test-Path -Path $OneDrivePath) {
                    # Check the size of the $OneDrivePath folder
                    $OneDriveSize = (Get-ChildItem -Path $OneDrivePath -Recurse | Measure-Object -Property Length -Sum).Sum
                    if ($OneDriveSize -gt 1MB) {
                        WriteLog "Found OneDrive download in $OneDrivePath, skipping download"
                        $OneDriveDownloaded = $true
                    }
                }
                if (-not $OneDriveDownloaded) {
                    #Check if $OneDrivePath exists, if not, create it
                    If (-not (Test-Path -Path $OneDrivePath)) {
                        WriteLog "Creating $OneDrivePath"
                        New-Item -Path $OneDrivePath -ItemType Directory -Force | Out-Null
                    }
                    WriteLog "Downloading latest OneDrive client"
                    if ($WindowsArch -eq 'x64') {
                        $OneDriveURL = 'https://go.microsoft.com/fwlink/?linkid=844652'
                    }
                    elseif ($WindowsArch -eq 'ARM64') {
                        $OneDriveURL = 'https://go.microsoft.com/fwlink/?linkid=2271260'
                    }
                    try {
                        Start-BitsTransferWithRetry -Source $OneDriveURL -Destination "$OneDrivePath\OneDriveSetup.exe"
                        WriteLog "OneDrive client downloaded to $OneDrivePath\OneDriveSetup.exe"
                    }
                    catch {
                        Write-Host "Downloading OneDrive client Failed"
                        WriteLog "Downloading OneDrive client Failed with error $_"
                        throw $_
                    }

                    # Create Update-OneDrive.ps1
                    $installODPath = Join-Path -Path $orchestrationPath -ChildPath "Update-OneDrive.ps1"
                    WriteLog "Creating $installODPath"
                    $installODCommand = "& `"`$env:FFUAppsRoot\OneDrive\OneDriveSetup.exe`" /allusers /silent"
                    # Back up any pre-existing script with the same name before overwrite.
                    Backup-RunFile -FFUDevelopmentPath $FFUDevelopmentPath -Path $installODPath
                    Set-Content -Path $installODPath -Value $installODCommand -Force
                    # Validate that the file created successfully
                    if (Test-Path -Path $installODPath) {
                        WriteLog "$installODPath created successfully"
                    }
                    else {
                        WriteLog "$installODPath failed to create"
                        throw "$installODPath failed to create"
                    }
                }
               
            }

            #Download and Install Edge Stable
            if ($UpdateEdge) {
                WriteLog "`$UpdateEdge is set to true, checking for latest Edge Stable $WindowsArch release"
                # Check if Edge has already been downloaded, if so, skip download
                if (Test-Path -Path $EdgePath) {
                    # Check the size of the $EdgePath folder
                    $EdgeSize = (Get-ChildItem -Path $EdgePath -Recurse | Measure-Object -Property Length -Sum).Sum
                    if ($EdgeSize -gt 1MB) {
                        WriteLog "Found Edge download in $EdgePath, skipping download"
                        $EdgeDownloaded = $true
                    }
                }
                if (-not $EdgeDownloaded) {
                    # Create the search string for Edge based on Windows architecture
                    $Name = "microsoft edge stable -extended $WindowsArch"
                    #Check if $EdgePath exists, if not, create it
                    If (-not (Test-Path -Path $EdgePath)) {
                        WriteLog "Creating $EdgePath"
                        New-Item -Path $EdgePath -ItemType Directory -Force | Out-Null
                    }
                    WriteLog "Searching for $Name from Microsoft Update Catalog and saving to $EdgePath"
                    $KBFilePath = Save-KB -Name $Name -Path $EdgePath
                    $EdgeCABFilePath = "$EdgePath\$KBFilePath"
                    WriteLog "Latest Edge Stable $WindowsArch release saved to $EdgeCABFilePath"
                
                    #Extract Edge cab file to same folder as $EdgeFilePath
                    $EdgeMSIFileName = "MicrosoftEdgeEnterprise$WindowsArch.msi"
                    $EdgeFullFilePath = "$EdgePath\$EdgeMSIFileName"
                    WriteLog "Expanding $EdgeCABFilePath"
                    Invoke-Process Expand "$EdgeCABFilePath -F:*.msi $EdgeFullFilePath" | Out-Null
                    WriteLog "Expansion complete"

                    #Remove Edge CAB file
                    WriteLog "Removing $EdgeCABFilePath"
                    Remove-Item -Path $EdgeCABFilePath -Force
                    WriteLog "Removal complete"

                    # Create Update-Edge.ps1
                    $installEdgePath = Join-Path -Path $orchestrationPath -ChildPath "Update-Edge.ps1"
                    WriteLog "Creating $installEdgePath"
                    $installEdgeCommand = "& `"`$env:FFUAppsRoot\Edge\$EdgeMSIFileName`" /quiet /norestart"
                    # Back up any pre-existing script with the same name before overwrite.
                    Backup-RunFile -FFUDevelopmentPath $FFUDevelopmentPath -Path $installEdgePath
                    Set-Content -Path $installEdgePath -Value $installEdgeCommand -Force
                    # Validate that the file created successfully
                    if (Test-Path -Path $installEdgePath) {
                        WriteLog "$installEdgePath created successfully"
                    }
                    else {
                        WriteLog "$installEdgePath failed to create"
                        throw "$installEdgePath failed to create"
                    }
                }
                        
            }
        
            # Process AppsScriptVariables - Create json file
            if ($AppsScriptVariables) {
                # Back up any pre-existing JSON before overwrite.
                Backup-RunFile -FFUDevelopmentPath $FFUDevelopmentPath -Path $appsScriptVarsJsonPath
                $AppsScriptVariables | ConvertTo-Json | Out-File -FilePath $appsScriptVarsJsonPath -Encoding UTF8
                WriteLog "AppsScriptVariables exported to $appsScriptVarsJsonPath for use during orchestration"
            }
        
            #Create Apps ISO
            # Inject Unattend.xml into Apps if requested and applicable
            if ($InstallApps -and $InjectUnattend) {
                $unattendSource = Get-UnattendSourcePath -UnattendFolder $UnattendFolder -WindowsArch $WindowsArch -UnattendX64FilePath $UnattendX64FilePath -UnattendArm64FilePath $UnattendArm64FilePath

                # Ensure target folder exists under Apps
                $targetFolder = Join-Path $AppsPath 'Unattend'
                if (-not (Test-Path -Path $targetFolder -PathType Container)) {
                    New-Item -Path $targetFolder -ItemType Directory -Force | Out-Null
                }

                # Copy if source exists; otherwise log and skip
                if (Test-Path -Path $unattendSource -PathType Leaf) {
                    $destination = Join-Path $targetFolder 'Unattend.xml'
                    Save-StagedUnattendFile -SourcePath $unattendSource -DestinationPath $destination -DeviceNamingMode $DeviceNamingMode -DeviceNameTemplate $normalizedDeviceNameTemplate -WindowsArch $WindowsArch
                    WriteLog "Injected unattend file into Apps: $unattendSource -> $destination"
                }
                else {
                    WriteLog "InjectUnattend is true but source file missing: $unattendSource. Skipping unattend injection."
                }
            }
            # Stage the configured BYO app list and runtime config before creating the Apps ISO.
            Sync-UserAppListForOrchestration -SourcePath $UserAppListPath -AppsPath $AppsPath -OrchestrationPath $OrchestrationPath -AppInstallConfigPath $appInstallConfigPath

            Set-Progress -Percentage 10 -Message "Creating Apps ISO..."
            WriteLog "Creating $AppsISO file"
            New-AppsISO
            WriteLog "$AppsISO created successfully"
        }
        catch {
            Write-Host "Creating Apps ISO Failed"
            WriteLog "Creating Apps ISO Failed with error $_"
            throw $_
        }
    }
	else {
		WriteLog "Will use existing ISO"
	}
}

#Create VHDX
$windowsIsoMountedByBuild = $false
try {
    Set-Progress -Percentage 11 -Message "Checking for required Windows Updates..."
    $requiredUpdates = [System.Collections.Generic.List[pscustomobject]]::new()
    $ssuUpdateInfos = [System.Collections.Generic.List[pscustomobject]]::new()
    $cuUpdateInfos = [System.Collections.Generic.List[pscustomobject]]::new()
    $cupUpdateInfos = [System.Collections.Generic.List[pscustomobject]]::new()
    $netUpdateInfos = [System.Collections.Generic.List[pscustomobject]]::new()
    $netFeatureUpdateInfos = [System.Collections.Generic.List[pscustomobject]]::new()
    $microcodeUpdateInfos = [System.Collections.Generic.List[pscustomobject]]::new()
    $cachedIncludedUpdateNames = [System.Collections.Generic.List[string]]::new()

    $esdMetadata = $null
    $esdVersion = $null
    $cuKbWindowsVersion = $null
    $cupKbWindowsVersion = $null

    if ($WindowsRelease -eq 11 -and -not $ISOPath) {
        try {
            $esdMetadata = Get-WindowsESDMetadata -WindowsRelease $WindowsRelease -WindowsArch $WindowsArch -WindowsLang $WindowsLang -MediaType $mediaType
            if ($esdMetadata -and $esdMetadata.Version) {
                $esdVersion = $esdMetadata.Version
                WriteLog "ESD version identified as $esdVersion"
            }
            elseif ($esdMetadata) {
                WriteLog "ESD metadata resolved but no version could be parsed from filename."
            }
        }
        catch {
            WriteLog "Failed to resolve Windows ESD metadata: $($_.Exception.Message)"
        }
    }

    if ($UpdateLatestCU -or $UpdatePreviewCU -or $UpdateLatestNet -or $UpdateLatestMicrocode) {
        # Determine required updates without downloading them yet
        $cuKbArticleId = $null
        $cupKbArticleId = $null
        $netKbArticleId = $null

        if ($UpdateLatestCU -and -not $UpdatePreviewCU) {
            Writelog "`$UpdateLatestCU is set to true, checking for latest CU"
            if ($WindowsRelease -in 10, 11) { $Name = """Cumulative update for Windows $WindowsRelease Version $WindowsVersion for $WindowsArch""" }
            if ($WindowsRelease -eq 2025) { $Name = """Cumulative Update for Microsoft server operating system, version 24h2 for $WindowsArch""" }
            if ($WindowsRelease -eq 2022) { $Name = """Cumulative Update for Microsoft server operating system, version 21h2 for $WindowsArch""" }
            if ($WindowsRelease -in 2016, 2019 -and $installationType -eq "Server") { $Name = """Cumulative update for Windows Server $WindowsRelease for $WindowsArch""" }
            if ($WindowsRelease -in 2016, 2019, 2021 -and $isLTSC) {
                $today = Get-Date; $firstDayOfMonth = Get-Date -Year $today.Year -Month $today.Month -Day 1; $secondTuesday = $firstDayOfMonth.AddDays(((2 - [int]$firstDayOfMonth.DayOfWeek + 7) % 7) + 7); $updateDate = if ($today -gt $secondTuesday) { $today } else { $today.AddMonths(-1) }
                $Name = """$($updateDate.ToString('yyyy-MM')) Cumulative update for Windows 10 Version $WindowsVersion for $WindowsArch"""
            }
            if ($WindowsRelease -eq 2024 -and $isLTSC) { $Name = """Cumulative update for Windows 11 Version $WindowsVersion for $WindowsArch""" }
            
            if (($WindowsRelease -eq 2016 -and $installationType -eq "Server") -or ($WindowsRelease -in 2016, 2019, 2021 -and $isLTSC)) {
                $SSUName = if ($isLTSC) { """Servicing Stack Update for Windows 10 Version $WindowsVersion for $WindowsArch""" } else { """Servicing stack update for Windows Server $WindowsRelease for $WindowsArch""" }
                WriteLog "Searching for $SSUName from Microsoft Update Catalog"
                (Get-UpdateFileInfo -Name $SSUName) | ForEach-Object { $ssuUpdateInfos.Add($_) }
            }
            WriteLog "Searching for $Name from Microsoft Update Catalog"
            (Get-UpdateFileInfo -Name $Name) | ForEach-Object { $cuUpdateInfos.Add($_) }
            $cuKbArticleId = $global:LastKBArticleID
            $cuKbWindowsVersion = $global:LastKBWindowsVersion
        }

        if ($UpdatePreviewCU -and $installationType -eq 'Client' -and $WindowsSKU -notlike "*LTSC") {
            Writelog "`$UpdatePreviewCU is set to true, checking for latest Preview CU"
            $Name = """Cumulative update Preview for Windows $WindowsRelease Version $WindowsVersion for $WindowsArch"""
            WriteLog "Searching for $Name from Microsoft Update Catalog"
            (Get-UpdateFileInfo -Name $Name) | ForEach-Object { $cupUpdateInfos.Add($_) }
            $cupKbArticleId = $global:LastKBArticleID
            $cupKbWindowsVersion = $global:LastKBWindowsVersion
        }

        if ($UpdateLatestNet) {
            Writelog "`$UpdateLatestNet is set to true, checking for latest .NET Framework"
            if ($WindowsRelease -in 2016, 2019, 2021 -and $isLTSC) {
                if ($ssuUpdateInfos.Count -eq 0) {
                    $SSUName = """Servicing Stack Update for Windows 10 Version $WindowsVersion for $WindowsArch"""
                    WriteLog "Searching for $SSUName from Microsoft Update Catalog"
                    (Get-UpdateFileInfo -Name $SSUName) | ForEach-Object { $ssuUpdateInfos.Add($_) }
                }
                if ($WindowsRelease -in 2016) { $name = """Cumulative Update for .NET Framework 4.8 for Windows 10 version $WindowsVersion for $WindowsArch""" }
                if ($WindowsRelease -eq 2019) { $name = """Cumulative Update for .NET Framework 3.5, 4.7.2 and 4.8 for Windows 10 Version $WindowsVersion for $WindowsArch""" }
                if ($WindowsRelease -eq 2021) { $name = """Cumulative Update for .NET Framework 3.5, 4.8 and 4.8.1 for Windows 10 Version $WindowsVersion for $WindowsArch""" }
                WriteLog "Searching for $name from Microsoft Update Catalog"
                (Get-UpdateFileInfo -Name $name) | ForEach-Object { $netUpdateInfos.Add($_) }
                $netKbArticleId = $global:LastKBArticleID

                if ($WindowsRelease -eq 2021) { $NETFeatureName = """Microsoft .NET Framework 4.8.1 for Windows 10 Version 21H2 for x64""" }
                if ($WindowsRelease -in 2016, 2019) { $NETFeatureName = """Microsoft .NET Framework 4.8 for Windows 10 Version $WindowsVersion and Windows Server $WindowsRelease for x64""" }
                WriteLog "Checking for latest .NET Framework feature pack: $NETFeatureName"
                (Get-UpdateFileInfo -Name $NETFeatureName) | ForEach-Object { $netFeatureUpdateInfos.Add($_) }
            }
            else {
                if ($WindowsRelease -eq 2024 -and $isLTSC) { $Name = "Cumulative update for .NET framework windows 11 $WindowsVersion $WindowsArch -preview" }
                if ($WindowsRelease -in 10, 11) { $Name = "Cumulative update for .NET framework windows $WindowsRelease $WindowsVersion $WindowsArch -preview" }
                if ($WindowsRelease -eq 2025 -and $installationType -eq "Server") { $Name = """Cumulative Update for .NET Framework"" ""3.5 and 4.8.1"" for Windows 11 24H2 x64 -preview" }
                if ($WindowsRelease -eq 2022 -and $installationType -eq "Server") { $Name = """Cumulative Update for .NET Framework 3.5, 4.8 and 4.8.1"" ""operating system version 21H2 for x64""" }
                if ($WindowsRelease -eq 2019 -and $installationType -eq "Server") { $Name = """Cumulative Update for .NET Framework 3.5, 4.7.2 and 4.8 for Windows Server 2019 for x64""" }
                if ($WindowsRelease -eq 2016 -and $installationType -eq "Server") { $Name = """Cumulative Update for .NET Framework 4.8 for Windows Server 2016 for x64""" }
                WriteLog "Searching for $Name from Microsoft Update Catalog"
                (Get-UpdateFileInfo -Name $Name) | ForEach-Object { $netUpdateInfos.Add($_) }
                $netKbArticleId = $global:LastKBArticleID
            }
        }

        if ($UpdateLatestMicrocode -and $WindowsRelease -in 2016, 2019) {
            WriteLog "`$UpdateLatestMicrocode is set to true, checking for latest Microcode"
            if ($WindowsRelease -eq 2016) { $name = "KB4589210 $windowsArch" }
            if ($WindowsRelease -eq 2019) { $name = "KB4589208 $windowsArch" }
            WriteLog "Searching for $name from Microsoft Update Catalog"
            (Get-UpdateFileInfo -Name $name) | ForEach-Object { $microcodeUpdateInfos.Add($_) }
        }
        
        $esdVerObj = $null
        $cuVerObj = $null
        $cupVerObj = $null
        if ($esdVersion) { try { $esdVerObj = [version]$esdVersion } catch { } }
        if ($cuKbWindowsVersion) { try { $cuVerObj = [version]$cuKbWindowsVersion } catch { } }
        if ($cupKbWindowsVersion) { try { $cupVerObj = [version]$cupKbWindowsVersion } catch { } }
        
        if ($esdVerObj -and $cuVerObj) {
            if ($esdVerObj -eq $cuVerObj -or $esdVerObj -gt $cuVerObj) {
                $skipReason = if ($esdVerObj -eq $cuVerObj) { 'matches' } else { 'is newer than' }
                WriteLog "Windows 11 ESD version $esdVersion $skipReason CU version $cuKbWindowsVersion. Skipping CU download and installation."
                if ($AllowVHDXCaching -and $cuUpdateInfos -and $cuUpdateInfos.Count -gt 0) {
                    foreach ($cuUpdateInfo in $cuUpdateInfos) {
                        if (-not [string]::IsNullOrWhiteSpace($cuUpdateInfo.Name) -and -not $cachedIncludedUpdateNames.Contains($cuUpdateInfo.Name)) {
                            $cachedIncludedUpdateNames.Add($cuUpdateInfo.Name)
                        }
                    }
                }
                $cuUpdateInfos.Clear()
                $UpdateLatestCU = $false
                $CUPath = $null
            }
        }
        if ($esdVerObj -and $cupVerObj) {
            if ($esdVerObj -eq $cupVerObj -or $esdVerObj -gt $cupVerObj) {
                $skipReason = if ($esdVerObj -eq $cupVerObj) { 'matches' } else { 'is newer than' }
                WriteLog "Windows 11 ESD version $esdVersion $skipReason Preview CU version $cupKbWindowsVersion. Skipping Preview CU download and installation."
                if ($AllowVHDXCaching -and $cupUpdateInfos -and $cupUpdateInfos.Count -gt 0) {
                    foreach ($cupUpdateInfo in $cupUpdateInfos) {
                        if (-not [string]::IsNullOrWhiteSpace($cupUpdateInfo.Name) -and -not $cachedIncludedUpdateNames.Contains($cupUpdateInfo.Name)) {
                            $cachedIncludedUpdateNames.Add($cupUpdateInfo.Name)
                        }
                    }
                }
                $cupUpdateInfos.Clear()
                $UpdatePreviewCU = $false
                $CUPPath = $null
            }
        }
        
        $requiredUpdates.AddRange($ssuUpdateInfos)
        $requiredUpdates.AddRange($cuUpdateInfos)
        $requiredUpdates.AddRange($cupUpdateInfos)
        $requiredUpdates.AddRange($netUpdateInfos)
        $requiredUpdates.AddRange($netFeatureUpdateInfos)
        $requiredUpdates.AddRange($microcodeUpdateInfos)
    }

    #Search for cached VHDX and skip VHDX creation if there's a cached version
    if ($AllowVHDXCaching) {
        WriteLog 'AllowVHDXCaching is true, checking for cached VHDX file'
        if (Test-Path -Path $VHDXCacheFolder) {
            WriteLog "Found $VHDXCacheFolder"
            $vhdxJsons = @(Get-ChildItem -File -Path $VHDXCacheFolder -Filter '*_config.json' | Sort-Object -Property CreationTime -Descending)
            WriteLog "Found $($vhdxJsons.Count) cached VHDX config files"
            
            # Build comparison list from update names and cached names
            $requiredUpdateFileNames = @()
            if ($requiredUpdates.Count -gt 0) {
                $requiredUpdateFileNames += $requiredUpdates | ForEach-Object {
                    if (-not [string]::IsNullOrWhiteSpace($_.Name)) {
                        $_.Name
                    }
                    elseif (-not [string]::IsNullOrWhiteSpace($_.Url)) {
                        ($_.Url -split '/')[-1]
                    }
                }
            }
            if ($cachedIncludedUpdateNames.Count -gt 0) {
                $requiredUpdateFileNames += $cachedIncludedUpdateNames
            }
            if ($requiredUpdateFileNames.Count -gt 0) {
                $requiredUpdateFileNames = @($requiredUpdateFileNames | Where-Object { $_ } | Sort-Object -Unique)
            }

            foreach ($vhdxJson in $vhdxJsons) {
                try {
                    WriteLog "Processing $($vhdxJson.FullName)"
                    $vhdxCacheItem = Get-Content -Path $vhdxJson.FullName -Raw | ConvertFrom-Json

                    if ($vhdxCacheItem.WindowsSKU -ne $WindowsSKU) { WriteLog 'WindowsSKU mismatch, continuing'; continue }
                    if ($vhdxCacheItem.LogicalSectorSizeBytes -ne $LogicalSectorSizeBytes) { WriteLog 'LogicalSectorSizeBytes mismatch, continuing'; continue }
                    if ($vhdxCacheItem.WindowsRelease -ne $WindowsRelease) { WriteLog 'WindowsRelease mismatch, continuing'; continue }
                    if ($vhdxCacheItem.WindowsVersion -ne $WindowsVersion) { WriteLog 'WindowsVersion mismatch, continuing'; continue }
                    if ($vhdxCacheItem.OptionalFeatures -ne $OptionalFeatures) { WriteLog 'OptionalFeatures mismatch, continuing'; continue }
                    if ($vhdxCacheItem.PSObject.Properties.Name -notcontains 'Disksize') { WriteLog 'Disksize missing in cached config, continuing'; continue }
                    [uint64]$cachedDisksize = 0
                    if (-not [uint64]::TryParse([string]$vhdxCacheItem.Disksize, [ref]$cachedDisksize)) { WriteLog "Disksize invalid in cached config ($($vhdxCacheItem.Disksize)), continuing"; continue }
                    if ($cachedDisksize -ne $Disksize) { WriteLog "Disksize mismatch (cached: $cachedDisksize, current: $Disksize), continuing"; continue }
                    if ($vhdxCacheItem.PSObject.Properties.Name -notcontains 'CreateRecoveryPartition') { WriteLog 'CreateRecoveryPartition missing in cached config, continuing'; continue }
                    if ([bool]$vhdxCacheItem.CreateRecoveryPartition -ne $CreateRecoveryPartition) { WriteLog "CreateRecoveryPartition mismatch (cached: $($vhdxCacheItem.CreateRecoveryPartition), current: $CreateRecoveryPartition), continuing"; continue }
                    if ($vhdxCacheItem.PSObject.Properties.Name -notcontains 'PartitionLayoutSignatureVersion') { WriteLog 'PartitionLayoutSignatureVersion missing in cached config, continuing'; continue }
                    if ([uint32]$vhdxCacheItem.PartitionLayoutSignatureVersion -ne $partitionLayoutSignatureVersion) { WriteLog "PartitionLayoutSignatureVersion mismatch (cached: $($vhdxCacheItem.PartitionLayoutSignatureVersion), current: $partitionLayoutSignatureVersion), continuing"; continue }
                    if ($vhdxCacheItem.PSObject.Properties.Name -notcontains 'PartitionLayoutSignature') { WriteLog 'PartitionLayoutSignature missing in cached config, continuing'; continue }
                    if ($vhdxCacheItem.PartitionLayoutSignature -ne $partitionLayoutSignature) { WriteLog 'PartitionLayoutSignature mismatch, continuing'; continue }

                    $cachedUpdateNames = @()
                    if ($vhdxCacheItem.IncludedUpdates -and $vhdxCacheItem.IncludedUpdates.Count -gt 0) {
                        $cachedUpdateNames = @($vhdxCacheItem.IncludedUpdates.Name | Sort-Object)
                    }
                    
                    # Manually compare the two sorted arrays of update names
                    $updatesMatch = $false
                    if ($requiredUpdateFileNames.Count -eq $cachedUpdateNames.Count) {
                        $updatesMatch = $true # Assume true and prove false
                        for ($i = 0; $i -lt $requiredUpdateFileNames.Count; $i++) {
                            if ($requiredUpdateFileNames[$i] -ne $cachedUpdateNames[$i]) {
                                $updatesMatch = $false
                                break
                            }
                        }
                    }

                    if (-not $updatesMatch) {
                        WriteLog 'IncludedUpdates mismatch, continuing'
                        continue
                    }

                    WriteLog "Found cached VHDX file $vhdxCacheFolder\$($vhdxCacheItem.VhdxFileName) with matching parameters and included updates"
                    $cachedVHDXFileFound = $true
                    $cachedVHDXInfo = $vhdxCacheItem
                    break
                }
                catch {
                    WriteLog "Reading $vhdxJson Failed with error $_"
                }
            }
        }
    }

    # If no cached VHDX is found, download the required updates now
    if (-Not $cachedVHDXFileFound -and $requiredUpdates.Count -gt 0) {
        Set-Progress -Percentage 12 -Message "Downloading Windows Updates..."
        WriteLog "No suitable VHDX cache found. Downloading $($requiredUpdates.Count) update(s)."
        
        If (-not (Test-Path -Path $KBPath)) {
            WriteLog "Creating $KBPath"; New-Item -Path $KBPath -ItemType Directory -Force | Out-Null
        }

        # Create an OS/version-scoped KB cache folder under KB
        # This allows caching across multiple Windows targets without mixing MSUs
        $kbCacheReleaseFolder = $null
        $kbCacheVersionFolder = $null

        if ($installationType -eq 'Server') {
            $kbCacheReleaseFolder = "Server$WindowsRelease"
            $kbCacheVersionFolder = $null
        }
        else {
            # Client
            if ($WindowsRelease -eq 10) {
                $kbCacheReleaseFolder = 'Windows10'
                $kbCacheVersionFolder = $WindowsVersion
            }
            elseif ($WindowsRelease -eq 11) {
                $kbCacheReleaseFolder = 'Windows11'
                $kbCacheVersionFolder = $WindowsVersion
            }
            elseif ($isLTSC -and $WindowsRelease -in 2016, 2019, 2021) {
                $kbCacheReleaseFolder = 'Windows10'
                $kbCacheVersionFolder = "LTSC$WindowsRelease"
            }
            elseif ($isLTSC -and $WindowsRelease -eq 2024) {
                # Windows 11 LTSC 2024 shares the same CU branch as Windows 11 24H2/25H2
                $kbCacheReleaseFolder = 'Windows11'
                $kbCacheVersionFolder = '24H2'
            }
            else {
                $kbCacheReleaseFolder = "Windows$WindowsRelease"
                $kbCacheVersionFolder = $WindowsVersion
            }
        }

        # Force Windows 11 25H2 to share Windows 11 24H2 cache folder (same CU branch)
        if ($kbCacheReleaseFolder -eq 'Windows11' -and $kbCacheVersionFolder -match '(?i)^25H2$') {
            $kbCacheVersionFolder = '24H2'
        }

        # For Server, do not create a WindowsVersion subfolder (unnecessary for Server releases)
        if ($installationType -eq 'Server') {
            $kbCacheRootPath = Join-Path -Path $KBPath -ChildPath $kbCacheReleaseFolder
        }
        else {
            if ([string]::IsNullOrWhiteSpace($kbCacheVersionFolder)) {
                $kbCacheVersionFolder = 'UnknownVersion'
            }
            $kbCacheRootPath = Join-Path -Path (Join-Path -Path $KBPath -ChildPath $kbCacheReleaseFolder) -ChildPath $kbCacheVersionFolder
        }

        if (-not (Test-Path -Path $kbCacheRootPath)) {
            WriteLog "Creating KB cache folder $kbCacheRootPath"
            New-Item -Path $kbCacheRootPath -ItemType Directory -Force | Out-Null
        }

        # Remove older MSU files for update types included in the current run
        # This avoids DISM considering multiple stale MSUs as local sources during servicing
        try {
            # Build allow-lists from update metadata for this run
            $expectedWindowsMsuNames = [System.Collections.Generic.HashSet[string]]::new([System.StringComparer]::OrdinalIgnoreCase)
            foreach ($updateInfo in $ssuUpdateInfos) {
                if (-not [string]::IsNullOrWhiteSpace($updateInfo.Name) -and $updateInfo.Name.ToLowerInvariant().EndsWith('.msu')) {
                    [void]$expectedWindowsMsuNames.Add($updateInfo.Name)
                }
            }
            foreach ($updateInfo in $cuUpdateInfos) {
                if (-not [string]::IsNullOrWhiteSpace($updateInfo.Name) -and $updateInfo.Name.ToLowerInvariant().EndsWith('.msu')) {
                    [void]$expectedWindowsMsuNames.Add($updateInfo.Name)
                }
            }
            foreach ($updateInfo in $cupUpdateInfos) {
                if (-not [string]::IsNullOrWhiteSpace($updateInfo.Name) -and $updateInfo.Name.ToLowerInvariant().EndsWith('.msu')) {
                    [void]$expectedWindowsMsuNames.Add($updateInfo.Name)
                }
            }

            $expectedNetMsuNames = [System.Collections.Generic.HashSet[string]]::new([System.StringComparer]::OrdinalIgnoreCase)
            foreach ($updateInfo in $netUpdateInfos) {
                if (-not [string]::IsNullOrWhiteSpace($updateInfo.Name) -and $updateInfo.Name.ToLowerInvariant().EndsWith('.msu')) {
                    [void]$expectedNetMsuNames.Add($updateInfo.Name)
                }
            }
            foreach ($updateInfo in $netFeatureUpdateInfos) {
                if (-not [string]::IsNullOrWhiteSpace($updateInfo.Name) -and $updateInfo.Name.ToLowerInvariant().EndsWith('.msu')) {
                    [void]$expectedNetMsuNames.Add($updateInfo.Name)
                }
            }

            $expectedMicrocodeMsuNames = [System.Collections.Generic.HashSet[string]]::new([System.StringComparer]::OrdinalIgnoreCase)
            foreach ($updateInfo in $microcodeUpdateInfos) {
                if (-not [string]::IsNullOrWhiteSpace($updateInfo.Name) -and $updateInfo.Name.ToLowerInvariant().EndsWith('.msu')) {
                    [void]$expectedMicrocodeMsuNames.Add($updateInfo.Name)
                }
            }

            # Prune older Windows update MSUs (CU/SSU/Preview CU) from this OS cache root (exclude .NET "ndp" MSUs)
            if ($expectedWindowsMsuNames.Count -gt 0) {
                WriteLog "Pruning older Windows update MSU files in $kbCacheRootPath"
                $existingWindowsMsus = @(Get-ChildItem -Path $kbCacheRootPath -Filter '*.msu' -File -ErrorAction SilentlyContinue | Where-Object {
                        $_.Name -match '(?i)^windows\d+\.0-kb\d+.*\.msu$' -and $_.Name -notmatch '(?i)ndp'
                    })
                foreach ($msu in $existingWindowsMsus) {
                    if (-not $expectedWindowsMsuNames.Contains($msu.Name)) {
                        WriteLog "Removing old Windows update MSU: $($msu.FullName)"
                        Remove-Item -LiteralPath $msu.FullName -Force -ErrorAction SilentlyContinue
                    }
                }
            }

            # Prune older .NET MSUs from NET folder (always store .NET updates under NET)
            if ($expectedNetMsuNames.Count -gt 0) {
                $netFolder = Join-Path -Path $kbCacheRootPath -ChildPath 'NET'

                # Ensure NET folder exists
                if (-not (Test-Path -Path $netFolder)) {
                    WriteLog "Creating .NET updates folder $netFolder"
                    New-Item -Path $netFolder -ItemType Directory -Force | Out-Null
                }

                # Migrate legacy layout: move expected .NET MSUs from cache root into NET folder (if present)
                $legacyRootNetMsus = @(Get-ChildItem -Path $kbCacheRootPath -Filter '*.msu' -File -ErrorAction SilentlyContinue | Where-Object { $_.Name -match '(?i)ndp' })
                foreach ($msu in $legacyRootNetMsus) {
                    $destPath = Join-Path -Path $netFolder -ChildPath $msu.Name
                    if ($expectedNetMsuNames.Contains($msu.Name) -and -not (Test-Path -LiteralPath $destPath)) {
                        WriteLog "Moving .NET MSU into NET folder: $($msu.FullName) -> $destPath"
                        Move-Item -LiteralPath $msu.FullName -Destination $destPath -Force -ErrorAction SilentlyContinue
                    }
                }

                # Prune any stale .NET MSUs from NET folder
                WriteLog "Pruning older .NET MSU files in $netFolder"
                $existingNetMsus = @(Get-ChildItem -Path $netFolder -Filter '*.msu' -File -ErrorAction SilentlyContinue)
                foreach ($msu in $existingNetMsus) {
                    if (-not $expectedNetMsuNames.Contains($msu.Name)) {
                        WriteLog "Removing old .NET MSU: $($msu.FullName)"
                        Remove-Item -LiteralPath $msu.FullName -Force -ErrorAction SilentlyContinue
                    }
                }

                # Remove stale legacy .NET MSUs that may still exist in the cache root
                $legacyRootNetMsus = @(Get-ChildItem -Path $kbCacheRootPath -Filter '*.msu' -File -ErrorAction SilentlyContinue | Where-Object { $_.Name -match '(?i)ndp' })
                foreach ($msu in $legacyRootNetMsus) {
                    if (-not $expectedNetMsuNames.Contains($msu.Name)) {
                        WriteLog "Removing old .NET MSU: $($msu.FullName)"
                        Remove-Item -LiteralPath $msu.FullName -Force -ErrorAction SilentlyContinue
                    }
                }
            }

            # Prune older Microcode MSUs from Microcode folder (only when Microcode is part of this run)
            if ($expectedMicrocodeMsuNames.Count -gt 0) {
                $microcodeFolder = Join-Path -Path $kbCacheRootPath -ChildPath 'Microcode'
                if (Test-Path -Path $microcodeFolder) {
                    WriteLog "Pruning older Microcode MSU files in $microcodeFolder"
                    $existingMicrocodeMsus = @(Get-ChildItem -Path $microcodeFolder -Filter '*.msu' -File -ErrorAction SilentlyContinue)
                    foreach ($msu in $existingMicrocodeMsus) {
                        if (-not $expectedMicrocodeMsuNames.Contains($msu.Name)) {
                            WriteLog "Removing old Microcode MSU: $($msu.FullName)"
                            Remove-Item -LiteralPath $msu.FullName -Force -ErrorAction SilentlyContinue
                        }
                    }
                }
            }
        }
        catch {
            WriteLog "Failed to prune old MSU files in $($kbCacheRootPath): $($_.Exception.Message). Continuing."
        }

        foreach ($update in $requiredUpdates) {
            $destinationPath = $kbCacheRootPath
            if (($netUpdateInfos -and ($netUpdateInfos.Name -contains $update.Name)) -or `
                ($netFeatureUpdateInfos -and ($netFeatureUpdateInfos.Name -contains $update.Name))) {
                $destinationPath = Join-Path -Path $kbCacheRootPath -ChildPath "NET"
            }
            if ($microcodeUpdateInfos -and ($microcodeUpdateInfos.Name -contains $update.Name)) {
                $destinationPath = Join-Path -Path $kbCacheRootPath -ChildPath "Microcode"
            }

            if (-not (Test-Path -Path $destinationPath)) {
                New-Item -Path $destinationPath -ItemType Directory -Force | Out-Null
            }

            # Skip download if expected file is already present
            $expectedFilePath = Join-Path -Path $destinationPath -ChildPath $update.Name

            if (Test-Path -LiteralPath $expectedFilePath) {
                WriteLog "Update already exists at $expectedFilePath, skipping download"
                continue
            }

            WriteLog "Downloading $($update.Name) to $destinationPath"
            Start-BitsTransferWithRetry -Source $update.Url -Destination $destinationPath
        }
    

        # Set file path variables for the patching process
        if ($ssuUpdateInfos.Count -gt 0) {
            $SSUFile = $ssuUpdateInfos[0].Name
            $SSUFilePath = Join-Path -Path $kbCacheRootPath -ChildPath $SSUFile
            WriteLog "Latest SSU identified as $SSUFilePath"
        }
        if ($cuUpdateInfos.Count -gt 0) {
            if (-not $CUPath) {
                $CUPath = (Get-ChildItem -Path $kbCacheRootPath -Filter "*$cuKbArticleId*" -Recurse | Select-Object -First 1).FullName
            }
            WriteLog "Latest CU identified as $CUPath"
        }
        if ($cupUpdateInfos.Count -gt 0) {
            if (-not $CUPPath) {
                $CUPPath = (Get-ChildItem -Path $kbCacheRootPath -Filter "*$cupKbArticleId*" -Recurse | Select-Object -First 1).FullName
            }
            WriteLog "Latest Preview CU identified as $CUPPath"
        }
        if ($netUpdateInfos.Count -gt 0 -or $netFeatureUpdateInfos.Count -gt 0) {
            $NETPath = Join-Path -Path $kbCacheRootPath -ChildPath "NET"
            if (-not (Test-Path -Path $NETPath)) {
                WriteLog "Creating .NET updates folder $NETPath"
                New-Item -Path $NETPath -ItemType Directory -Force | Out-Null
            }
            WriteLog ".NET updates are in $NETPath"
        }
        if ($microcodeUpdateInfos.Count -gt 0) {
            $MicrocodePath = Join-Path -Path $kbCacheRootPath -ChildPath "Microcode"
            WriteLog "Microcode updates are in $MicrocodePath"
        }
    }
    
    if (-Not $cachedVHDXFileFound) {
        Set-Progress -Percentage 15 -Message "Creating VHDX and applying base Windows image..."
        if ($ISOPath) {
			$reservedBuildDriveLetters = @($partitionDriveLetters.PSObject.Properties | ForEach-Object { [string]$_.Value })
			$wimPath = Get-WimFromISO -ImagePath $ISOPath -ExcludedDriveLetters $reservedBuildDriveLetters
			$windowsIsoMountedByBuild = $true
        }
        else {
            $wimPath = Get-WindowsESD -WindowsRelease $WindowsRelease -WindowsArch $WindowsArch -WindowsLang $WindowsLang -MediaType $mediaType -Metadata $esdMetadata
        }
        #If index not specified by user, try and find based on WindowsSKU
        if (-not($index) -and ($WindowsSKU)) {
            $requestedWindowsSKU = $WindowsSKU
            $previousInstallationType = $installationType
            $previousWindowsVersion = $WindowsVersion
            $previousIsLTSC = [bool]$isLTSC
            $windowsImageSelection = Get-WindowsImageSelection -WindowsImagePath $wimPath -WindowsSKU $WindowsSKU -WindowsRelease $WindowsRelease
            $index = $windowsImageSelection.ImageIndex

            if (-not [string]::IsNullOrWhiteSpace($windowsImageSelection.ResolvedWindowsSKU)) {
                $WindowsSKU = $windowsImageSelection.ResolvedWindowsSKU
                $windowsTargetRuntimeState = Get-WindowsTargetRuntimeState -WindowsRelease $WindowsRelease -WindowsSKU $WindowsSKU -CurrentWindowsVersion $WindowsVersion -UpdateLatestCU:$UpdateLatestCU
                $installationType = $windowsTargetRuntimeState.InstallationType
                $WindowsVersion = $windowsTargetRuntimeState.WindowsVersion
                $isLTSC = $windowsTargetRuntimeState.IsLTSC
                $isWindows10LtscClient = $windowsTargetRuntimeState.IsWindows10LtscClient
                $installLatestCuInVm = $windowsTargetRuntimeState.InstallLatestCuInVm

                if ($requestedWindowsSKU -ne $WindowsSKU) {
                    WriteLog "Resolved WindowsSKU from '$requestedWindowsSKU' to '$WindowsSKU' based on image selection '$($windowsImageSelection.ImageName)'."
                }

                if (($previousInstallationType -ne $installationType) -or ($previousWindowsVersion -ne $WindowsVersion) -or ($previousIsLTSC -ne [bool]$isLTSC)) {
                    WriteLog "Updated Windows target state after image selection: InstallationType='$installationType', WindowsVersion='$WindowsVersion', IsLTSC='$isLTSC'."
                }

                if (($InstallApps -eq $false) -and ($installLatestCuInVm -eq $true)) {
                    WriteLog 'You have selected to update Defender, Malicious Software Removal Tool, OneDrive, Edge, or the latest Windows 10 LTSB/LTSC cumulative update, however you are setting InstallApps to false. These updates require the InstallApps variable to be set to true. Please set InstallApps to true and try again.'
                    throw "InstallApps variable must be set to `$true to update Defender, OneDrive, Edge, MSRT, or the latest Windows 10 LTSB/LTSC cumulative update"
                }
            }
            else {
                WriteLog "Could not resolve a friendly WindowsSKU for selected image '$($windowsImageSelection.ImageName)'. Continuing with requested SKU '$WindowsSKU'."
            }
        }

        $vhdxDisk = New-ScratchVhdx -VhdxPath $VHDXPath -SizeBytes $disksize -LogicalSectorSizeBytes $LogicalSectorSizeBytes

        $createdSystemPartitionDriveLetter = New-SystemPartition -VhdxDisk $vhdxDisk -DriveLetter $SystemPartitionDriveLetter
    
        New-MSRPartition -VhdxDisk $vhdxDisk
    
        Set-Progress -Percentage 16 -Message "Applying base Windows image to VHDX..."
        $osPartition = New-OSPartition -VhdxDisk $vhdxDisk -OSPartitionSize $OSPartitionSize -WimPath $WimPath -WimIndex $index -DriveLetter $WindowsPartitionDriveLetter
        $osPartitionObject = $osPartition[1]
        if ($OSPartitionSize -le 0) {
			$osPartitionObject = Resize-OSPartitionForDataPartitions -OsPartition $osPartitionObject -DataPartitions $normalizedAdditionalDataPartitions
		}
        $osPartitionDriveLetter = $osPartitionObject.DriveLetter
        $WindowsPartition = $osPartitionDriveLetter + ':\'

        if ($CreateRecoveryPartition) {
            $recoveryPartition = New-RecoveryPartition -VhdxDisk $vhdxDisk -OsPartition $osPartitionObject -RecoveryPartitionSize $RecoveryPartitionSize -DriveLetter $RecoveryPartitionDriveLetter -OsPartitionUsesMaximumSize ($OSPartitionSize -le 0)
        }
        else {
            WriteLog 'CreateRecoveryPartition is false. Skipping Windows Recovery partition creation.'
        }

        foreach ($additionalDataPartition in $normalizedAdditionalDataPartitions) {
            New-DataPartition -VhdxDisk $vhdxDisk -DataPartition $additionalDataPartition | Out-Null
        }

        WriteLog 'All necessary partitions created.'

        Add-BootFiles -OsPartitionDriveLetter $osPartitionDriveLetter -SystemPartitionDriveLetter $createdSystemPartitionDriveLetter -AdkPath $adkPath -WindowsArch $WindowsArch
    
        #Add Windows packages
        if ($UpdateLatestCU -or $UpdateLatestNet -or $UpdatePreviewCU ) {
            try {
                Set-Progress -Percentage 25 -Message "Applying Windows Updates to VHDX..."
                WriteLog "Adding KBs to $WindowsPartition"
                WriteLog 'This can take 10+ minutes depending on how old the media is and the size of the KB. Please be patient'
                # If WindowsRelease is 2016, we need to add the SSU first
                if ($WindowsRelease -eq 2016 -and $installationType -eq "Server") {
                    WriteLog 'WindowsRelease is 2016, adding SSU first'
                    WriteLog "Adding SSU to $WindowsPartition"
                    # Add-WindowsPackage -Path $WindowsPartition -PackagePath $SSUFilePath | Out-Null
                    Invoke-Process cmd "/c ""$DandIEnv"" && dism /Image:$WindowsPartition /Add-Package /PackagePath:$SSUFilePath" | Out-Null
                    WriteLog "SSU added to $WindowsPartition"
                    # WriteLog "Removing $SSUFilePath"
                    # Remove-Item -Path $SSUFilePath -Force | Out-Null
                    # WriteLog 'SSU removed'
                }
                if ($WindowsRelease -in 2016, 2019, 2021 -and $isLTSC) {
                    WriteLog "WindowsRelease is $WindowsRelease and is $WindowsSKU, adding SSU first"
                    WriteLog "Adding SSU to $WindowsPartition"
                    # Add-WindowsPackage -Path $WindowsPartition -PackagePath $SSUFilePath | Out-Null
                    Invoke-Process cmd "/c ""$DandIEnv"" && dism /Image:$WindowsPartition /Add-Package /PackagePath:$SSUFilePath" | Out-Null
                    WriteLog "SSU added to $WindowsPartition"
                    # WriteLog "Removing $SSUFilePath"
                    # Remove-Item -Path $SSUFilePath -Force | Out-Null
                    # WriteLog 'SSU removed'
                }
                # Break out CU and NET updates to be added separately to abide by Checkpoint Update recommendations
                if ($UpdateLatestCU) {
                    if ($installLatestCuInVm) {
                        WriteLog "Skipping offline CU install for Windows 10 LTSB/LTSC. CU will be installed in VM from Apps ISO."
                    }
                    else {
                        WriteLog "Adding $CUPath to $WindowsPartition"
                        # Add-WindowsPackage -Path $WindowsPartition -PackagePath $CUPath | Out-Null
                        Invoke-Process cmd "/c ""$DandIEnv"" && dism /Image:$WindowsPartition /Add-Package /PackagePath:$CUPath" | Out-Null
                        WriteLog "$CUPath added to $WindowsPartition"
                    }
                }
                if ($UpdatePreviewCU) {
                    WriteLog "Adding $CUPPath to $WindowsPartition"
                    # Add-WindowsPackage -Path $WindowsPartition -PackagePath $CUPPath | Out-Null
                    Invoke-Process cmd "/c ""$DandIEnv"" && dism /Image:$WindowsPartition /Add-Package /PackagePath:$CUPPath" | Out-Null
                    WriteLog "$CUPPath added to $WindowsPartition"
                }
                if ($UpdateLatestNet) {
                    WriteLog "Adding $NETPath to $WindowsPartition"
                    # Add-WindowsPackage -Path $WindowsPartition -PackagePath $NETPath | Out-Null
                    Invoke-Process cmd "/c ""$DandIEnv"" && dism /Image:$WindowsPartition /Add-Package /PackagePath:$NETPath" | Out-Null
                    WriteLog "$NETPath added to $WindowsPartition"
                }
                if ($UpdateLatestMicrocode -and $WindowsRelease -in 2016, 2019) {
                    WriteLog "Adding $MicrocodePath to $WindowsPartition"
                    # Add-WindowsPackage -Path $WindowsPartition -PackagePath $MicrocodePath | Out-Null
                    Invoke-Process cmd "/c ""$DandIEnv"" && dism /Image:$WindowsPartition /Add-Package /PackagePath:$MicrocodePath" | Out-Null
                    WriteLog "$MicrocodePath added to $WindowsPartition"
                }
                WriteLog "KBs added to $WindowsPartition"
                if ($AllowVHDXCaching) {
                    $cachedVHDXInfo = [VhdxCacheItem]::new()
                    # Record only updates from this build (current required updates and any cached names carried forward)
                    $includedUpdateNames = [System.Collections.Generic.List[string]]::new()
                    foreach ($includedUpdate in $requiredUpdates) {
                        if (-not [string]::IsNullOrWhiteSpace($includedUpdate.Name)) {
                            $includedUpdateNames.Add($includedUpdate.Name)
                        }
                    }
                    foreach ($cachedName in $cachedIncludedUpdateNames) {
                        if (-not [string]::IsNullOrWhiteSpace($cachedName)) {
                            $includedUpdateNames.Add($cachedName)
                        }
                    }
                    foreach ($includedName in ($includedUpdateNames | Sort-Object -Unique)) {
                        if (-not ($cachedVHDXInfo.IncludedUpdates | Where-Object { $_.Name -eq $includedName })) {
                            $cachedVHDXInfo.IncludedUpdates += ([VhdxCacheUpdateItem]::new($includedName))
                        }
                    }
                }
                WriteLog 'Clean Up the WinSxS Folder'
                WriteLog 'This can take 10+ minutes depending on how old the media is and the size of the KB. Please be patient'
                Dism /Image:$WindowsPartition /Cleanup-Image /StartComponentCleanup /ResetBase | Out-Null
                WriteLog 'Clean Up the WinSxS Folder completed'
            }
            catch {
                Write-Host "Adding KB to VHDX failed with error $_"
                WriteLog "Adding KB to VHDX failed with error $_"
                if ($_.Exception.HResult -eq -2146498525) {
                    Write-Host 'Missing latest Servicing Stack Update'
                    Write-Host 'Media likely older than 2023-09 for Windows Server 2022 (KB5030216), or 2021-08 for Windows Server 2019 (KB5005112)'
                    Write-Host 'Recommended to use the latest media'
                    WriteLog 'Missing latest Servicing Stack Update'
                    WriteLog 'Media likely older than 2023-09 for Windows Server 2022 (KB5030216), or 2021-08 for Windows Server 2019 (KB5005112)'
                    WriteLog 'Recommended to use the latest media'
                }
                throw $_
            }  
        }

        #Enable Windows Optional Features (e.g. .Net3, etc)
        If ($OptionalFeatures) {
            $Source = Join-Path (Split-Path $wimpath) 'sxs'
            Enable-WindowsFeaturesByName -FeatureNames $OptionalFeatures -Source $Source
        }
		If ($windowsIsoMountedByBuild) {
			Dismount-WindowsIsoSource -ImagePath $ISOPath
			$windowsIsoMountedByBuild = $false
        }
        # If $wimPath is an ESD file, remove it only when configured
        If ($wimPath -match '\.esd$') {
            if ($RemoveDownloadedESD -and (Test-Path -Path $wimPath)) {
                WriteLog "Deleting $wimPath file"
                Remove-Item -Path $wimPath -Force
                WriteLog "$wimPath deleted"
            }
            else {
                WriteLog "Keeping downloaded ESD file $wimPath"
            }
        }
    
    }
    else {
        #Use cached vhdx file
        WriteLog 'Using cached VHDX file to speed up build process'
        WriteLog "VHDX file is: $($cachedVHDXInfo.VhdxFileName)"

        Robocopy.exe $($VHDXCacheFolder) $($VMPath) $($cachedVHDXInfo.VhdxFileName) /E /COPY:DAT /R:5 /W:5 /J | Out-Null
        $VHDXPath = Join-Path $($VMPath) $($cachedVHDXInfo.VhdxFileName)

        $vhdxDisk = Get-VHD -Path $VHDXPath | Mount-VHD -Passthru | Get-Disk
        $partitionLayout = Resolve-VhdxPartitionLayout -Disk $vhdxDisk -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions
        $partitionLayout = Set-VhdxBuildPartitionDriveLetters -Layout $partitionLayout -SystemPartitionDriveLetter $SystemPartitionDriveLetter -WindowsPartitionDriveLetter $WindowsPartitionDriveLetter -RecoveryPartitionDriveLetter $RecoveryPartitionDriveLetter -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions
        $osPartition = $partitionLayout.WindowsPartition
        $osPartitionDriveLetter = $osPartition.DriveLetter
        $WindowsPartition = $osPartitionDriveLetter + ':\'

    }

    #Set Product key
    If ($ProductKey) {
        WriteLog "Setting Windows Product Key"
        Set-WindowsProductKey -Path $WindowsPartition -ProductKey $ProductKey
    }

    Set-Progress -Percentage 40 -Message "Finalizing VHDX..."
    if ($AllowVHDXCaching -and !$cachedVHDXFileFound) {
        WriteLog 'Caching VHDX file'

        # Run full VHDX optimization after servicing/cleanup and before cache copy
        WriteLog 'Optimizing VHDX before copying to cache dir'
        Optimize-FFUCaptureDrive -VhdxPath $VHDXPath -EnableVolumeRetrim $true -SystemPartitionDriveLetter $SystemPartitionDriveLetter -WindowsPartitionDriveLetter $WindowsPartitionDriveLetter -RecoveryPartitionDriveLetter $RecoveryPartitionDriveLetter -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions

        WriteLog 'Copying to cache dir'

        #Assuming there are now name collisons
        Robocopy.exe $($VMPath) $($VHDXCacheFolder) $("$VMName.vhdx") /E /COPY:DAT /R:5 /W:5 /J | Out-Null

        #Only create new instance if not created during patching
        if ($null -eq $cachedVHDXInfo) {
            $cachedVHDXInfo = [VhdxCacheItem]::new()
        }
        if ($AllowVHDXCaching -and $cachedIncludedUpdateNames.Count -gt 0) {
            foreach ($cachedName in $cachedIncludedUpdateNames) {
                if (-not ($cachedVHDXInfo.IncludedUpdates | Where-Object { $_.Name -eq $cachedName })) {
                    $cachedVHDXInfo.IncludedUpdates += ([VhdxCacheUpdateItem]::new($cachedName))
                }
            }
        }
        $cachedVHDXInfo.VhdxFileName = $("$VMName.vhdx")
        $cachedVHDXInfo.LogicalSectorSizeBytes = $LogicalSectorSizeBytes
        $cachedVHDXInfo.Disksize = $Disksize
        $cachedVHDXInfo.CreateRecoveryPartition = $CreateRecoveryPartition
        $cachedVHDXInfo.PartitionLayoutSignatureVersion = $partitionLayoutSignatureVersion
        $cachedVHDXInfo.PartitionLayoutSignature = $partitionLayoutSignature
        $cachedVHDXInfo.WindowsSKU = $WindowsSKU
        $cachedVHDXInfo.WindowsRelease = $WindowsRelease
        $cachedVHDXInfo.WindowsVersion = $WindowsVersion
        $cachedVHDXInfo.OptionalFeatures = $OptionalFeatures
        
        $cachedVHDXInfo | ConvertTo-Json | Out-File -FilePath ("{0}\{1}_config.json" -f $($VHDXCacheFolder), $VMName)
        WriteLog "Cached VHDX file $("$VMName.vhdx")"

        #Remount the VHDX file if $installapps is false so the VHDX can be captured to an FFU
        If (-not $InstallApps) {
            Mount-Vhd -Path $VHDXPath
        }
    } 
}
catch {
    Write-Host 'Creating VHDX Failed'
    WriteLog "Creating VHDX Failed with error $_"
    WriteLog "Dismounting $VHDXPath"
    Dismount-ScratchVhdx -VhdxPath $VHDXPath
    WriteLog "Removing $VMPath"
    Remove-Item -Path $VMPath -Force -Recurse | Out-Null
    WriteLog 'Removal complete'
	If ($windowsIsoMountedByBuild) {
		Dismount-WindowsIsoSource -ImagePath $ISOPath
		$windowsIsoMountedByBuild = $false
    }
    if (-not $ISOPath) {
        # Remove ESD file only when configured
        if ($RemoveDownloadedESD -and -not [string]::IsNullOrWhiteSpace($wimPath) -and ($wimPath -match '\.esd$') -and (Test-Path -Path $wimPath)) {
            WriteLog "Deleting ESD file $wimPath"
            Remove-Item -Path $wimPath -Force
            WriteLog "ESD File deleted"
        }
        else {
            WriteLog "Keeping downloaded ESD file $wimPath"
        }
    }
    throw $_
    
}

# Prepare Windows 10 LTSB/LTSC CU assets for in-VM install when required
if ($InstallApps -and $installLatestCuInVm) {
    try {
        # Ensure CUPath is resolved and the CU package is available locally
        if (([string]::IsNullOrWhiteSpace($CUPath)) -or -not (Test-Path -Path $CUPath)) {
            $ltscKbCacheRootPath = Join-Path -Path (Join-Path -Path $KBPath -ChildPath 'Windows10') -ChildPath "LTSC$WindowsRelease"
            if (-not (Test-Path -Path $ltscKbCacheRootPath)) {
                WriteLog "Creating LTSC KB cache folder $ltscKbCacheRootPath"
                New-Item -Path $ltscKbCacheRootPath -ItemType Directory -Force | Out-Null
            }

            foreach ($cuUpdateInfo in $cuUpdateInfos) {
                $expectedCuName = $cuUpdateInfo.Name
                if ([string]::IsNullOrWhiteSpace($expectedCuName)) {
                    continue
                }

                $expectedCuPath = Join-Path -Path $ltscKbCacheRootPath -ChildPath $expectedCuName
                if (-not (Test-Path -Path $expectedCuPath)) {
                    WriteLog "Downloading $expectedCuName to $ltscKbCacheRootPath"
                    Start-BitsTransferWithRetry -Source $cuUpdateInfo.Url -Destination $ltscKbCacheRootPath
                }
            }

            if (-not [string]::IsNullOrWhiteSpace($cuKbArticleId)) {
                $resolvedCu = Get-ChildItem -Path $ltscKbCacheRootPath -Filter "*$cuKbArticleId*" -Recurse -File -ErrorAction SilentlyContinue | Select-Object -First 1
                if ($null -ne $resolvedCu) {
                    $CUPath = $resolvedCu.FullName
                }
            }

            if (([string]::IsNullOrWhiteSpace($CUPath)) -and $cuUpdateInfos.Count -gt 0) {
                $fallbackCuName = $cuUpdateInfos[0].Name
                if (-not [string]::IsNullOrWhiteSpace($fallbackCuName)) {
                    $resolvedCu = Get-ChildItem -Path $ltscKbCacheRootPath -Filter $fallbackCuName -Recurse -File -ErrorAction SilentlyContinue | Select-Object -First 1
                    if ($null -ne $resolvedCu) {
                        $CUPath = $resolvedCu.FullName
                    }
                }
            }

            if (([string]::IsNullOrWhiteSpace($CUPath)) -or -not (Test-Path -Path $CUPath)) {
                throw "Unable to resolve LTSC CU package path for in-VM install."
            }
        }

        # Stage CU payload into Apps content so it is included in Apps ISO
        if (-not (Test-Path -Path $LtscCUStagePath)) {
            WriteLog "Creating LTSC CU staging folder $LtscCUStagePath"
            New-Item -Path $LtscCUStagePath -ItemType Directory -Force | Out-Null
        }

        $ltscCuFileName = Split-Path -Path $CUPath -Leaf
        $stagedLtscCuPath = Join-Path -Path $LtscCUStagePath -ChildPath $ltscCuFileName
        Copy-Item -Path $CUPath -Destination $stagedLtscCuPath -Force | Out-Null
        WriteLog "Staged LTSC CU package to $stagedLtscCuPath"

        # Create Install-LTSCUpdate.ps1 for in-VM execution via orchestrator
        $installLtscUpdateCommand = @"
# Validate LTSC CU package exists on Apps ISO mount
`$kbPath = Join-Path -Path `$env:FFUAppsRoot -ChildPath "LTSCUpdate\$ltscCuFileName"

# Extract KB ID from filename for idempotent checks
`$kbFileName = Split-Path -Path `$kbPath -Leaf
`$kbMatch = [regex]::Match(`$kbFileName, 'KB\d+', [System.Text.RegularExpressions.RegexOptions]::IgnoreCase)
`$kbId = if (`$kbMatch.Success) { `$kbMatch.Value.ToUpperInvariant() } else { `$null }

# Detect whether Windows has a pending reboot
function Test-PendingReboot {
    if (Test-Path -Path 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\Component Based Servicing\RebootPending') { return `$true }
    if (Test-Path -Path 'HKLM:\SOFTWARE\Microsoft\Windows\CurrentVersion\WindowsUpdate\Auto Update\RebootRequired') { return `$true }

    `$sessionManagerPath = 'HKLM:\SYSTEM\CurrentControlSet\Control\Session Manager'
    try {
        `$pendingRename = Get-ItemProperty -Path `$sessionManagerPath -Name 'PendingFileRenameOperations' -ErrorAction SilentlyContinue
        if (`$null -ne `$pendingRename) { return `$true }
    }
    catch {
    }

    return `$false
}

# Skip if the target KB is already installed
if (`$kbId -and (Get-HotFix -Id `$kbId -ErrorAction SilentlyContinue)) {
    Write-Host "`$kbId is already installed. Skipping LTSC CU install."
    if (Test-PendingReboot) {
        Write-Host "Pending reboot detected. Restarting before continuing orchestration..."
        Restart-Computer -Force
    }
    return
}

# Stop with a non-zero code if the package is missing
if (-not (Test-Path -Path `$kbPath)) {
    Write-Host "LTSC CU package not found at `$kbPath"
    exit 1
}

# Start silent LTSC CU installation and capture exit code
Write-Host "Starting LTSC CU install from `$kbPath"
`$wusaArguments = @(`$kbPath, '/quiet', '/norestart')
`$wusaProcess = Start-Process -FilePath "wusa.exe" -ArgumentList `$wusaArguments -PassThru -Wait -ErrorAction Stop
`$wusaExitCode = `$wusaProcess.ExitCode
Write-Host "WUSA exit code: `$wusaExitCode"

# 3010 means reboot required; restart now so sysprep won't fail later
if (`$wusaExitCode -eq 3010) {
    Write-Host "LTSC CU requires reboot. Restarting now..."
    Restart-Computer -Force
    return
}

# 0 = success, 2359302 = not applicable/already present
if (`$wusaExitCode -eq 0 -or `$wusaExitCode -eq 2359302) {
    if (Test-PendingReboot) {
        Write-Host "Pending reboot detected after LTSC CU step. Restarting now..."
        Restart-Computer -Force
        return
    }

    Write-Host "LTSC CU install completed."
    return
}

throw "LTSC CU install failed with WUSA exit code `$wusaExitCode."
"@

        # Back up any pre-existing script with the same name before overwrite.
        Backup-RunFile -FFUDevelopmentPath $FFUDevelopmentPath -Path $InstallLTSCUpdatePath
        Set-Content -Path $InstallLTSCUpdatePath -Value $installLtscUpdateCommand -Force
        if (-not (Test-Path -Path $InstallLTSCUpdatePath)) {
            throw "Failed to create $InstallLTSCUpdatePath"
        }

        WriteLog "$InstallLTSCUpdatePath created successfully"
        $refreshAppsIsoForLtscCu = $true
        WriteLog "Prepared Windows 10 LTSB/LTSC CU assets for in-VM installation"
    }
    catch {
        WriteLog "Preparing Windows 10 LTSB/LTSC CU assets for in-VM installation failed with error $_"
        throw $_
    }
}

#Inject unattend after caching so cached VHDX never contains audit-mode unattend
if ($InstallApps) {
    # Determine mount state and only mount if needed to avoid redundant mount/dismount cycles
    $vhdMeta = Get-VHD -Path $VHDXPath
    if ($vhdMeta.Attached) {
        WriteLog 'VHDX already mounted; reusing existing mount for unattend injection'
        $disk = Get-Disk -Number $vhdMeta.DiskNumber
    }
    else {
        WriteLog 'Mounting VHDX to inject unattend for audit-mode boot'
        $disk = Mount-VHD -Path $VHDXPath -Passthru | Get-Disk
    }
    $partitionLayout = Resolve-VhdxPartitionLayout -Disk $disk -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions
    $partitionLayout = Set-VhdxBuildPartitionDriveLetters -Layout $partitionLayout -SystemPartitionDriveLetter $SystemPartitionDriveLetter -WindowsPartitionDriveLetter $WindowsPartitionDriveLetter -RecoveryPartitionDriveLetter $RecoveryPartitionDriveLetter -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions
    $osPartition = $partitionLayout.WindowsPartition
    $osPartitionDriveLetter = $osPartition.DriveLetter
    WriteLog 'Copying unattend file to boot to audit mode'
    New-Item -Path "$($osPartitionDriveLetter):\Windows\Panther\Unattend" -ItemType Directory -Force | Out-Null
    $orchestrationBootstrapSourcePath = Join-Path -Path $FFUDevelopmentPath -ChildPath 'BuildFFUUnattend\Start-FFUOrchestration.ps1'
    if (-not (Test-Path -Path $orchestrationBootstrapSourcePath -PathType Leaf)) {
        throw "Orchestration bootstrap script not found at $orchestrationBootstrapSourcePath"
    }
    $orchestrationBootstrapTargetFolder = "$($osPartitionDriveLetter):\Windows\Setup\Scripts"
    New-Item -Path $orchestrationBootstrapTargetFolder -ItemType Directory -Force | Out-Null
    Copy-Item -Path $orchestrationBootstrapSourcePath -Destination (Join-Path -Path $orchestrationBootstrapTargetFolder -ChildPath 'Start-FFUOrchestration.ps1') -Force | Out-Null
    if ($enforceBuildVmDataPartitionDriveLetters) {
        $windowsPartitionRoot = "$($osPartitionDriveLetter):\"
        $null = Add-FFUDataPartitionDriveLetterArtifacts -WindowsPartitionRoot $windowsPartitionRoot -Layout $partitionLayout -AdditionalDataPartitions $normalizedAdditionalDataPartitions -WindowsArch $WindowsArch -ManifestPurpose BuildVm -Optimize $Optimize -OptimizeFFUPartitionNumber $OptimizeFFUPartitionNumber -FFUDevelopmentPath $FFUDevelopmentPath
    }
    if ($WindowsArch -eq 'x64') {
        Copy-Item -Path "$FFUDevelopmentPath\BuildFFUUnattend\unattend_x64.xml" -Destination "$($osPartitionDriveLetter):\Windows\Panther\Unattend\Unattend.xml" -Force | Out-Null
    }
    else {
        Copy-Item -Path "$FFUDevelopmentPath\BuildFFUUnattend\unattend_arm64.xml" -Destination "$($osPartitionDriveLetter):\Windows\Panther\Unattend\Unattend.xml" -Force | Out-Null
    }
    WriteLog 'Copy completed'
    # Always dismount so downstream VM creation logic has a clean starting point
    Dismount-ScratchVhdx -VhdxPath $VHDXPath
}
elseif ($configureDeployedDataPartitionDriveLetters) {
    $vhdMeta = Get-VHD -Path $VHDXPath
    if ($vhdMeta.Attached) {
        WriteLog 'VHDX already mounted; reusing existing mount for data partition drive-letter staging'
        $disk = Get-Disk -Number $vhdMeta.DiskNumber
    }
    else {
        WriteLog 'Mounting VHDX to stage data partition drive-letter assignment'
        $disk = Mount-VHD -Path $VHDXPath -Passthru | Get-Disk
    }
    $partitionLayout = Resolve-VhdxPartitionLayout -Disk $disk -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions
    $partitionLayout = Set-VhdxBuildPartitionDriveLetters -Layout $partitionLayout -SystemPartitionDriveLetter $SystemPartitionDriveLetter -WindowsPartitionDriveLetter $WindowsPartitionDriveLetter -RecoveryPartitionDriveLetter $RecoveryPartitionDriveLetter -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions
    $osPartitionDriveLetter = $partitionLayout.WindowsPartition.DriveLetter
    $windowsPartitionRoot = "$($osPartitionDriveLetter):\"
    $null = Add-FFUDataPartitionDriveLetterArtifacts -WindowsPartitionRoot $windowsPartitionRoot -Layout $partitionLayout -AdditionalDataPartitions $normalizedAdditionalDataPartitions -WindowsArch $WindowsArch -Optimize $Optimize -OptimizeFFUPartitionNumber $OptimizeFFUPartitionNumber -FFUDevelopmentPath $FFUDevelopmentPath
    $deploymentUnattendPath = Join-Path -Path $windowsPartitionRoot -ChildPath 'Windows\Panther\Unattend.xml'
    Add-FFUDataPartitionDriveLetterCommandToUnattend -UnattendPath $deploymentUnattendPath -ProcessorArchitecture $WindowsArch
}

#If installing apps (Office or 3rd party), we need to build a VM and capture that FFU, if not, just cut the FFU from the VHDX file
if ($InstallApps) {
    Set-Progress -Percentage 41 -Message "Starting VM for app installation..."

    # Refresh Apps ISO to include LTSC CU assets when staged for in-VM install
    if ($refreshAppsIsoForLtscCu) {
        try {
            WriteLog "Refreshing Apps ISO to include LTSC CU assets"
            if (Test-Path -Path $AppsISO) {
                WriteLog "Removing $AppsISO"
                Remove-Item -Path $AppsISO -Force -ErrorAction SilentlyContinue
                WriteLog 'Removal complete'
            }
            Sync-UserAppListForOrchestration -SourcePath $UserAppListPath -AppsPath $AppsPath -OrchestrationPath $OrchestrationPath -AppInstallConfigPath $appInstallConfigPath
            New-AppsISO
            WriteLog "Apps ISO refreshed with LTSC CU assets"
        }
        catch {
            WriteLog "Refreshing Apps ISO for LTSC CU assets failed with error $_"
            throw $_
        }
    }

    #Create VM and attach VHDX
    try {
        WriteLog 'Creating new FFU VM'
        $FFUVM = New-FFUVM
        WriteLog 'FFU VM Created'
    }
    catch {
        Write-Host 'VM creation failed'
        Writelog "VM creation failed with error $_"
        Remove-FFUVM -VMName $VMName
        throw $_
        
    }
}
#Capture FFU file
try {
    #Check for FFU Folder and create it if it's missing
    If (-not (Test-Path -Path $FFUCaptureLocation)) {
        WriteLog "Creating FFU capture location at $FFUCaptureLocation"
        New-Item -Path $FFUCaptureLocation -ItemType Directory -Force
        WriteLog "Successfully created FFU capture location at $FFUCaptureLocation"
    }
    #Shorten Windows SKU for use in FFU file name to remove spaces and long names
    WriteLog "Shortening Windows SKU: $WindowsSKU for FFU file name"
    $shortenedWindowsSKU = Get-ShortenedWindowsSKU -WindowsSKU $WindowsSKU
    WriteLog "Shortened Windows SKU: $shortenedWindowsSKU"
    #Check if VM is done provisioning
    If ($InstallApps) {
        Set-Progress -Percentage 50 -Message "Installing applications in VM; please wait for VM to shut down..."
        do {
            $FFUVM = Get-VM -Name $FFUVM.Name
            Start-Sleep -Seconds 10
            WriteLog 'Waiting for VM to shutdown'
        } while ($FFUVM.State -ne 'Off')
        WriteLog 'VM Shutdown'
        if ($enforceBuildVmDataPartitionDriveLetters) {
            $vhdMeta = Get-VHD -Path $VHDXPath
            if ($vhdMeta.Attached) {
                WriteLog 'VHDX already mounted; reusing existing mount to validate data partition drive-letter audit results'
                $disk = Get-Disk -Number $vhdMeta.DiskNumber
            }
            else {
                WriteLog 'Mounting VHDX to validate data partition drive-letter audit results'
                $disk = Mount-VHD -Path $VHDXPath -Passthru | Get-Disk
            }
            $partitionLayout = Resolve-VhdxPartitionLayout -Disk $disk -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions
            $partitionLayout = Set-VhdxBuildPartitionDriveLetters -Layout $partitionLayout -SystemPartitionDriveLetter $SystemPartitionDriveLetter -WindowsPartitionDriveLetter $WindowsPartitionDriveLetter -RecoveryPartitionDriveLetter $RecoveryPartitionDriveLetter -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions
            $osPartitionDriveLetter = $partitionLayout.WindowsPartition.DriveLetter
            $windowsPartitionRoot = "$($osPartitionDriveLetter):\"
            Test-FFUDataPartitionDriveLetterAuditResult -WindowsPartitionRoot $windowsPartitionRoot

            $runtimeDirectory = Join-Path -Path $windowsPartitionRoot -ChildPath 'Windows\Setup\Scripts\FFUDL'
            $runtimeScriptPath = Join-Path -Path $runtimeDirectory -ChildPath 'Apply.ps1'
            $manifestPath = Join-Path -Path $runtimeDirectory -ChildPath 'Manifest.json'
            $null = Add-FFUDataPartitionDriveLetterArtifacts -WindowsPartitionRoot $windowsPartitionRoot -Layout $partitionLayout -AdditionalDataPartitions $normalizedAdditionalDataPartitions -WindowsArch $WindowsArch -Optimize $Optimize -OptimizeFFUPartitionNumber $OptimizeFFUPartitionNumber -FFUDevelopmentPath $FFUDevelopmentPath
            if (-not (Test-Path -LiteralPath $runtimeScriptPath -PathType Leaf) -or -not (Test-Path -LiteralPath $manifestPath -PathType Leaf)) {
                throw "Data partition drive-letter runtime artifacts are missing from $runtimeDirectory after audit validation."
            }

            $deploymentUnattendPath = Join-Path -Path $windowsPartitionRoot -ChildPath 'Windows\Panther\Unattend.xml'
            if (-not (Test-Path -LiteralPath $deploymentUnattendPath -PathType Leaf) -and $InjectUnattend) {
                $deploymentUnattendDirectory = Split-Path -Path $deploymentUnattendPath -Parent
                New-Item -Path $deploymentUnattendDirectory -ItemType Directory -Force | Out-Null
                $stagedInjectedUnattendPath = Join-Path -Path $AppsPath -ChildPath 'Unattend\Unattend.xml'
                if (Test-Path -LiteralPath $stagedInjectedUnattendPath -PathType Leaf) {
                    Copy-Item -LiteralPath $stagedInjectedUnattendPath -Destination $deploymentUnattendPath -Force
                    WriteLog "Restored staged deployment unattend to $deploymentUnattendPath before adding data partition drive-letter assignment."
                }
                else {
                    $unattendSource = Get-UnattendSourcePath -UnattendFolder $UnattendFolder -WindowsArch $WindowsArch -UnattendX64FilePath $UnattendX64FilePath -UnattendArm64FilePath $UnattendArm64FilePath
                    Save-StagedUnattendFile -SourcePath $unattendSource -DestinationPath $deploymentUnattendPath -DeviceNamingMode $DeviceNamingMode -DeviceNameTemplate $normalizedDeviceNameTemplate -WindowsArch $WindowsArch
                    WriteLog "Restaged deployment unattend to $deploymentUnattendPath before adding data partition drive-letter assignment."
                }
            }
            Add-FFUDataPartitionDriveLetterCommandToUnattend -UnattendPath $deploymentUnattendPath -ProcessorArchitecture $WindowsArch
        }
        Set-Progress -Percentage 65 -Message "Optimizing VHDX before capture..."
        Optimize-FFUCaptureDrive -VhdxPath $VHDXPath -SystemPartitionDriveLetter $SystemPartitionDriveLetter -WindowsPartitionDriveLetter $WindowsPartitionDriveLetter -RecoveryPartitionDriveLetter $RecoveryPartitionDriveLetter -CreateRecoveryPartition $CreateRecoveryPartition -AdditionalDataPartitions $normalizedAdditionalDataPartitions
        #Capture FFU file
        New-FFU
    }
    else {
        #Create FFU file
        New-FFU
    }    
}
Catch {
    Write-Host 'Capturing FFU file failed'
    Writelog "Capturing FFU file failed with error $_"
    If ($InstallApps) {
        Remove-FFUVM -VMName $VMName
    }
    else {
        Remove-FFUVM
    }
    
    throw $_
    
}
#Clean up VM or VHDX
try {
    Remove-FFUVM
    WriteLog 'FFU build complete!'
}
catch {
    Write-Host 'VM or vhdx cleanup failed'
    Writelog "VM or vhdx cleanup failed with error $_"
    throw $_
}


#Create Deployment Media
If ($CreateDeploymentMedia) {
    Set-Progress -Percentage 91 -Message "Creating deployment media..."
    try {
        New-PEMedia
    }
    catch {
        Write-Host 'Creating deployment media failed'
        WriteLog "Creating deployment media failed with error $_"
        throw $_
    
    }
}
If ($BuildUSBDrive) {
    Set-Progress -Percentage 95 -Message "Building USB drive..."
    try {
        If (Test-Path -Path $DeployISO) {
            $ffuFilesToCopy = @()

            # Always include the FFU that was just built (fallback to most recent .ffu in capture folder)
            $currentFFU = $null
            if ($null -ne $FFUFile -and -not [string]::IsNullOrWhiteSpace($FFUFile) -and (Test-Path -LiteralPath $FFUFile)) {
                $currentFFU = $FFUFile
            }
            else {
                try {
                    $ffuDir = if (-not [string]::IsNullOrWhiteSpace($FFUCaptureLocation)) { $FFUCaptureLocation } else { Join-Path $FFUDevelopmentPath 'FFU' }
                    if (Test-Path -LiteralPath $ffuDir) {
                        $latest = Get-ChildItem -Path $ffuDir -Filter '*.ffu' -File | Sort-Object LastWriteTime -Descending | Select-Object -First 1
                        if ($null -ne $latest) { $currentFFU = $latest.FullName }
                    }
                }
                catch {
                    WriteLog "Failed to resolve latest FFU file to copy: $($_.Exception.Message)"
                }
            }
            if ($null -ne $currentFFU) {
                $ffuFilesToCopy += $currentFFU
            }

            if ($CopyAdditionalFFUFiles -and ($null -ne $AdditionalFFUFiles) -and ($AdditionalFFUFiles.Count -gt 0)) {
                $ffuFilesToCopy += $AdditionalFFUFiles
            }

            $ffuFilesToCopy = $ffuFilesToCopy | Where-Object { -not [string]::IsNullOrWhiteSpace($_) } | Select-Object -Unique
            New-DeploymentUSB -CopyFFU -FFUFilesToCopy $ffuFilesToCopy
        }
        else {
            WriteLog "$BuildUSBDrive set to true, however unable to find $DeployISO. USB drive not built."
        }
        
    }
    catch {
        Write-Host 'Building USB deployment drive failed'
        Writelog "Building USB deployment drive failed with error $_"
        throw $_
    }
}

Set-Progress -Percentage 99 -Message "Finalizing and cleaning up..."
# Delegated post-build cleanup to common module
Invoke-FFUPostBuildCleanup -RootPath $FFUDevelopmentPath -AppsPath $AppsPath -DriversPath $DriversFolder -FFUCapturePath $FFUCaptureLocation -DeployISOPath $DeployISO -AppsISOPath $AppsISO -RemoveDeployISO:$CleanupDeployISO -RemoveAppsISO:$CleanupAppsISO -RemoveDrivers:$CleanupDrivers -RemoveFFU:$RemoveFFU -RemoveApps:$RemoveApps -RemoveUpdates:$RemoveUpdates -RemoveDownloadedESD:$RemoveDownloadedESD -KBPath:$KBPath


# Remove WinGetWin32Apps.json so it is always rebuilt next run
if (Test-Path -Path $wingetWin32jsonFile -PathType Leaf) {
    WriteLog "Removing $wingetWin32jsonFile"
    Remove-Item -Path $wingetWin32jsonFile -Force -ErrorAction SilentlyContinue
    WriteLog "Removal complete"
}

# Remove AppInstallConfig.json so it is always rebuilt next run
if (Test-Path -Path $appInstallConfigPath -PathType Leaf) {
    WriteLog "Removing $appInstallConfigPath"
    Remove-Item -Path $appInstallConfigPath -Force -ErrorAction SilentlyContinue
    WriteLog "Removal complete"
}
#Set $LongPathsEnabled registry value back to original value. $LongPathsEnabled could be $null if the registry value was not found
if ($null -eq $LongPathsEnabled) {
    Remove-ItemProperty -Path 'HKLM:\SYSTEM\CurrentControlSet\Control\FileSystem' -Name 'LongPathsEnabled' -ErrorAction SilentlyContinue
}
else {
    Set-ItemProperty -Path 'HKLM:\SYSTEM\CurrentControlSet\Control\FileSystem' -Name 'LongPathsEnabled' -Value $LongPathsEnabled
}

#Clean up dirty.txt file
$dirtyFilePath = Join-Path -Path $FFUDevelopmentPath -ChildPath 'dirty.txt'
Remove-Item -Path $dirtyFilePath -Force | out-null
# Remove per-run session folder if present
$sessionDir = Join-Path $FFUDevelopmentPath '.session'
if (Test-Path -Path $sessionDir) { 
    Remove-Item -Path $sessionDir -Recurse -Force -ErrorAction SilentlyContinue 
}
if ($VerbosePreference -ne 'Continue') {
    Write-Host 'Script complete'
}
Set-Progress -Percentage 100 -Message "Build process complete."
# Record the end time
$endTime = Get-Date
Write-Host "FFU build process completed at" $endTime

# Calculate the total run time
$runTime = $endTime - $startTime

# Format the runtime with hours, minutes, and seconds
if ($runTime.TotalHours -ge 1) {
    $runTimeFormatted = 'Duration: {0:hh} hr {0:mm} min {0:ss} sec' -f $runTime
}
else {
    $runTimeFormatted = 'Duration: {0:mm} min {0:ss} sec' -f $runTime
}

if ($VerbosePreference -ne 'Continue') {
    Write-Host $runTimeFormatted
}
WriteLog 'Script complete'
WriteLog $runTimeFormatted
