<#
.SYNOPSIS
    Contains functions for loading and saving UI configuration.
.DESCRIPTION
    This module provides the core logic for loading and saving the UI configuration. It includes functions to gather settings from the various UI controls, save them to a JSON file, and load settings from a JSON file to populate the UI. This allows users to persist their build configurations and easily switch between different setups.
#>
function Get-UIConfig {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )
    Sync-DiskLayoutRowsToControls -State $State

    # Create hash to store configuration
    $config = [ordered]@{
        AllowExternalHardDiskMedia     = $State.Controls.chkAllowExternalHardDiskMedia.IsChecked
        AllowVHDXCaching               = $State.Controls.chkAllowVHDXCaching.IsChecked
        AppListPath                    = $State.Controls.txtAppListJsonPath.Text
        AppsPath                       = $State.Controls.txtApplicationPath.Text
        AppsScriptVariables            = if ($State.Controls.chkDefineAppsScriptVariables.IsChecked) {
            $vars = @{}
            foreach ($item in $State.Data.appsScriptVariablesDataList) {
                $vars[$item.Key] = $item.Value
            }
            if ($vars.Count -gt 0) { $vars } else { $null }
        }
        else { $null }
        BuildUSBDrive                  = $State.Controls.chkBuildUSBDriveEnable.IsChecked
        CleanupAppsISO                 = $State.Controls.chkCleanupAppsISO.IsChecked
        CleanupDeployISO               = $State.Controls.chkCleanupDeployISO.IsChecked
        CleanupDrivers                 = $State.Controls.chkCleanupDrivers.IsChecked
        CompactOS                      = $State.Controls.chkCompactOS.IsChecked
        CompressDownloadedDriversToWim = $State.Controls.chkCompressDriversToWIM.IsChecked
        CopyAutopilot                  = $State.Controls.chkCopyAutopilot.IsChecked
        CopyDrivers                    = $State.Controls.chkCopyDrivers.IsChecked
        CopyOfficeConfigXML            = $State.Controls.chkCopyOfficeConfigXML.IsChecked
        CopyPEDrivers                  = $State.Controls.chkCopyPEDrivers.IsChecked
        UseDriversAsPEDrivers          = $State.Controls.chkUseDriversAsPEDrivers.IsChecked
        CopyPPKG                       = $State.Controls.chkCopyPPKG.IsChecked
        CopyUnattend                   = $State.Controls.chkCopyUnattend.IsChecked
        DeviceNamingMode               = Get-ConfiguredDeviceNamingMode -State $State
        DeviceNameTemplate             = $State.Controls.txtDeviceNameTemplate.Text
        DeviceNamePrefixesPath         = $State.Controls.txtDeviceNamePrefixesPath.Text
        DeviceNamePrefixes             = @(Get-DeviceNamePrefixes -State $State)
        DeviceNameSerialComputerNamesPath = $State.Controls.txtDeviceNameSerialComputerNamesPath.Text
        DeviceNameSerialComputerNames  = @(Get-SerialComputerNamesLines -State $State)
        CopyAdditionalFFUFiles         = $State.Controls.chkCopyAdditionalFFUFiles.IsChecked
        CreateDeploymentMedia          = $State.Controls.chkCreateDeploymentMedia.IsChecked
        InjectUnattend                 = $State.Controls.chkInjectUnattend.IsChecked
        UnattendX64FilePath            = $State.Controls.txtUnattendX64FilePath.Text
        UnattendArm64FilePath          = $State.Controls.txtUnattendArm64FilePath.Text
        CustomFFUNameTemplate          = $State.Controls.txtCustomFFUNameTemplate.Text
        Disksize                       = [int64]$State.Controls.txtDiskSize.Text * 1GB
        OSPartitionSize                = ConvertTo-PartitionSizeBytesFromGBText -Text $State.Controls.txtOSPartitionSizeGB.Text -FieldName 'Windows Partition Size'
        RecoveryPartitionSize          = ConvertTo-PartitionSizeBytesFromGBText -Text $State.Controls.txtRecoveryPartitionSizeGB.Text -FieldName 'Recovery Partition Size'
        CreateRecoveryPartition        = [bool]$State.Data.createRecoveryPartition
        AdditionalDataPartitions       = @(Get-AdditionalDataPartitionConfigRows -State $State)
        DownloadDrivers                = $State.Controls.chkDownloadDrivers.IsChecked
        DriversFolder                  = $State.Controls.txtDriversFolder.Text
        DriversJsonPath                = $State.Controls.txtDriversJsonPath.Text
        EnableVMNetworking             = $State.Controls.chkEnableVMNetworking.IsChecked
        FFUCaptureLocation             = $State.Controls.txtFFUCaptureLocation.Text
        FFUDevelopmentPath             = $State.Controls.txtFFUDevPath.Text
        FFUPrefix                      = $State.Controls.txtVMNamePrefix.Text
        InstallApps                    = $State.Controls.chkInstallApps.IsChecked
        InstallDrivers                 = $State.Controls.chkInstallDrivers.IsChecked
        InstallOffice                  = $State.Controls.chkInstallOffice.IsChecked
        InstallWingetApps              = $State.Controls.chkInstallWingetApps.IsChecked
        ISOPath                        = $State.Controls.txtISOPath.Text
        WindowsMediaSource             = if ($null -ne $State.Controls.rbProvideISO -and $State.Controls.rbProvideISO.IsChecked) { "Provide Windows ISO" } else { "Download Windows ESD" }
        SystemPartitionDriveLetter     = $State.Controls.cmbSystemPartitionDriveLetter.SelectedItem.Content
        WindowsPartitionDriveLetter    = $State.Controls.cmbWindowsPartitionDriveLetter.SelectedItem.Content
        RecoveryPartitionDriveLetter   = $State.Controls.cmbRecoveryPartitionDriveLetter.SelectedItem.Content
        LogicalSectorSizeBytes         = [int]$State.Controls.cmbLogicalSectorSize.SelectedItem.Content
        # Make                           = $null
        MediaType                      = $State.Controls.cmbMediaType.SelectedItem
        Memory                         = [int64]$State.Controls.txtMemory.Text * 1GB
        # Model                          = if ($State.Controls.chkDownloadDrivers.IsChecked) {
        #     $selectedModels = $State.Controls.lstDriverModels.Items | Where-Object { $_.IsSelected }
        #     if ($selectedModels.Count -ge 1) {
        #         $selectedModels[0].Model
        #     }
        #     else {
        #         $null
        #     }
        # }
        # else {
        #     $null
        # }
        OfficeConfigXMLFile            = $State.Controls.txtOfficeConfigXMLFilePath.Text
        OfficePath                     = $State.Controls.txtOfficePath.Text
        Optimize                       = $State.Controls.chkOptimize.IsChecked
        OptimizeFFUPartitionNumber     = 0
        OptionalFeatures               = (($State.Controls.featureCheckBoxes.GetEnumerator() | Where-Object { $_.Value.IsChecked } | ForEach-Object { $_.Key } | Sort-Object) -join ';')
        OrchestrationPath              = "$($State.Controls.txtApplicationPath.Text)\Orchestration"
        PEDriversFolder                = $State.Controls.txtPEDriversFolder.Text
        Processors                     = [int]$State.Controls.txtProcessors.Text
        ProductKey                     = $State.Controls.txtProductKey.Text
        PromptExternalHardDiskMedia    = $State.Controls.chkPromptExternalHardDiskMedia.IsChecked
        RemoveApps                     = $State.Controls.chkRemoveApps.IsChecked
        RemoveFFU                      = $State.Controls.chkRemoveFFU.IsChecked
        RemoveUpdates                  = $State.Controls.chkRemoveUpdates.IsChecked
        RemoveDownloadedESD            = $State.Controls.chkRemoveDownloadedESD.IsChecked
        UpdateADK                      = $State.Controls.chkUpdateADK.IsChecked
        UpdateEdge                     = $State.Controls.chkUpdateEdge.IsChecked
        UpdateLatestCU                 = $State.Controls.chkUpdateLatestCU.IsChecked
        UpdateLatestDefender           = $State.Controls.chkUpdateLatestDefender.IsChecked
        UpdateLatestMicrocode          = $State.Controls.chkUpdateLatestMicrocode.IsChecked
        UpdateLatestMSRT               = $State.Controls.chkUpdateLatestMSRT.IsChecked
        UpdateLatestNet                = $State.Controls.chkUpdateLatestNet.IsChecked
        UpdateOneDrive                 = $State.Controls.chkUpdateOneDrive.IsChecked
        UpdatePreviewCU                = $State.Controls.chkUpdatePreviewCU.IsChecked
        UserAppListPath                = $State.Controls.txtUserAppListPath.Text
        USBDriveList                   = @{}
        Threads                        = [int]$State.Controls.txtThreads.Text
        BitsPriority                    = $State.Controls.cmbBitsPriority.SelectedItem
        MaxUSBDrives                   = [int]$State.Controls.txtMaxUSBDrives.Text
        ThemeMode                      = if ($null -ne $State.Controls.cmbThemeMode -and $null -ne $State.Controls.cmbThemeMode.SelectedItem) { $State.Controls.cmbThemeMode.SelectedItem } else { "System" }
        Verbose                        = $State.Controls.chkVerbose.IsChecked
        VMLocation                     = $State.Controls.txtVMLocation.Text
        VMSwitchName                   = if ($State.Controls.cmbVMSwitchName.SelectedItem -eq 'Other') {
            $State.Controls.txtCustomVMSwitchName.Text
        }
        else {
            $State.Controls.cmbVMSwitchName.SelectedItem
        }
        WindowsArch                    = $State.Controls.cmbWindowsArch.SelectedItem
        WindowsLang                    = $State.Controls.cmbWindowsLang.SelectedItem
        WindowsRelease                 = [int]$State.Controls.cmbWindowsRelease.SelectedItem.Value
        WindowsSKU                     = $State.Controls.cmbWindowsSKU.SelectedItem
        WindowsVersion                 = $State.Controls.cmbWindowsVersion.SelectedItem
    }

    # Save selected USB drives using UniqueId for reliable identification
    # Multiple physical drives can share the same Model, so store an array of UniqueIds per Model.
    $State.Controls.lstUSBDrives.Items | Where-Object { $_.IsSelected } | ForEach-Object {
        $modelName = $_.Model
        $uniqueId = $_.UniqueId

        if ([string]::IsNullOrWhiteSpace($modelName) -or [string]::IsNullOrWhiteSpace($uniqueId)) {
            return
        }

        # Ensure the hashtable value is always an array so multiple same-model drives are preserved
        $existingUniqueIds = $config.USBDriveList[$modelName]
        if ($null -eq $existingUniqueIds) {
            $config.USBDriveList[$modelName] = @($uniqueId)
            return
        }

        $existingUniqueIds = @($existingUniqueIds)
        if (-not ($existingUniqueIds -contains $uniqueId)) {
            $existingUniqueIds += $uniqueId
        }
        $config.USBDriveList[$modelName] = $existingUniqueIds
    }

    # Additional FFU file selections
    $config.AdditionalFFUFiles = @()
    if ($State.Controls.chkCopyAdditionalFFUFiles.IsChecked) {
        $config.AdditionalFFUFiles = @(
            $State.Controls.lstAdditionalFFUs.Items |
                Where-Object { $_.IsSelected } |
                ForEach-Object { $_.FullName }
        )
    }
    WriteLog "Get-UIConfig: Saving $($config.AdditionalDataPartitions.Count) additional data partition(s)."
    
    return $config
}

function Set-UIValue {
    param(
        [string]$ControlName,
        [string]$PropertyName,
        [object]$ConfigObject,
        [string]$ConfigKey,
        [scriptblock]$TransformValue = $null, # Optional scriptblock to transform the value from config
        [psobject]$State # Pass the $State object
    )

    $control = $State.Controls[$ControlName]
    if ($null -eq $control) {
        WriteLog "LoadConfig Error: Control '$ControlName' not found in the state object."
        return
    }

    # Robust check for property existence.
    $keyExists = $false
    if ($ConfigObject -is [System.Management.Automation.PSCustomObject] -and $null -ne $ConfigObject.PSObject.Properties) {
        # Use the Match() method, which returns a collection of matching properties.
        # If the count is greater than 0, the key exists.
        try {
            if (($ConfigObject.PSObject.Properties.Match($ConfigKey)).Count -gt 0) {
                $keyExists = $true
            }
        }
        catch {
            WriteLog "ERROR: Exception while trying to Match key '$ConfigKey' on ConfigObject.PSObject.Properties. Error: $($_.Exception.Message)"
            # $keyExists remains false
        }
    }

    if (-not $keyExists) {
        WriteLog "LoadConfig Info: Key '$ConfigKey' not found in configuration object. Skipping '$ControlName.$PropertyName'."
        return
    }

    $valueFromConfig = $ConfigObject.$ConfigKey
    WriteLog "LoadConfig: Preparing to set '$ControlName.$PropertyName'. Config key: '$ConfigKey', Raw value: '$valueFromConfig'."

    $finalValue = $valueFromConfig
    if ($null -ne $TransformValue) {
        try {
            $finalValue = Invoke-Command -ScriptBlock $TransformValue -ArgumentList $valueFromConfig
            WriteLog "LoadConfig: Transformed value for '$ControlName.$PropertyName' (from key '$ConfigKey') is: '$finalValue'."
        }
        catch {
            WriteLog "LoadConfig Error: Failed to transform value for '$ControlName.$PropertyName' from key '$ConfigKey'. Error: $($_.Exception.Message)"
            return
        }
    }

    try {
        # Handle ComboBox SelectedItem specifically
        if ($control -is [System.Windows.Controls.ComboBox] -and $PropertyName -eq 'SelectedItem') {
            $itemToSelect = $null
            # Iterate through the Items collection of the ComboBox
            foreach ($item in $control.Items) {
                $itemValue = $null
                if ($item -is [System.Windows.Controls.ComboBoxItem]) {
                    $itemValue = $item.Content
                }
                elseif ($item -is [pscustomobject] -and $item.PSObject.Properties['Value']) {
                    $itemValue = $item.Value
                }
                elseif ($item -is [pscustomobject] -and $item.PSObject.Properties['Display']) {
                    # Assuming 'Display' might be used if 'Value' isn't
                    $itemValue = $item.Display
                }
                else {
                    $itemValue = $item # For simple string items or direct object comparison
                }

                # Compare, ensuring types are compatible or converting $finalValue if necessary
                if (($null -ne $itemValue -and $itemValue.ToString() -eq $finalValue.ToString()) -or ($item -eq $finalValue)) {
                    $itemToSelect = $item
                    break
                }
            }

            if ($null -ne $itemToSelect) {
                $control.SelectedItem = $itemToSelect
                WriteLog "LoadConfig: Successfully set '$ControlName.SelectedItem' by finding matching item for value '$finalValue'."
            }
            elseif ($control.IsEditable -and ($finalValue -is [string] -or $finalValue -is [int] -or $finalValue -is [long])) {
                $control.Text = $finalValue.ToString()
                WriteLog "LoadConfig: Set '$ControlName.Text' to '$($finalValue.ToString())' as SelectedItem match failed (editable ComboBox)."
            }
            else {
                $itemsString = ""
                try {
                    # Safer way to get item strings
                    $itemStrings = @()
                    foreach ($cbItem in $control.Items) {
                        if ($null -ne $cbItem) { $itemStrings += $cbItem.ToString() } else { $itemStrings += "[NULL_ITEM]" }
                    }
                    $itemsString = $itemStrings -join "; "
                }
                catch { $itemsString = "Error retrieving item strings." }
                WriteLog "LoadConfig Warning: Could not find or set item matching value '$finalValue' for '$ControlName.SelectedItem'. Current items: [$itemsString]"
            }
        }
        else {
            # For other properties or controls
            $control.$PropertyName = $finalValue
            WriteLog "LoadConfig: Successfully set '$ControlName.$PropertyName' to '$finalValue'."
        }
    }
    catch {
        WriteLog "LoadConfig Error: Failed to set '$ControlName.$PropertyName' to '$finalValue'. Error: $($_.Exception.Message)"
    }
}

function Get-ComboBoxSelectedContent {
    param(
        [System.Windows.Controls.ComboBox]$ComboBox
    )

    if ($null -eq $ComboBox -or $null -eq $ComboBox.SelectedItem) {
        return $null
    }

    if ($ComboBox.SelectedItem -is [System.Windows.Controls.ComboBoxItem]) {
        return [string]$ComboBox.SelectedItem.Content
    }

    return [string]$ComboBox.SelectedItem
}

function ConvertTo-PartitionSizeBytesFromGBText {
    param(
        [AllowNull()]
        [object]$Text,
        [string]$FieldName
    )

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

    $trimmedText = ([string]$Text).Trim()
    [decimal]$sizeGb = 0
    if (-not [decimal]::TryParse($trimmedText, [System.Globalization.NumberStyles]::Number, [System.Globalization.CultureInfo]::InvariantCulture, [ref]$sizeGb)) {
        throw "$FieldName must be a number in GB."
    }

    if ($sizeGb -lt 0) {
        throw "$FieldName cannot be negative."
    }

    return [int64]($sizeGb * 1GB)
}

function Get-PartitionSizeGBDisplay {
    param(
        [object]$SizeBytes
    )

    [decimal]$parsedSizeBytes = 0
    if ($null -eq $SizeBytes -or -not [decimal]::TryParse([string]$SizeBytes, [ref]$parsedSizeBytes) -or $parsedSizeBytes -le 0) {
        return ''
    }

    $sizeGb = $parsedSizeBytes / 1GB
    if (($parsedSizeBytes % 1GB) -eq 0) {
        return ([int64]$sizeGb).ToString([System.Globalization.CultureInfo]::InvariantCulture)
    }

    return $sizeGb.ToString('0.##', [System.Globalization.CultureInfo]::InvariantCulture)
}

function Get-PartitionDriveLetterOptions {
    return @('A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z')
}

function New-DiskLayoutPartitionRow {
    param(
        [Parameter(Mandatory = $true)]
        [string]$PartitionType,
        [Parameter(Mandatory = $true)]
        [string]$Name,
        [string]$Label,
        [string]$DriveLetter,
        [string]$SizeGB,
        [int64]$SizeBytes = 0,
        [bool]$FillRemaining,
        [string]$FileSystem = 'NTFS',
        [bool]$CanSelect,
        [bool]$CanEditDriveLetter,
        [bool]$CanEditSize,
        [bool]$CanEditFillRemaining,
        [string]$FillRemainingVisibility = 'Visible',
        [bool]$CanRemove,
        [bool]$CanReorder
    )

    if ([string]::IsNullOrWhiteSpace($Label)) {
        $Label = $Name
    }

    return [PSCustomObject]@{
        IsSelected           = $false
        CanSelect            = $CanSelect
        PartitionType        = $PartitionType
        Name                 = $Name
        Label                = $Label
        DriveLetter          = $DriveLetter
        DriveLetterOptions   = @(Get-PartitionDriveLetterOptions)
        SizeGB               = $SizeGB
        SizeBytes            = $SizeBytes
        FillRemaining        = $FillRemaining
        FillRemainingVisibility = $FillRemainingVisibility
        FileSystem           = $FileSystem
        CanEditDriveLetter   = $CanEditDriveLetter
        CanEditSize          = $CanEditSize
        CanEditFillRemaining = $CanEditFillRemaining
        CanRemove            = $CanRemove
        CanReorder           = $CanReorder
    }
}

function Get-ComboBoxDriveLetterValue {
    param(
        [object]$ComboBox,
        [string]$DefaultValue
    )

    $driveLetter = [string](Get-ComboBoxSelectedContent -ComboBox $ComboBox)
    if ([string]::IsNullOrWhiteSpace($driveLetter)) {
        $driveLetter = $DefaultValue
    }

    return $driveLetter.Trim().TrimEnd(':').ToUpperInvariant()
}

function Get-DiskLayoutPartitionRows {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    if ($null -eq $State.Data.additionalDataPartitionsDataList) {
        $State.Data.additionalDataPartitionsDataList = [System.Collections.Generic.List[PSCustomObject]]::new()
    }
    if ($null -eq $State.Data.createRecoveryPartition) {
        $State.Data.createRecoveryPartition = $true
    }

    $hasDataPartitions = @($State.Data.additionalDataPartitionsDataList).Count -gt 0
    $windowsSizeText = [string]$State.Controls.txtOSPartitionSizeGB.Text
    $windowsCanFillRemaining = -not $hasDataPartitions
    $windowsFillRemaining = $windowsCanFillRemaining -and [string]::IsNullOrWhiteSpace($windowsSizeText)
    $windowsFillRemainingVisibility = if ($windowsCanFillRemaining) { 'Visible' } else { 'Hidden' }

    $rows = [System.Collections.Generic.List[PSCustomObject]]::new()
    $rows.Add((New-DiskLayoutPartitionRow -PartitionType 'System' -Name 'System' -DriveLetter (Get-ComboBoxDriveLetterValue -ComboBox $State.Controls.cmbSystemPartitionDriveLetter -DefaultValue 'S') -SizeGB '0.26' -SizeBytes 260MB -CanEditDriveLetter $true -CanEditSize $false -CanEditFillRemaining $false -FillRemainingVisibility 'Hidden' -CanRemove $false -CanReorder $false))
    $rows.Add((New-DiskLayoutPartitionRow -PartitionType 'MSR' -Name 'MSR' -SizeGB '0.016' -SizeBytes 16MB -CanEditDriveLetter $false -CanEditSize $false -CanEditFillRemaining $false -FillRemainingVisibility 'Hidden' -CanRemove $false -CanReorder $false))
    $rows.Add((New-DiskLayoutPartitionRow -PartitionType 'Windows' -Name 'Windows' -DriveLetter (Get-ComboBoxDriveLetterValue -ComboBox $State.Controls.cmbWindowsPartitionDriveLetter -DefaultValue 'W') -SizeGB $windowsSizeText -SizeBytes (ConvertTo-PartitionSizeBytesFromGBText -Text $windowsSizeText -FieldName 'Windows Partition Size') -FillRemaining $windowsFillRemaining -CanEditDriveLetter $true -CanEditSize $true -CanEditFillRemaining $windowsCanFillRemaining -FillRemainingVisibility $windowsFillRemainingVisibility -CanRemove $false -CanReorder $false))

    if ([bool]$State.Data.createRecoveryPartition) {
        $rows.Add((New-DiskLayoutPartitionRow -PartitionType 'Recovery' -Name 'Recovery' -DriveLetter (Get-ComboBoxDriveLetterValue -ComboBox $State.Controls.cmbRecoveryPartitionDriveLetter -DefaultValue 'R') -SizeGB ([string]$State.Controls.txtRecoveryPartitionSizeGB.Text) -SizeBytes (ConvertTo-PartitionSizeBytesFromGBText -Text $State.Controls.txtRecoveryPartitionSizeGB.Text -FieldName 'Recovery Partition Size') -CanSelect $true -CanEditDriveLetter $true -CanEditSize $true -CanEditFillRemaining $false -FillRemainingVisibility 'Hidden' -CanRemove $true -CanReorder $false))
    }

    foreach ($dataPartition in @($State.Data.additionalDataPartitionsDataList)) {
        $rows.Add((New-DiskLayoutPartitionRow -PartitionType 'Data' -Name $dataPartition.Name -Label $dataPartition.Label -DriveLetter $dataPartition.DriveLetter -SizeGB $dataPartition.SizeGB -SizeBytes ([int64]$dataPartition.SizeBytes) -FillRemaining ([bool]$dataPartition.FillRemaining) -FileSystem $dataPartition.FileSystem -CanSelect $true -CanEditDriveLetter $true -CanEditSize $true -CanEditFillRemaining $true -CanRemove $true -CanReorder $true))
    }

    return $rows.ToArray()
}

function Sync-DiskLayoutRowsToControls {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    if ($null -eq $State.Controls.lstDataPartitions) { return }
    if ($null -eq $State.Data.createRecoveryPartition) {
        $State.Data.createRecoveryPartition = $true
    }

    $layoutRows = @($State.Controls.lstDataPartitions.ItemsSource)
    if ($layoutRows.Count -eq 0) { return }

    $State.Data.createRecoveryPartition = [bool](@($layoutRows | Where-Object { $_.PartitionType -eq 'Recovery' }).Count -gt 0)
    $updatedDataPartitions = [System.Collections.Generic.List[PSCustomObject]]::new()

    foreach ($row in $layoutRows) {
        switch ($row.PartitionType) {
            'System' {
                $row.DriveLetter = ([string]$row.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
                $State.Controls.cmbSystemPartitionDriveLetter.SelectedItem = ($State.Controls.cmbSystemPartitionDriveLetter.Items | Where-Object { $_.Content -eq $row.DriveLetter } | Select-Object -First 1)
            }
            'Windows' {
                $row.DriveLetter = ([string]$row.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
                $State.Controls.cmbWindowsPartitionDriveLetter.SelectedItem = ($State.Controls.cmbWindowsPartitionDriveLetter.Items | Where-Object { $_.Content -eq $row.DriveLetter } | Select-Object -First 1)
                if ($row.FillRemaining) {
                    $State.Controls.txtOSPartitionSizeGB.Clear()
                }
                else {
                    $State.Controls.txtOSPartitionSizeGB.Text = [string]$row.SizeGB
                }
            }
            'Recovery' {
                $row.DriveLetter = ([string]$row.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
                $State.Controls.cmbRecoveryPartitionDriveLetter.SelectedItem = ($State.Controls.cmbRecoveryPartitionDriveLetter.Items | Where-Object { $_.Content -eq $row.DriveLetter } | Select-Object -First 1)
                $State.Controls.txtRecoveryPartitionSizeGB.Text = [string]$row.SizeGB
            }
            'Data' {
                $row.DriveLetter = ([string]$row.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
                if ($row.FillRemaining) {
                    $row.SizeGB = ''
                    $row.SizeBytes = 0
                }
                else {
                    $row.SizeBytes = ConvertTo-PartitionSizeBytesFromGBText -Text $row.SizeGB -FieldName "Data Partition '$($row.Name)' Size"
                }
                $updatedDataPartitions.Add([PSCustomObject]@{
                        Name          = $row.Name
                        Label         = $row.Label
                        DriveLetter   = $row.DriveLetter
                        SizeGB        = $row.SizeGB
                        SizeBytes     = [int64]$row.SizeBytes
                        FillRemaining = [bool]$row.FillRemaining
                        FileSystem    = $row.FileSystem
                    })
            }
        }
    }

    if ($null -eq $State.Data.additionalDataPartitionsDataList) {
        $State.Data.additionalDataPartitionsDataList = [System.Collections.Generic.List[PSCustomObject]]::new()
    }
    $State.Data.additionalDataPartitionsDataList.Clear()
    foreach ($dataPartition in $updatedDataPartitions) {
        $State.Data.additionalDataPartitionsDataList.Add($dataPartition)
    }
}

function Update-AdditionalDataPartitionsListView {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    if ($null -eq $State.Data.additionalDataPartitionsDataList) {
        $State.Data.additionalDataPartitionsDataList = [System.Collections.Generic.List[PSCustomObject]]::new()
    }

    if ($null -ne $State.Controls.lstDataPartitions) {
        $State.Controls.lstDataPartitions.ItemsSource = @(Get-DiskLayoutPartitionRows -State $State)
        $State.Controls.lstDataPartitions.Items.Refresh()
        Request-ListViewColumnAutoResize -ListView $State.Controls.lstDataPartitions
        if ($null -ne $State.Controls.chkSelectAllDataPartitions) {
            Update-SelectAllHeaderCheckBoxState -ListView $State.Controls.lstDataPartitions -HeaderCheckBox $State.Controls.chkSelectAllDataPartitions
        }
        Update-DiskLayoutActionButtonsState -State $State
        Update-DiskLayoutCapacityStatus -State $State
    }
}

function Get-DiskLayoutSizeDisplay {
    param(
        [int64]$SizeBytes
    )

    $sizeGb = [decimal]$SizeBytes / 1GB
    return $sizeGb.ToString('0.##', [System.Globalization.CultureInfo]::InvariantCulture)
}

function Get-DiskLayoutCapacityStatus {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    try {
        $diskSizeBytes = ConvertTo-PartitionSizeBytesFromGBText -Text $State.Controls.txtDiskSize.Text -FieldName 'Disk Size'
    }
    catch {
        return [PSCustomObject]@{ Level = 'Red'; Message = $_.Exception.Message }
    }

    if ($diskSizeBytes -le 0) {
        return [PSCustomObject]@{ Level = 'Red'; Message = 'Disk Size must be greater than 0 GB.' }
    }

    $partitionRows = @($State.Controls.lstDataPartitions.ItemsSource)
    if ($partitionRows.Count -eq 0) {
        $partitionRows = @(Get-DiskLayoutPartitionRows -State $State)
    }

    $fixedPartitionSizeBytes = [int64]0
    $fillRemainingPartitionNames = [System.Collections.Generic.List[string]]::new()
    $hasAutoSizedRecovery = $false
    $dataPartitionCount = @($partitionRows | Where-Object { $_.PartitionType -eq 'Data' }).Count

    foreach ($partitionRow in $partitionRows) {
        if ($partitionRow.PartitionType -eq 'Recovery' -and -not [bool]$State.Data.createRecoveryPartition) { continue }

        if ($partitionRow.FillRemaining) {
            $fillRemainingPartitionNames.Add([string]$partitionRow.Name)
            continue
        }

        if ($partitionRow.PartitionType -eq 'Windows' -and $dataPartitionCount -gt 0 -and [string]::IsNullOrWhiteSpace([string]$partitionRow.SizeGB)) {
            return [PSCustomObject]@{ Level = 'Red'; Message = 'Set a fixed Windows partition size before adding data partitions.' }
        }

        if ($partitionRow.PartitionType -eq 'Recovery' -and [string]::IsNullOrWhiteSpace([string]$partitionRow.SizeGB)) {
            $hasAutoSizedRecovery = $true
            continue
        }

        if ($partitionRow.PartitionType -eq 'Data' -and [string]::IsNullOrWhiteSpace([string]$partitionRow.SizeGB)) {
            return [PSCustomObject]@{ Level = 'Red'; Message = "Data partition '$($partitionRow.Name)' must have a size or Fill Remaining selected." }
        }

        if (-not [bool]$partitionRow.CanEditSize -and [int64]$partitionRow.SizeBytes -gt 0) {
            $fixedPartitionSizeBytes += [int64]$partitionRow.SizeBytes
            continue
        }

        try {
            $partitionSizeBytes = ConvertTo-PartitionSizeBytesFromGBText -Text $partitionRow.SizeGB -FieldName "$($partitionRow.Name) Partition Size"
        }
        catch {
            return [PSCustomObject]@{ Level = 'Red'; Message = $_.Exception.Message }
        }

        if ($partitionSizeBytes -gt 0) {
            $fixedPartitionSizeBytes += $partitionSizeBytes
        }
    }

    if ($fillRemainingPartitionNames.Count -gt 1) {
        return [PSCustomObject]@{ Level = 'Red'; Message = 'Only one partition can fill remaining disk space.' }
    }

    if ($fixedPartitionSizeBytes -gt $diskSizeBytes) {
        $overByBytes = $fixedPartitionSizeBytes - $diskSizeBytes
        return [PSCustomObject]@{ Level = 'Red'; Message = "Partition sizes exceed Disk Size by $(Get-DiskLayoutSizeDisplay -SizeBytes $overByBytes) GB. Fixed total: $(Get-DiskLayoutSizeDisplay -SizeBytes $fixedPartitionSizeBytes) GB of $(Get-DiskLayoutSizeDisplay -SizeBytes $diskSizeBytes) GB." }
    }

    $remainingBytes = $diskSizeBytes - $fixedPartitionSizeBytes
    $remainingPurpose = if ($fillRemainingPartitionNames.Count -gt 0 -and $hasAutoSizedRecovery) {
        "remaining for $($fillRemainingPartitionNames -join ', ') and the auto-sized Recovery partition"
    }
    elseif ($fillRemainingPartitionNames.Count -gt 0) {
        "remaining for $($fillRemainingPartitionNames -join ', ')"
    }
    elseif ($hasAutoSizedRecovery) {
        'remaining for the auto-sized Recovery partition and unallocated space'
    }
    else {
        'left unallocated'
    }

    return [PSCustomObject]@{ Level = 'Green'; Message = "Partition sizes fit within Disk Size. $(Get-DiskLayoutSizeDisplay -SizeBytes $remainingBytes) GB $remainingPurpose." }
}

function Update-DiskLayoutCapacityStatus {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    if ($null -eq $State.Controls.ellipseDiskLayoutCapacityStatus -or $null -eq $State.Controls.txtDiskLayoutCapacityStatusValue) { return }
    if ($State.Flags -is [System.Collections.IDictionary] -and $true -eq $State.Flags['updatingDiskLayoutCapacityStatus']) { return }

    try {
        if ($State.Flags -is [System.Collections.IDictionary]) {
            $State.Flags['updatingDiskLayoutCapacityStatus'] = $true
        }

        $capacityStatus = Get-DiskLayoutCapacityStatus -State $State
        $State.Controls.ellipseDiskLayoutCapacityStatus.Fill = switch ($capacityStatus.Level) {
            'Green' { [System.Windows.Media.Brushes]::LimeGreen }
            'Red' { [System.Windows.Media.Brushes]::IndianRed }
            default { [System.Windows.Media.Brushes]::Gold }
        }
        $State.Controls.txtDiskLayoutCapacityStatusValue.Text = $capacityStatus.Message
    }
    finally {
        if ($State.Flags -is [System.Collections.IDictionary]) {
            $State.Flags['updatingDiskLayoutCapacityStatus'] = $false
        }
    }
}

function Update-DiskLayoutFillRemainingState {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State,
        [Parameter(Mandatory = $true)]
        [object]$PartitionRow
    )

    if ($PartitionRow.FillRemaining) {
        $PartitionRow.SizeGB = ''
        $PartitionRow.SizeBytes = 0
        if ($PartitionRow.PartitionType -eq 'Windows') {
            $State.Controls.txtOSPartitionSizeGB.Clear()
        }
    }

    if ($null -ne $State.Controls.lstDataPartitions) {
        $State.Controls.lstDataPartitions.Items.Refresh()
    }
    Update-DiskLayoutCapacityStatus -State $State
}

function Update-DiskLayoutActionButtonsState {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    if ($null -eq $State.Controls.lstDataPartitions) { return }

    $selectedItem = $State.Controls.lstDataPartitions.SelectedItem
    $selectedDataRows = @($State.Controls.lstDataPartitions.Items | Where-Object { $_.PartitionType -eq 'Data' })
    $selectedDataIndex = -1
    for ($dataIndex = 0; $dataIndex -lt $selectedDataRows.Count; $dataIndex++) {
        if ([object]::ReferenceEquals($selectedDataRows[$dataIndex], $selectedItem)) {
            $selectedDataIndex = $dataIndex
            break
        }
    }

    $canMove = $selectedDataIndex -ge 0
    if ($null -ne $State.Controls.btnMoveDataPartitionTop) { $State.Controls.btnMoveDataPartitionTop.IsEnabled = $canMove -and $selectedDataIndex -gt 0 }
    if ($null -ne $State.Controls.btnMoveDataPartitionUp) { $State.Controls.btnMoveDataPartitionUp.IsEnabled = $canMove -and $selectedDataIndex -gt 0 }
    if ($null -ne $State.Controls.btnMoveDataPartitionDown) { $State.Controls.btnMoveDataPartitionDown.IsEnabled = $canMove -and $selectedDataIndex -lt ($selectedDataRows.Count - 1) }
    if ($null -ne $State.Controls.btnMoveDataPartitionBottom) { $State.Controls.btnMoveDataPartitionBottom.IsEnabled = $canMove -and $selectedDataIndex -lt ($selectedDataRows.Count - 1) }

    if ($null -ne $State.Controls.btnRestoreRecoveryPartition) {
        $State.Controls.btnRestoreRecoveryPartition.Visibility = if ([bool]$State.Data.createRecoveryPartition) { 'Collapsed' } else { 'Visible' }
        $State.Controls.btnRestoreRecoveryPartition.IsEnabled = -not [bool]$State.Data.createRecoveryPartition
    }
}

function Move-DataPartitionRow {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State,
        [Parameter(Mandatory = $true)]
        [ValidateSet('Top', 'Up', 'Down', 'Bottom')]
        [string]$Direction
    )

    $selectedItem = $State.Controls.lstDataPartitions.SelectedItem
    if ($null -eq $selectedItem -or $selectedItem.PartitionType -ne 'Data') { return }

    Sync-DiskLayoutRowsToControls -State $State
    $dataRows = @($State.Controls.lstDataPartitions.Items | Where-Object { $_.PartitionType -eq 'Data' })
    $currentIndex = -1
    for ($dataIndex = 0; $dataIndex -lt $dataRows.Count; $dataIndex++) {
        if ([object]::ReferenceEquals($dataRows[$dataIndex], $selectedItem)) {
            $currentIndex = $dataIndex
            break
        }
    }

    if ($currentIndex -lt 0) { return }

    $targetIndex = switch ($Direction) {
        'Top' { 0 }
        'Up' { [Math]::Max(0, $currentIndex - 1) }
        'Down' { [Math]::Min($State.Data.additionalDataPartitionsDataList.Count - 1, $currentIndex + 1) }
        'Bottom' { $State.Data.additionalDataPartitionsDataList.Count - 1 }
    }

    if ($targetIndex -eq $currentIndex) { return }

    $movingItem = $State.Data.additionalDataPartitionsDataList[$currentIndex]
    $State.Data.additionalDataPartitionsDataList.RemoveAt($currentIndex)
    $State.Data.additionalDataPartitionsDataList.Insert($targetIndex, $movingItem)
    Update-AdditionalDataPartitionsListView -State $State

    $updatedDataRows = @($State.Controls.lstDataPartitions.Items | Where-Object { $_.PartitionType -eq 'Data' })
    if ($targetIndex -ge 0 -and $targetIndex -lt $updatedDataRows.Count) {
        $State.Controls.lstDataPartitions.SelectedItem = $updatedDataRows[$targetIndex]
    }

    Update-DiskLayoutActionButtonsState -State $State
}

function Clear-AdditionalDataPartitions {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    $result = [System.Windows.MessageBox]::Show("Are you sure you want to clear all additional data partitions?", "Clear Data Partitions", [System.Windows.MessageBoxButton]::YesNo, [System.Windows.MessageBoxImage]::Question)
    if ($result -ne [System.Windows.MessageBoxResult]::Yes) { return }

    if ($null -eq $State.Data.additionalDataPartitionsDataList) {
        $State.Data.additionalDataPartitionsDataList = [System.Collections.Generic.List[PSCustomObject]]::new()
    }
    $State.Data.additionalDataPartitionsDataList.Clear()
    Clear-AdditionalDataPartitionForm -State $State
    Update-AdditionalDataPartitionsListView -State $State

    if ($null -ne $State.Controls.txtStatus) {
        $State.Controls.txtStatus.Text = "Additional data partitions list cleared."
    }
}

function Restore-RecoveryPartition {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    $State.Data.createRecoveryPartition = $true
    Update-AdditionalDataPartitionsListView -State $State
}

function Test-DiskLayoutConfiguration {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State,
        [Parameter(Mandatory = $true)]
        [System.Collections.IDictionary]$Config
    )

    $errors = [System.Collections.Generic.List[string]]::new()
    $driveLetterEntries = [System.Collections.Generic.List[PSCustomObject]]::new()

    $partitionRows = @($State.Controls.lstDataPartitions.ItemsSource)
    foreach ($partitionRow in $partitionRows) {
        if ($partitionRow.PartitionType -eq 'MSR') { continue }
        if ($partitionRow.PartitionType -eq 'Recovery' -and -not [bool]$Config.CreateRecoveryPartition) { continue }

        $driveLetter = ([string]$partitionRow.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
        if ([string]::IsNullOrWhiteSpace($driveLetter) -or $driveLetter -notmatch '^[A-Z]$') {
            $errors.Add("$($partitionRow.Name) must use a single drive letter from A to Z.")
            continue
        }

        $driveLetterEntries.Add([PSCustomObject]@{
                Name        = $partitionRow.Name
                DriveLetter = $driveLetter
            })
    }

    $duplicateDriveLetters = @($driveLetterEntries | Group-Object -Property DriveLetter | Where-Object { $_.Count -gt 1 })
    foreach ($duplicateDriveLetter in $duplicateDriveLetters) {
        $partitionNames = ($duplicateDriveLetter.Group | ForEach-Object { $_.Name }) -join ', '
        $errors.Add("Drive letter $($duplicateDriveLetter.Name) is used by multiple partitions: $partitionNames.")
    }

    $dataPartitions = @($Config.AdditionalDataPartitions)
    if ($dataPartitions.Count -gt 0 -and $Config.OSPartitionSize -le 0) {
        $errors.Add('Windows must use a fixed size when data partitions are configured.')
    }

    $fillRemainingCount = 0
    if ($Config.OSPartitionSize -le 0) {
        $fillRemainingCount++
    }
    $fillRemainingCount += @($dataPartitions | Where-Object { $_.FillRemaining }).Count
    if ($fillRemainingCount -gt 1) {
        $errors.Add('Only one partition can fill remaining disk space.')
    }

    $duplicateDataNames = @($dataPartitions | Group-Object -Property Name | Where-Object { $_.Count -gt 1 })
    foreach ($duplicateDataName in $duplicateDataNames) {
        $errors.Add("Data partition name '$($duplicateDataName.Name)' is used more than once.")
    }

    [uint64]$fixedPartitionSizeBytes = 260MB + 16MB
    if ($Config.OSPartitionSize -gt 0) {
        $fixedPartitionSizeBytes += [uint64]$Config.OSPartitionSize
    }
    if ([bool]$Config.CreateRecoveryPartition -and $Config.RecoveryPartitionSize -gt 0) {
        $fixedPartitionSizeBytes += [uint64]$Config.RecoveryPartitionSize
    }
    foreach ($dataPartition in $dataPartitions) {
        if ($dataPartition.FillRemaining) { continue }
        if ([uint64]$dataPartition.SizeBytes -le 0) {
            $errors.Add("Data partition '$($dataPartition.Name)' must have a size or Fill Remaining selected.")
            continue
        }
        $fixedPartitionSizeBytes += [uint64]$dataPartition.SizeBytes
    }

    if ([uint64]$Config.Disksize -le 0) {
        $errors.Add('Disk Size must be greater than 0 GB.')
    }
    elseif ($fixedPartitionSizeBytes -gt [uint64]$Config.Disksize) {
        $fixedPartitionSizeGB = [Math]::Round(($fixedPartitionSizeBytes / 1GB), 2)
        $diskSizeGB = [Math]::Round(([uint64]$Config.Disksize / 1GB), 2)
        $errors.Add("Fixed partition sizes total $fixedPartitionSizeGB GB, which exceeds the configured disk size of $diskSizeGB GB.")
    }

    if ($errors.Count -gt 0) {
        [System.Windows.MessageBox]::Show(($errors -join [System.Environment]::NewLine), "Disk Layout Validation", "OK", "Warning") | Out-Null
        return $false
    }

    return $true
}

function Clear-AdditionalDataPartitionForm {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    $State.Controls.txtDataPartitionName.Clear()
    $State.Controls.txtDataPartitionSizeGB.Clear()
    $State.Controls.chkDataPartitionFillRemaining.IsChecked = $false
    $State.Controls.txtDataPartitionSizeGB.IsEnabled = $true
}

function Add-AdditionalDataPartition {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    try {
        Sync-DiskLayoutRowsToControls -State $State
    }
    catch {
        [System.Windows.MessageBox]::Show($_.Exception.Message, "Disk Layout", "OK", "Warning") | Out-Null
        return $false
    }

    if ($null -eq $State.Data.additionalDataPartitionsDataList) {
        $State.Data.additionalDataPartitionsDataList = [System.Collections.Generic.List[PSCustomObject]]::new()
    }

    $partitionName = ([string]$State.Controls.txtDataPartitionName.Text).Trim()
    if ([string]::IsNullOrWhiteSpace($partitionName)) {
        $partitionName = "Data$($State.Data.additionalDataPartitionsDataList.Count + 1)"
    }

    if ($State.Data.additionalDataPartitionsDataList | Where-Object { $_.Name -ieq $partitionName } | Select-Object -First 1) {
        [System.Windows.MessageBox]::Show("A data partition named '$partitionName' already exists.", "Duplicate Data Partition", "OK", "Warning") | Out-Null
        return $false
    }

    $driveLetter = ([string](Get-ComboBoxSelectedContent -ComboBox $State.Controls.cmbDataPartitionDriveLetter)).Trim().TrimEnd(':').ToUpperInvariant()
    if ([string]::IsNullOrWhiteSpace($driveLetter) -or $driveLetter -notmatch '^[A-Z]$') {
        [System.Windows.MessageBox]::Show("Select a drive letter for the data partition.", "Data Partition Drive Letter", "OK", "Warning") | Out-Null
        return $false
    }

    $reservedDriveLetters = @(
        Get-ComboBoxSelectedContent -ComboBox $State.Controls.cmbSystemPartitionDriveLetter
        Get-ComboBoxSelectedContent -ComboBox $State.Controls.cmbWindowsPartitionDriveLetter
    ) | ForEach-Object { ([string]$_).Trim().TrimEnd(':').ToUpperInvariant() } | Where-Object { -not [string]::IsNullOrWhiteSpace($_) }
    if ([bool]$State.Data.createRecoveryPartition) {
        $reservedDriveLetters += ([string](Get-ComboBoxSelectedContent -ComboBox $State.Controls.cmbRecoveryPartitionDriveLetter)).Trim().TrimEnd(':').ToUpperInvariant()
    }

    if ([string]::IsNullOrWhiteSpace([string]$State.Controls.txtOSPartitionSizeGB.Text)) {
        [System.Windows.MessageBox]::Show("Set a fixed Windows partition size before adding data partitions. Windows can fill remaining space only when no data partitions are configured.", "Windows Partition Size Required", "OK", "Warning") | Out-Null
        return $false
    }

    if ($reservedDriveLetters -contains $driveLetter) {
        [System.Windows.MessageBox]::Show("Drive letter $driveLetter is already used by a required build partition.", "Duplicate Drive Letter", "OK", "Warning") | Out-Null
        return $false
    }

    if ($State.Data.additionalDataPartitionsDataList | Where-Object { $_.DriveLetter -eq $driveLetter } | Select-Object -First 1) {
        [System.Windows.MessageBox]::Show("Drive letter $driveLetter is already used by another data partition.", "Duplicate Drive Letter", "OK", "Warning") | Out-Null
        return $false
    }

    $fillRemaining = $true -eq $State.Controls.chkDataPartitionFillRemaining.IsChecked
    if ($fillRemaining -and ($State.Data.additionalDataPartitionsDataList | Where-Object { $_.FillRemaining } | Select-Object -First 1)) {
        [System.Windows.MessageBox]::Show("Only one data partition can fill the remaining VHDX space.", "Fill Remaining Already Used", "OK", "Warning") | Out-Null
        return $false
    }

    $sizeBytes = [int64]0
    $sizeGbDisplay = ''
    if (-not $fillRemaining) {
        try {
            $sizeBytes = ConvertTo-PartitionSizeBytesFromGBText -Text $State.Controls.txtDataPartitionSizeGB.Text -FieldName 'Data Partition Size'
        }
        catch {
            [System.Windows.MessageBox]::Show($_.Exception.Message, "Data Partition Size", "OK", "Warning") | Out-Null
            return $false
        }

        if ($sizeBytes -le 0) {
            [System.Windows.MessageBox]::Show("Enter a data partition size or select Fill Remaining.", "Data Partition Size", "OK", "Warning") | Out-Null
            return $false
        }

        $sizeGbDisplay = Get-PartitionSizeGBDisplay -SizeBytes $sizeBytes
    }

    $newItem = [PSCustomObject]@{
        Name          = $partitionName
        Label         = $partitionName
        DriveLetter   = $driveLetter
        SizeGB        = $sizeGbDisplay
        SizeBytes     = $sizeBytes
        FillRemaining = $fillRemaining
        FileSystem    = 'NTFS'
    }

    $State.Data.additionalDataPartitionsDataList.Add($newItem)
    WriteLog "Added additional data partition '$partitionName' (DriveLetter=$driveLetter, SizeBytes=$sizeBytes, FillRemaining=$fillRemaining)."
    Update-AdditionalDataPartitionsListView -State $State
    Clear-AdditionalDataPartitionForm -State $State
    return $true
}

function Add-PendingAdditionalDataPartition {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    $hasPendingDataPartitionInput = -not [string]::IsNullOrWhiteSpace([string]$State.Controls.txtDataPartitionName.Text) -or
        -not [string]::IsNullOrWhiteSpace([string]$State.Controls.txtDataPartitionSizeGB.Text) -or
        ($true -eq $State.Controls.chkDataPartitionFillRemaining.IsChecked)

    if (-not $hasPendingDataPartitionInput) {
        return $true
    }

    WriteLog 'Detected pending data partition form input. Adding it before saving build configuration.'
    return ($true -eq (Add-AdditionalDataPartition -State $State))
}

function Remove-SelectedDataPartition {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    $itemsToRemove = @($State.Controls.lstDataPartitions.Items | Where-Object { $_.IsSelected -and $_.CanRemove })
    if ($itemsToRemove.Count -eq 0) {
        [System.Windows.MessageBox]::Show("Select one or more removable partitions to remove.", "Selection Required", "OK", "Warning") | Out-Null
        return
    }

    foreach ($itemToRemove in $itemsToRemove) {
        if ($itemToRemove.PartitionType -eq 'Recovery') {
            $State.Data.createRecoveryPartition = $false
        }
        elseif ($itemToRemove.PartitionType -eq 'Data') {
            $dataItem = @($State.Data.additionalDataPartitionsDataList | Where-Object { $_.Name -eq $itemToRemove.Name -and $_.DriveLetter -eq $itemToRemove.DriveLetter } | Select-Object -First 1)
            if ($dataItem.Count -gt 0) {
                $State.Data.additionalDataPartitionsDataList.Remove($dataItem[0]) | Out-Null
            }
        }
    }

    Update-AdditionalDataPartitionsListView -State $State
}

function Get-AdditionalDataPartitionConfigRows {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    if ($null -eq $State.Data.additionalDataPartitionsDataList) {
        return @()
    }

    Sync-DiskLayoutRowsToControls -State $State

    return @($State.Data.additionalDataPartitionsDataList | ForEach-Object {
            [PSCustomObject]@{
                Name          = $_.Name
                Label         = $_.Label
                DriveLetter   = $_.DriveLetter
                SizeBytes     = [int64]$_.SizeBytes
                FillRemaining = [bool]$_.FillRemaining
                FileSystem    = $_.FileSystem
            }
        })
}

function Import-AdditionalDataPartitionsFromConfig {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State,
        [Parameter(Mandatory = $true)]
        [object]$ConfigContent
    )

    if ($null -eq $State.Data.additionalDataPartitionsDataList) {
        $State.Data.additionalDataPartitionsDataList = [System.Collections.Generic.List[PSCustomObject]]::new()
    }

    $State.Data.additionalDataPartitionsDataList.Clear()

    $keyExists = $false
    if ($ConfigContent -is [System.Management.Automation.PSCustomObject] -and $null -ne $ConfigContent.PSObject.Properties) {
        try {
            if (($ConfigContent.PSObject.Properties.Match('AdditionalDataPartitions')).Count -gt 0) {
                $keyExists = $true
            }
        }
        catch {
            WriteLog "ERROR: Exception while trying to Match key 'AdditionalDataPartitions'. Error: $($_.Exception.Message)"
        }
    }

    if (-not $keyExists -or $null -eq $ConfigContent.AdditionalDataPartitions) {
        Update-AdditionalDataPartitionsListView -State $State
        return
    }

    $partitionIndex = 0
    foreach ($partition in @($ConfigContent.AdditionalDataPartitions)) {
        if ($null -eq $partition) { continue }
        $partitionIndex++

        $partitionName = [string]$partition.Name
        $partitionLabel = [string]$partition.Label
        if ([string]::IsNullOrWhiteSpace($partitionName)) { $partitionName = $partitionLabel }
        if ([string]::IsNullOrWhiteSpace($partitionName)) { $partitionName = "Data$partitionIndex" }
        if ([string]::IsNullOrWhiteSpace($partitionLabel)) { $partitionLabel = $partitionName }

        $driveLetter = ([string]$partition.DriveLetter).Trim().TrimEnd(':').ToUpperInvariant()
        $fileSystem = [string]$partition.FileSystem
        if ([string]::IsNullOrWhiteSpace($fileSystem)) { $fileSystem = 'NTFS' }

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

        [int64]$sizeBytes = 0
        if ($partition.PSObject.Properties.Name -contains 'SizeBytes') {
            [int64]::TryParse([string]$partition.SizeBytes, [ref]$sizeBytes) | Out-Null
        }
        elseif ($partition.PSObject.Properties.Name -contains 'SizeGB') {
            $sizeBytes = ConvertTo-PartitionSizeBytesFromGBText -Text $partition.SizeGB -FieldName "Data Partition '$partitionName' Size"
        }

        $State.Data.additionalDataPartitionsDataList.Add([PSCustomObject]@{
                Name          = $partitionName
                Label         = $partitionLabel
                DriveLetter   = $driveLetter
                SizeGB        = Get-PartitionSizeGBDisplay -SizeBytes $sizeBytes
                SizeBytes     = $sizeBytes
                FillRemaining = $fillRemaining
                FileSystem    = $fileSystem
            })
    }

    Update-AdditionalDataPartitionsListView -State $State
}

function Get-ConfigDriverBaseName {
    param(
        [string]$RawName
    )

    if ([string]::IsNullOrWhiteSpace($RawName)) {
        return $RawName
    }

    if ($RawName -match '^(.*?)\s*\((.+)\)\s*$') {
        return $matches[1].Trim()
    }

    return $RawName.Trim()
}

function Get-ConfigDriverDisplayName {
    param(
        [string]$Make,
        [string]$StoredName,
        [string]$ProductName,
        [string]$SystemId,
        [string]$MachineType
    )

    $baseName = if (-not [string]::IsNullOrWhiteSpace($ProductName)) { $ProductName } else { Get-ConfigDriverBaseName -RawName $StoredName }

    switch ($Make) {
        'Dell' {
            if ([string]::IsNullOrWhiteSpace($baseName)) { $baseName = $StoredName }
            if ([string]::IsNullOrWhiteSpace($SystemId)) { return $baseName }
            return "{0} ({1})" -f $baseName.Trim(), $SystemId.Trim()
        }
        'HP' {
            if ([string]::IsNullOrWhiteSpace($baseName)) { $baseName = $StoredName }
            if ([string]::IsNullOrWhiteSpace($SystemId)) { return $baseName }
            return "{0} ({1})" -f $baseName.Trim(), $SystemId.Trim()
        }
        'Lenovo' {
            if ([string]::IsNullOrWhiteSpace($baseName)) { $baseName = $StoredName }
            if ([string]::IsNullOrWhiteSpace($MachineType)) { return $baseName }
            return "{0} ({1})" -f $baseName.Trim(), $MachineType.Trim()
        }
        default {
            return $StoredName
        }
    }
}

function Invoke-LoadConfiguration {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )
    try {
        $filePath = Invoke-BrowseAction -Type 'OpenFile' -Filter "JSON files (*.json)|*.json|All files (*.*)|*.*" -Title "Load Configuration File"
        if (-not $filePath) {
            WriteLog "Load configuration cancelled by user."
            return
        }
        WriteLog "Loading configuration from: $filePath"
        $raw = $null
        try {
            $raw = Get-Content -Path $filePath -Raw -ErrorAction Stop
        }
        catch {
            WriteLog "LoadConfig Error: Failed reading file $filePath : $($_.Exception.Message)"
            [System.Windows.MessageBox]::Show("Failed to read the configuration file.`n$($_.Exception.Message)", "Load Error", "OK", "Error")
            return
        }
        if ([string]::IsNullOrWhiteSpace($raw)) {
            WriteLog "LoadConfig Error: File $filePath is empty."
            [System.Windows.MessageBox]::Show("The selected configuration file is empty.", "Load Error", "OK", "Error")
            return
        }
        $configContent = $null
        try {
            $configContent = $raw | ConvertFrom-Json -ErrorAction Stop
        }
        catch {
            WriteLog "LoadConfig Error: JSON parse failure for $filePath : $($_.Exception.Message)"
            [System.Windows.MessageBox]::Show("Failed to parse the configuration file (invalid JSON).`n$($_.Exception.Message)", "Load Error", "OK", "Error")
            return
        }
        if ($null -eq $configContent) {
            WriteLog "LoadConfig Error: Parsed config object is null after $filePath."
            [System.Windows.MessageBox]::Show("Parsed configuration object was null.", "Load Error", "OK", "Error")
            return
        }
        WriteLog "LoadConfig: Successfully parsed config file. Top-level keys: $($configContent.PSObject.Properties.Name -join ', ')"
        Update-UIFromConfig -ConfigContent $configContent -State $State
        $State.Data.lastConfigFilePath = $filePath
        Import-ConfigSupplementalAssets -ConfigContent $configContent -State $State -ShowWarnings:$true
    }
    catch {
        WriteLog "LoadConfig FATAL Error: $($_.Exception.ToString())"
        [System.Windows.MessageBox]::Show("Error loading config file:`n$($_.Exception.Message)", "Error", "OK", "Error")
    }
}

function Select-VMSwitchFromConfig {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State,
        [Parameter(Mandatory = $true)]
        [psobject]$ConfigContent
    )

    # Select VM switch based on configuration; fall back to 'Other' with custom name.
    $combo = $State.Controls.cmbVMSwitchName
    if ($null -eq $combo) {
        WriteLog "LoadConfig Error: 'cmbVMSwitchName' control not found."
        return
    }

    $configSwitch = $ConfigContent.VMSwitchName
    if ($null -eq $configSwitch -or [string]::IsNullOrWhiteSpace($configSwitch)) {
        WriteLog "LoadConfig Info: VMSwitchName in config was empty or null. Leaving selection unchanged."
        return
    }

    $itemFound = $false
    foreach ($item in $combo.Items) {
        if ($null -ne $item -and $item.ToString().Equals($configSwitch, [System.StringComparison]::OrdinalIgnoreCase)) {
            $itemFound = $true
            break
        }
    }

    if ($itemFound) {
        $combo.SelectedItem = ($combo.Items | Where-Object { $_.ToString().Equals($configSwitch, [System.StringComparison]::OrdinalIgnoreCase) } | Select-Object -First 1)
        $State.Controls.txtCustomVMSwitchName.Visibility = 'Collapsed'
        WriteLog "LoadConfig: Selected existing VM switch '$configSwitch'."
    }
    else {
        # Ensure 'Other' exists
        $otherExists = $false
        foreach ($item in $combo.Items) {
            if ($null -ne $item -and $item.ToString() -eq 'Other') { $otherExists = $true; break }
        }
        if (-not $otherExists) { $combo.Items.Add('Other') | Out-Null }

        # Select 'Other' and populate custom name
        $combo.SelectedItem = 'Other'
        $State.Controls.txtCustomVMSwitchName.Visibility = 'Visible'
        $State.Controls.txtCustomVMSwitchName.Text = $configSwitch
        $State.Data.customVMSwitchName = $configSwitch
        WriteLog "LoadConfig: VMSwitchName '$configSwitch' not found. Selected 'Other' and populated custom VM Switch Name textbox."
    }
}

function Update-UIFromConfig {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$ConfigContent,
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )

    WriteLog "Applying loaded configuration to the UI."

    # Apply theme mode from config (must be done before other controls load for proper styling)
    if ($null -ne $ConfigContent.PSObject.Properties.Item('ThemeMode') -and $State.Flags.isFluentSupported) {
        $configTheme = $ConfigContent.ThemeMode
        if ($configTheme -in @("Light", "Dark", "System")) {
            Initialize-FluentTheme -Window $State.Window -ThemeMode $configTheme -State $State
            Set-UIValue -ControlName 'cmbThemeMode' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'ThemeMode' -State $State
        }
    }

    # Update Build tab values
    Set-UIValue -ControlName 'txtFFUDevPath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'FFUDevelopmentPath' -State $State
    Set-UIValue -ControlName 'txtCustomFFUNameTemplate' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'CustomFFUNameTemplate' -State $State
    Set-UIValue -ControlName 'txtFFUCaptureLocation' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'FFUCaptureLocation' -State $State
    Set-UIValue -ControlName 'txtThreads' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'Threads' -State $State
    Set-UIValue -ControlName 'cmbBitsPriority' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'BitsPriority' -State $State
    Set-UIValue -ControlName 'txtMaxUSBDrives' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'MaxUSBDrives' -State $State
    Set-UIValue -ControlName 'chkBuildUSBDriveEnable' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'BuildUSBDrive' -State $State
    Set-UIValue -ControlName 'chkCompactOS' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CompactOS' -State $State
    Set-UIValue -ControlName 'chkUpdateADK' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'UpdateADK' -State $State
    Set-UIValue -ControlName 'chkOptimize' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'Optimize' -State $State
    Set-UIValue -ControlName 'chkAllowVHDXCaching' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'AllowVHDXCaching' -State $State
    Set-UIValue -ControlName 'chkAllowExternalHardDiskMedia' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'AllowExternalHardDiskMedia' -State $State
    Set-UIValue -ControlName 'chkPromptExternalHardDiskMedia' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'PromptExternalHardDiskMedia' -State $State
    Set-UIValue -ControlName 'chkCopyAdditionalFFUFiles' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CopyAdditionalFFUFiles' -State $State
    Set-UIValue -ControlName 'chkCreateDeploymentMedia' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CreateDeploymentMedia' -State $State
    Set-UIValue -ControlName 'chkInjectUnattend' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'InjectUnattend' -State $State
    Set-UIValue -ControlName 'txtUnattendX64FilePath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'UnattendX64FilePath' -State $State
    Set-UIValue -ControlName 'txtUnattendArm64FilePath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'UnattendArm64FilePath' -State $State
    Set-UIValue -ControlName 'chkVerbose' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'Verbose' -State $State

    if ([string]::IsNullOrWhiteSpace($State.Controls.txtUnattendX64FilePath.Text)) {
        $State.Controls.txtUnattendX64FilePath.Text = Get-DefaultUnattendFilePath -FFUDevelopmentPath $State.Controls.txtFFUDevPath.Text -WindowsArch 'x64'
    }

    if ([string]::IsNullOrWhiteSpace($State.Controls.txtUnattendArm64FilePath.Text)) {
        $State.Controls.txtUnattendArm64FilePath.Text = Get-DefaultUnattendFilePath -FFUDevelopmentPath $State.Controls.txtFFUDevPath.Text -WindowsArch 'arm64'
    }

    # USB Drive Modification group (Build Tab)
    Set-UIValue -ControlName 'chkCopyAutopilot' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CopyAutopilot' -State $State
    Set-UIValue -ControlName 'chkCopyUnattend' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CopyUnattend' -State $State
    Set-UIValue -ControlName 'chkCopyPPKG' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CopyPPKG' -State $State
    Set-UIValue -ControlName 'txtDeviceNamePrefixesPath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'DeviceNamePrefixesPath' -State $State
    Set-UIValue -ControlName 'txtDeviceNameSerialComputerNamesPath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'DeviceNameSerialComputerNamesPath' -State $State
    Set-UIValue -ControlName 'txtDeviceNameTemplate' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'DeviceNameTemplate' -State $State
    Set-UIValue -ControlName 'txtDeviceNamePrefixes' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'DeviceNamePrefixes' -TransformValue { param($val) if ($val -is [System.Array]) { $val -join [System.Environment]::NewLine } else { [string]$val } } -State $State
    Set-UIValue -ControlName 'txtDeviceNameSerialComputerNames' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'DeviceNameSerialComputerNames' -TransformValue { param($val) if ($val -is [System.Array]) { $val -join [System.Environment]::NewLine } else { [string]$val } } -State $State

    if ([string]::IsNullOrWhiteSpace($State.Controls.txtDeviceNamePrefixesPath.Text)) {
        $State.Controls.txtDeviceNamePrefixesPath.Text = Get-DefaultDeviceNamePrefixesPath -FFUDevelopmentPath $State.Controls.txtFFUDevPath.Text
    }

    if ([string]::IsNullOrWhiteSpace($State.Controls.txtDeviceNameSerialComputerNamesPath.Text)) {
        $State.Controls.txtDeviceNameSerialComputerNamesPath.Text = Get-DefaultSerialComputerNamesPath -FFUDevelopmentPath $State.Controls.txtFFUDevPath.Text
    }

    $loadedDeviceNamingMode = $null
    if ($ConfigContent.PSObject.Properties.Name -contains 'DeviceNamingMode') {
        $candidateDeviceNamingMode = [string]$ConfigContent.DeviceNamingMode
        if ($candidateDeviceNamingMode -in @('Legacy', 'None', 'Prompt', 'Template', 'Prefixes', 'SerialComputerNames')) {
            $loadedDeviceNamingMode = $candidateDeviceNamingMode
        }
    }
    $displayDeviceNamingMode = if ($loadedDeviceNamingMode -in @('Prompt', 'Template', 'Prefixes', 'SerialComputerNames')) {
        $loadedDeviceNamingMode
    }
    else {
        'None'
    }
    Set-DeviceNamingModeState -State $State -DisplayMode $displayDeviceNamingMode -LoadedMode $loadedDeviceNamingMode
    Import-DeviceNamePrefixesFromConfiguredPath -State $State
    Import-SerialComputerNamesFromConfiguredPath -State $State
    Update-DeviceNamingControls -State $State
    
    # Post Build Cleanup group (Build Tab)
    Set-UIValue -ControlName 'chkCleanupAppsISO' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CleanupAppsISO' -State $State
    Set-UIValue -ControlName 'chkCleanupDeployISO' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CleanupDeployISO' -State $State
    Set-UIValue -ControlName 'chkCleanupDrivers' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CleanupDrivers' -State $State
    Set-UIValue -ControlName 'chkRemoveFFU' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'RemoveFFU' -State $State
    Set-UIValue -ControlName 'chkRemoveApps' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'RemoveApps' -State $State
    Set-UIValue -ControlName 'chkRemoveUpdates' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'RemoveUpdates' -State $State
    Set-UIValue -ControlName 'chkRemoveDownloadedESD' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'RemoveDownloadedESD' -State $State

    # Hyper-V Settings
    Set-UIValue -ControlName 'chkEnableVMNetworking' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'EnableVMNetworking' -State $State
    Select-VMSwitchFromConfig -State $State -ConfigContent $ConfigContent
    Set-UIValue -ControlName 'txtDiskSize' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'Disksize' -TransformValue { param($val) $val / 1GB } -State $State
    Set-UIValue -ControlName 'txtMemory' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'Memory' -TransformValue { param($val) $val / 1GB } -State $State
    Set-UIValue -ControlName 'txtProcessors' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'Processors' -State $State
    Set-UIValue -ControlName 'txtVMLocation' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'VMLocation' -State $State
    Set-UIValue -ControlName 'txtVMNamePrefix' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'FFUPrefix' -State $State
    Set-UIValue -ControlName 'txtOSPartitionSizeGB' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'OSPartitionSize' -TransformValue { param($val) Get-PartitionSizeGBDisplay -SizeBytes $val } -State $State
    Set-UIValue -ControlName 'txtRecoveryPartitionSizeGB' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'RecoveryPartitionSize' -TransformValue { param($val) Get-PartitionSizeGBDisplay -SizeBytes $val } -State $State
    if ($ConfigContent.PSObject.Properties.Name -contains 'CreateRecoveryPartition') {
        $State.Data.createRecoveryPartition = [System.Convert]::ToBoolean($ConfigContent.CreateRecoveryPartition)
    }
    else {
        $State.Data.createRecoveryPartition = $true
    }
    Set-UIValue -ControlName 'cmbSystemPartitionDriveLetter' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'SystemPartitionDriveLetter' -TransformValue { param($val) ([string]$val).Trim().TrimEnd(':').ToUpperInvariant() } -State $State
    Set-UIValue -ControlName 'cmbWindowsPartitionDriveLetter' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'WindowsPartitionDriveLetter' -TransformValue { param($val) ([string]$val).Trim().TrimEnd(':').ToUpperInvariant() } -State $State
    Set-UIValue -ControlName 'cmbRecoveryPartitionDriveLetter' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'RecoveryPartitionDriveLetter' -TransformValue { param($val) ([string]$val).Trim().TrimEnd(':').ToUpperInvariant() } -State $State
    Import-AdditionalDataPartitionsFromConfig -State $State -ConfigContent $ConfigContent
    Set-UIValue -ControlName 'cmbLogicalSectorSize' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'LogicalSectorSizeBytes' -TransformValue { param($val) $val.ToString() } -State $State
    $State.Controls.spVMNetworkingSettings.IsEnabled = $true -eq $State.Controls.chkEnableVMNetworking.IsChecked
    if (-not ($true -eq $State.Controls.chkEnableVMNetworking.IsChecked)) {
        $State.Controls.txtCustomVMSwitchName.Visibility = 'Collapsed'
    }

    # Windows Settings
    Set-UIValue -ControlName 'txtISOPath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'ISOPath' -State $State
    # Load Windows Media Source setting
    if ($null -ne $ConfigContent.PSObject.Properties.Item('WindowsMediaSource')) {
        if ($ConfigContent.WindowsMediaSource -eq 'Provide Windows ISO') {
            $State.Controls.rbProvideISO.IsChecked = $true
        }
        else {
            $State.Controls.rbDownloadESD.IsChecked = $true
        }
    }
    
    # Special handling for Windows Release and SKU due to value collision (e.g., 2019 for Server and LTSC)
    if (($null -ne $ConfigContent.PSObject.Properties.Item('WindowsRelease')) -and ($null -ne $ConfigContent.PSObject.Properties.Item('WindowsSKU'))) {
        $configReleaseValue = $ConfigContent.WindowsRelease
        $configSkuValue = $ConfigContent.WindowsSKU
        WriteLog "LoadConfig: Handling Windows Release/SKU selection. Release: '$configReleaseValue', SKU: '$configSkuValue'."

        $releaseCombo = $State.Controls.cmbWindowsRelease
        # The items in the combobox are PSCustomObjects with Display and Value properties
        $possibleReleases = $releaseCombo.Items | Where-Object { $_.Value -eq $configReleaseValue }

        $releaseToSelect = $null
        if ($possibleReleases.Count -gt 1) {
            WriteLog "LoadConfig: Ambiguous release value '$configReleaseValue' found. Using SKU to disambiguate."
            if ($configSkuValue -like '*LTS*') {
                $releaseToSelect = $possibleReleases | Where-Object { $_.Display -like '*LTS*' } | Select-Object -First 1
                WriteLog "LoadConfig: SKU contains 'LTS'. Selecting LTSC-related release: '$($releaseToSelect.Display)'."
            }
            else {
                $releaseToSelect = $possibleReleases | Where-Object { $_.Display -notlike '*LTS*' } | Select-Object -First 1
                WriteLog "LoadConfig: SKU does not contain 'LTS'. Selecting non-LTSC (Server) release: '$($releaseToSelect.Display)'."
            }
        }
        else {
            $releaseToSelect = $possibleReleases | Select-Object -First 1
            if ($null -ne $releaseToSelect) {
                WriteLog "LoadConfig: Found unique release match: '$($releaseToSelect.Display)'."
            }
        }

        if ($null -ne $releaseToSelect) {
            $releaseCombo.SelectedItem = $releaseToSelect
        }
        else {
            WriteLog "LoadConfig: Could not determine a specific Windows Release to select for value '$configReleaseValue'. Skipping."
        }
    }
    else {
        # Fallback to individual setting if only one key exists
        WriteLog "LoadConfig: WindowsRelease or WindowsSKU key not found in config. Falling back to simple assignment for WindowsRelease."
        Set-UIValue -ControlName 'cmbWindowsRelease' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'WindowsRelease' -State $State
    }

    Set-UIValue -ControlName 'cmbWindowsVersion' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'WindowsVersion' -State $State
    Set-UIValue -ControlName 'cmbWindowsArch' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'WindowsArch' -State $State
    Set-UIValue -ControlName 'cmbWindowsLang' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'WindowsLang' -State $State
    Set-UIValue -ControlName 'cmbWindowsSKU' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'WindowsSKU' -State $State
    Set-UIValue -ControlName 'cmbMediaType' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'MediaType' -State $State
    Set-UIValue -ControlName 'txtProductKey' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'ProductKey' -State $State
   
    # Update Optional Features checkboxes
    $loadedFeaturesString = $ConfigContent.OptionalFeatures
    if (-not [string]::IsNullOrWhiteSpace($loadedFeaturesString)) {
        $loadedFeaturesArray = $loadedFeaturesString.Split(';')
        WriteLog "LoadConfig: Updating Optional Features checkboxes. Loaded features: $($loadedFeaturesArray -join ', ')"
        foreach ($featureEntry in $State.Controls.featureCheckBoxes.GetEnumerator()) {
            $featureName = $featureEntry.Key
            $featureCheckbox = $featureEntry.Value
            if ($loadedFeaturesArray -contains $featureName) {
                $featureCheckbox.IsChecked = $true
                WriteLog "LoadConfig: Checked checkbox for feature '$featureName'."
            }
            else {
                $featureCheckbox.IsChecked = $false
            }
        }
    }
    else {
        # If no optional features are loaded, uncheck all
        WriteLog "LoadConfig: No optional features string loaded. Unchecking all feature checkboxes."
        foreach ($featureEntry in $State.Controls.featureCheckBoxes.GetEnumerator()) {
            $featureEntry.Value.IsChecked = $false
        }
    }

    # M365 Apps/Office tab
    Set-UIValue -ControlName 'chkInstallOffice' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'InstallOffice' -State $State
    Set-UIValue -ControlName 'txtOfficePath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'OfficePath' -State $State
    Set-UIValue -ControlName 'chkCopyOfficeConfigXML' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CopyOfficeConfigXML' -State $State
    Set-UIValue -ControlName 'txtOfficeConfigXMLFilePath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'OfficeConfigXMLFile' -State $State

    # Drivers tab
    Set-UIValue -ControlName 'chkInstallDrivers' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'InstallDrivers' -State $State
    Set-UIValue -ControlName 'chkDownloadDrivers' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'DownloadDrivers' -State $State
    Set-UIValue -ControlName 'chkCopyDrivers' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CopyDrivers' -State $State
    # Set-UIValue -ControlName 'cmbMake' -PropertyName 'SelectedItem' -ConfigObject $ConfigContent -ConfigKey 'Make' -State $State
    Set-UIValue -ControlName 'txtDriversFolder' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'DriversFolder' -State $State
    Set-UIValue -ControlName 'txtPEDriversFolder' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'PEDriversFolder' -State $State
    Set-UIValue -ControlName 'txtDriversJsonPath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'DriversJsonPath' -State $State
    Set-UIValue -ControlName 'chkCopyPEDrivers' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CopyPEDrivers' -State $State
    Set-UIValue -ControlName 'chkUseDriversAsPEDrivers' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'UseDriversAsPEDrivers' -State $State
    Set-UIValue -ControlName 'chkCompressDriversToWIM' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'CompressDownloadedDriversToWim' -State $State

    # Updates tab
    Set-UIValue -ControlName 'chkUpdateLatestCU' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'UpdateLatestCU' -State $State
    Set-UIValue -ControlName 'chkUpdateLatestNet' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'UpdateLatestNet' -State $State
    Set-UIValue -ControlName 'chkUpdateLatestDefender' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'UpdateLatestDefender' -State $State
    Set-UIValue -ControlName 'chkUpdateEdge' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'UpdateEdge' -State $State
    Set-UIValue -ControlName 'chkUpdateOneDrive' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'UpdateOneDrive' -State $State
    Set-UIValue -ControlName 'chkUpdateLatestMSRT' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'UpdateLatestMSRT' -State $State
    Set-UIValue -ControlName 'chkUpdateLatestMicrocode' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'UpdateLatestMicrocode' -State $State
    Set-UIValue -ControlName 'chkUpdatePreviewCU' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'UpdatePreviewCU' -State $State

    # Applications tab
    Set-UIValue -ControlName 'chkInstallApps' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'InstallApps' -State $State
    Set-UIValue -ControlName 'chkInstallWingetApps' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'InstallWingetApps' -State $State
    Set-UIValue -ControlName 'chkBringYourOwnApps' -PropertyName 'IsChecked' -ConfigObject $ConfigContent -ConfigKey 'BringYourOwnApps' -State $State
    Set-UIValue -ControlName 'txtApplicationPath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'AppsPath' -State $State
    Set-UIValue -ControlName 'txtAppListJsonPath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'AppListPath' -State $State
    Set-UIValue -ControlName 'txtUserAppListPath' -PropertyName 'Text' -ConfigObject $ConfigContent -ConfigKey 'UserAppListPath' -State $State

    # Handle AppsScriptVariables
    $appsScriptVarsKeyExists = $false
    if ($ConfigContent -is [System.Management.Automation.PSCustomObject] -and $null -ne $ConfigContent.PSObject.Properties) {
        try {
            if (($ConfigContent.PSObject.Properties.Match('AppsScriptVariables')).Count -gt 0) {
                $appsScriptVarsKeyExists = $true
            }
        }
        catch { WriteLog "ERROR: Exception while trying to Match key 'AppsScriptVariables'. Error: $($_.Exception.Message)" }
    }

    $lstAppsScriptVars = $State.Controls.lstAppsScriptVariables
    $chkDefineAppsScriptVars = $State.Controls.chkDefineAppsScriptVariables
    $appsScriptVarsPanel = $State.Controls.appsScriptVariablesPanel
    $State.Data.appsScriptVariablesDataList.Clear()

    if ($appsScriptVarsKeyExists -and $null -ne $ConfigContent.AppsScriptVariables -and $ConfigContent.AppsScriptVariables -is [System.Management.Automation.PSCustomObject]) {
        WriteLog "LoadConfig: Processing AppsScriptVariables from config."
        $loadedVars = $ConfigContent.AppsScriptVariables
        $hasVars = $false
        foreach ($prop in $loadedVars.PSObject.Properties) {
            $State.Data.appsScriptVariablesDataList.Add([PSCustomObject]@{ IsSelected = $false; Key = $prop.Name; Value = $prop.Value })
            $hasVars = $true
        }
        if ($hasVars) {
            $chkDefineAppsScriptVars.IsChecked = $true
            $appsScriptVarsPanel.Visibility = 'Visible'
            WriteLog "LoadConfig: Loaded AppsScriptVariables and checked 'Define Apps Script Variables'."
        }
        else {
            $chkDefineAppsScriptVars.IsChecked = $false
            $appsScriptVarsPanel.Visibility = 'Collapsed'
            WriteLog "LoadConfig: AppsScriptVariables key was present but empty. Unchecked 'Define Apps Script Variables'."
        }
    }
    elseif ($appsScriptVarsKeyExists -and $null -ne $ConfigContent.AppsScriptVariables -and $ConfigContent.AppsScriptVariables -is [hashtable]) {
        # Handle if it's already a hashtable (e.g., from older config or direct creation)
        WriteLog "LoadConfig: Processing AppsScriptVariables (Hashtable) from config."
        $loadedVars = $ConfigContent.AppsScriptVariables
        $hasVars = $false
        foreach ($keyName in $loadedVars.Keys) {
            $State.Data.appsScriptVariablesDataList.Add([PSCustomObject]@{ IsSelected = $false; Key = $keyName; Value = $loadedVars[$keyName] })
            $hasVars = $true
        }
        if ($hasVars) {
            $chkDefineAppsScriptVars.IsChecked = $true
            $appsScriptVarsPanel.Visibility = 'Visible'
            WriteLog "LoadConfig: Loaded AppsScriptVariables (Hashtable) and checked 'Define Apps Script Variables'."
        }
        else {
            $chkDefineAppsScriptVars.IsChecked = $false
            $appsScriptVarsPanel.Visibility = 'Collapsed'
            WriteLog "LoadConfig: AppsScriptVariables (Hashtable) key was present but empty. Unchecked 'Define Apps Script Variables'."
        }
    }
    else {
        $chkDefineAppsScriptVars.IsChecked = $false
        $appsScriptVarsPanel.Visibility = 'Collapsed'
        WriteLog "LoadConfig Info: Key 'AppsScriptVariables' not found, is null, or not a PSCustomObject/Hashtable. Unchecked 'Define Apps Script Variables'."
    }
    # Update the ListView's ItemsSource after populating the data list
    $lstAppsScriptVars.ItemsSource = $State.Data.appsScriptVariablesDataList.ToArray()
    Request-ListViewColumnAutoResize -ListView $lstAppsScriptVars
    # Update the header checkbox state
    if ($null -ne $State.Controls.chkSelectAllAppsScriptVariables) {
        Update-SelectAllHeaderCheckBoxState -ListView $lstAppsScriptVars -HeaderCheckBox $State.Controls.chkSelectAllAppsScriptVariables
    }

    # Update USB Drive selection if present in config
    $usbDriveListKeyExists = $false
    if ($ConfigContent -is [System.Management.Automation.PSCustomObject] -and $null -ne $ConfigContent.PSObject.Properties) {
        try {
            if (($ConfigContent.PSObject.Properties.Match('USBDriveList')).Count -gt 0) {
                $usbDriveListKeyExists = $true
            }
        }
        catch {
            WriteLog "ERROR: Exception while trying to Match key 'USBDriveList' on ConfigContent.PSObject.Properties. Error: $($_.Exception.Message)"
        }
    }

    if ($usbDriveListKeyExists -and $null -ne $ConfigContent.USBDriveList) {
        WriteLog "LoadConfig: Processing USBDriveList from config."
        # First click the Check USB Drives button to populate the list
        $State.Controls.btnCheckUSBDrives.RaiseEvent(
            [System.Windows.RoutedEventArgs]::new(
                [System.Windows.Controls.Button]::ClickEvent
            )
        )

        # Then select the drives that match the saved configuration
        foreach ($item in $State.Controls.lstUSBDrives.Items) {
            $propertyName = $item.Model
            $propertyExists = $false
            $propertyValue = $null

            # Ensure USBDriveList is a PSCustomObject before trying to access its properties dynamically
            if ($null -ne $ConfigContent.USBDriveList -and $ConfigContent.USBDriveList -is [System.Management.Automation.PSCustomObject]) {
                # Check if the property exists on the USBDriveList object
                if ($ConfigContent.USBDriveList.PSObject.Properties.Match($propertyName).Count -gt 0) {
                    $propertyExists = $true
                    # Access the value dynamically
                    $propertyValue = $ConfigContent.USBDriveList.$($propertyName)
                }
            }

            # Match USB drives by UniqueId instead of SerialNumber
            # USBDriveList values can be a single UniqueId (string) or an array of UniqueIds (multiple same-model drives)
            $isMatch = $false
            if ($propertyExists) {
                if ($propertyValue -is [string]) {
                    $isMatch = ($propertyValue -eq $item.UniqueId)
                }
                else {
                    $propertyValueArray = @($propertyValue)
                    $isMatch = ($propertyValueArray -contains $item.UniqueId)
                }
            }

            if ($isMatch) {
                WriteLog "LoadConfig: Selecting USB Drive Model '$($item.Model)' with UniqueId '$($item.UniqueId)'."
                $item.IsSelected = $true
            }
            else {
                if (-not $propertyExists -and ($null -ne $ConfigContent.USBDriveList)) {
                    WriteLog "LoadConfig: Property '$($propertyName)' not found on USBDriveList for item Model '$($item.Model)'."
                }
                $item.IsSelected = $false # Ensure others are deselected if not in config or value mismatch
            }
        }
        $State.Controls.lstUSBDrives.Items.Refresh()
        Request-ListViewColumnAutoResize -ListView $State.Controls.lstUSBDrives

        # Update the Select All header checkbox state
        $headerChk = $State.Controls.chkSelectAllUSBDrivesHeader
        if ($null -ne $headerChk) {
            Update-SelectAllHeaderCheckBoxState -ListView $State.Controls.lstUSBDrives -HeaderCheckBox $headerChk
        }
        WriteLog "LoadConfig: USBDriveList processing complete."
    }
    else {
        WriteLog "LoadConfig Info: Key 'USBDriveList' not found or is null in configuration file. Skipping USB drive selection."
    }

    # If BuildUSBDrive is enabled and USBDriveList was present and not empty in the config,
    # ensure "Select Specific USB Drives" is checked to show the list.
    $shouldAutoCheckSpecificDrives = $false
    if ($State.Controls.chkBuildUSBDriveEnable.IsChecked -and $usbDriveListKeyExists -and ($null -ne $ConfigContent.USBDriveList)) {
        if ($ConfigContent.USBDriveList -is [System.Management.Automation.PSCustomObject]) {
            if ($ConfigContent.USBDriveList.PSObject.Properties.Count -gt 0) {
                $shouldAutoCheckSpecificDrives = $true
            }
        }
        elseif ($ConfigContent.USBDriveList -is [hashtable]) {
            # Fallback for older configs
            if ($ConfigContent.USBDriveList.Keys.Count -gt 0) {
                $shouldAutoCheckSpecificDrives = $true
            }
        }
    }

    if ($shouldAutoCheckSpecificDrives) {
        WriteLog "LoadConfig: Auto-checking 'Select Specific USB Drives' due to pre-selected USB drives in config."
        $State.Controls.chkSelectSpecificUSBDrives.IsChecked = $true
    }
    else {
        WriteLog "LoadConfig: Condition to auto-check 'Select Specific USB Drives' was NOT met."
    }
        # Populate additional FFU list and apply selections
        try {
            if ($State.Controls.chkCopyAdditionalFFUFiles.IsChecked) {
                $State.Controls.additionalFFUPanel.Visibility = 'Visible'
                if ($State.Controls.btnRefreshAdditionalFFUs) {
                    $State.Controls.btnRefreshAdditionalFFUs.RaiseEvent([System.Windows.RoutedEventArgs]::new([System.Windows.Controls.Button]::ClickEvent))
                }
                $selectedFiles = @()
                $addFFUKeyExists = $false
                if ($ConfigContent -is [System.Management.Automation.PSCustomObject] -and $null -ne $ConfigContent.PSObject.Properties) {
                    if (($ConfigContent.PSObject.Properties.Match('AdditionalFFUFiles')).Count -gt 0) {
                        $addFFUKeyExists = $true
                    }
                }
                if ($addFFUKeyExists -and $null -ne $ConfigContent.AdditionalFFUFiles) {
                    $selectedFiles = @($ConfigContent.AdditionalFFUFiles)
                }
                if ($selectedFiles.Count -gt 0) {
                    foreach ($item in $State.Controls.lstAdditionalFFUs.Items) {
                        if ($selectedFiles -contains $item.FullName) {
                            $item.IsSelected = $true
                        }
                    }
                    $State.Controls.lstAdditionalFFUs.Items.Refresh()
                    Request-ListViewColumnAutoResize -ListView $State.Controls.lstAdditionalFFUs
                    $headerChk = $State.Controls.chkSelectAllAdditionalFFUs
                    if ($null -ne $headerChk) {
                        Update-SelectAllHeaderCheckBoxState -ListView $State.Controls.lstAdditionalFFUs -HeaderCheckBox $headerChk
                    }
                }
            }
            else {
                $State.Controls.additionalFFUPanel.Visibility = 'Collapsed'
            }
            }
            catch {
                WriteLog "LoadConfig: Error applying Additional FFU selections: $($_.Exception.Message)"
            }
        
            Update-BitsPrioritySetting -State $State
            WriteLog "LoadConfig: Configuration loading process finished."
        }

function Invoke-SaveConfiguration {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )
    try {
        $config = Get-UIConfig -State $State
        $defaultConfigPath = Join-Path $config.FFUDevelopmentPath "config"
        if (-not (Test-Path $defaultConfigPath)) {
            New-Item -Path $defaultConfigPath -ItemType Directory -Force | Out-Null
        }
        
        $savePath = Invoke-BrowseAction -Type 'SaveFile' `
            -Title "Save Configuration File" `
            -Filter "JSON files (*.json)|*.json|All files (*.*)|*.*" `
            -InitialDirectory $defaultConfigPath `
            -FileName "FFUConfig.json" `
            -DefaultExt ".json"

        if ($savePath) {
            # Sort top-level keys alphabetically for consistent output
            $sortedConfig = [ordered]@{}
            foreach ($k in ($config.Keys | Sort-Object)) { $sortedConfig[$k] = $config[$k] }
            $sortedConfig | ConvertTo-Json -Depth 10 | Set-Content -Path $savePath -Encoding UTF8
            [System.Windows.MessageBox]::Show("Configuration file saved to:`n$savePath", "Success", "OK", "Information")
        }
    }
    catch {
        [System.Windows.MessageBox]::Show("Error saving config file:`n$($_.Exception.Message)", "Error", "OK", "Error")
    }
}

function Invoke-RestoreDefaults {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )
    try {
        $rootPath = $State.FFUDevelopmentPath

        # Normalize potential array values to single strings
        function Get-PathScalar {
            param([object]$value)
            if ($null -eq $value) { return $null }
            if ($value -is [System.Array]) {
                foreach ($v in $value) {
                    if (-not [string]::IsNullOrWhiteSpace([string]$v)) {
                        return [string]$v
                    }
                }
                return $null
            }
            return [string]$value
        }

        $appsPath = Join-Path $rootPath 'Apps'
        $driversRaw = Get-PathScalar -value $State.Controls.txtDriversFolder.Text
        if ([string]::IsNullOrWhiteSpace($driversRaw)) {
            $driversPath = Join-Path $rootPath 'Drivers'
        }
        else {
            $driversPath = $driversRaw
        }
        $ffuCaptureRaw = Get-PathScalar -value $State.Controls.txtFFUCaptureLocation.Text
        $ffuCapturePath = if ([string]::IsNullOrWhiteSpace($ffuCaptureRaw)) { Join-Path $rootPath 'FFU' } else { $ffuCaptureRaw }

        $deployISOPath = Join-Path $rootPath 'WinPEDeployFFUFiles\WinPE-Deploy.iso'
        $appsISOPath = Join-Path $rootPath 'Apps.iso'
        
        $msg = "Restore Defaults will:`n`n- Delete generated config and app/driver list JSON files`n- Remove ISO files (Deploy, Apps) if present`n- Remove Apps/Update/downloaded artifacts`n- Remove driver folder contents (not the folder)`n- Remove FFU files in the capture folder`n`nSample/template files and VM/VHDX cache are NOT removed.`n`nProceed?"
        $result = [System.Windows.MessageBox]::Show($msg, "Confirm Restore Defaults", "YesNo", "Warning")
        if ($result -ne [System.Windows.MessageBoxResult]::Yes) {
            WriteLog "RestoreDefaults: User cancelled."
            return
        }

        WriteLog "RestoreDefaults: Starting environment reset."
        WriteLog "RestoreDefaults: Paths -> Apps=$appsPath Drivers=$driversPath FFUCapture=$ffuCapturePath"

        # Remove JSON artifact files if present
        $artifactFiles = @(
            (Join-Path $rootPath 'config\FFUConfig.json'),
            (Join-Path $appsPath 'AppList.json'),
            (Join-Path $driversPath 'Drivers.json'),
            (Join-Path $appsPath 'UserAppList.json')
        ) | Where-Object { -not [string]::IsNullOrWhiteSpace($_) }

        foreach ($file in $artifactFiles) {
            if ((-not [string]::IsNullOrWhiteSpace($file)) -and (Test-Path -LiteralPath $file)) {
                try {
                    WriteLog "RestoreDefaults: Removing $file"
                    Remove-Item -LiteralPath $file -Force -ErrorAction Stop
                }
                catch {
                    WriteLog "RestoreDefaults: Failed removing $file : $($_.Exception.Message)"
                }
            }
        }

        # Force all cleanup flags true
        Invoke-FFUPostBuildCleanup `
            -RootPath $rootPath `
            -AppsPath $appsPath `
            -DriversPath $driversPath `
            -FFUCapturePath $ffuCapturePath `
            -DeployISOPath $deployISOPath `
            -AppsISOPath $appsISOPath `
            -KBPath (Join-Path $rootPath 'KB') `
            -RemoveDeployISO:$true `
            -RemoveAppsISO:$true `
            -RemoveDrivers:$true `
            -RemoveFFU:$true `
            -RemoveApps:$true `
            -RemoveUpdates:$true `
            -RemoveDownloadedESD:$true

        # Clear UI lists / state
        if ($null -ne $State.Data.allDriverModels) { $State.Data.allDriverModels.Clear() }
        if ($null -ne $State.Controls.lstDriverModels) { $State.Controls.lstDriverModels.Items.Refresh() }
        if ($null -ne $State.Controls.lstApplications) {
            try {
                if ($State.Controls.lstApplications.ItemsSource) { $State.Controls.lstApplications.ItemsSource = $null }
                $State.Controls.lstApplications.Items.Clear()
            } catch {}
        }
        if ($null -ne $State.Controls.lstWingetResults) {
            try { 
                if ($State.Controls.lstWingetResults.ItemsSource) { $State.Controls.lstWingetResults.ItemsSource = $null }
                $State.Controls.lstWingetResults.Items.Clear() 
            } catch {}
        }
        if ($null -ne $State.Controls.lstAppsScriptVariables) {
            try {
                if ($State.Controls.lstAppsScriptVariables.ItemsSource) { $State.Controls.lstAppsScriptVariables.ItemsSource = $null }
                $State.Controls.lstAppsScriptVariables.Items.Clear()
            } catch {}
        }

        $State.Data.lastConfigFilePath = $null

        Initialize-UIDefaults -State $State

        WriteLog "RestoreDefaults: Completed."
        [System.Windows.MessageBox]::Show("Environment restored to defaults.", "Restore Defaults", "OK", "Information")
    }
    catch {
        WriteLog "RestoreDefaults: Failed with $($_.Exception.Message)"
        [System.Windows.MessageBox]::Show("Restore Defaults failed:`n$($_.Exception.Message)", "Error", "OK", "Error")
    }
}

function Invoke-AutoLoadPreviousEnvironment {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$State
    )
    try {
        $ffuDevRoot = $State.FFUDevelopmentPath
        if ([string]::IsNullOrWhiteSpace($ffuDevRoot)) {
            WriteLog "AutoLoad: FFUDevelopmentPath not set; skipping."
            return
        }
        $configPath = Join-Path $ffuDevRoot "config\FFUConfig.json"
        if (-not (Test-Path -LiteralPath $configPath)) {
            WriteLog "AutoLoad: No existing FFUConfig.json found at $configPath."
            return
        }
        WriteLog "AutoLoad: Found config file at $configPath. Parsing..."
        $raw = Get-Content -Path $configPath -Raw -ErrorAction SilentlyContinue
        if ([string]::IsNullOrWhiteSpace($raw)) {
            WriteLog "AutoLoad: Config file empty; aborting."
            return
        }
        $configContent = $null
        try {
            $configContent = $raw | ConvertFrom-Json -ErrorAction Stop
        }
        catch {
            WriteLog "AutoLoad: JSON parse failed: $($_.Exception.Message)"
            return
        }
        if ($null -eq $configContent) {
            WriteLog "AutoLoad: Parsed object null; aborting."
            return
        }
        WriteLog "AutoLoad: Applying core configuration."
        Update-UIFromConfig -ConfigContent $configContent -State $State
        $State.Data.lastConfigFilePath = $configPath
        Import-ConfigSupplementalAssets -ConfigContent $configContent -State $State -ShowWarnings:$false
        WriteLog "AutoLoad: Completed supplemental import with warnings disabled."
    }
    catch {
        WriteLog "AutoLoad: Unexpected failure: $($_.Exception.ToString())"
    }
}

function Import-ConfigSupplementalAssets {
    param(
        [Parameter(Mandatory = $true)]
        [psobject]$ConfigContent,
        [Parameter(Mandatory = $true)]
        [psobject]$State,
        [Parameter()]
        [bool]$ShowWarnings = $false
    )
    WriteLog "SupplementalImport: Starting import of helper assets."
    $loadedWinget = $false
    $loadedBYO = $false
    $loadedDrivers = $false
    $missing = New-Object System.Collections.Generic.List[string]

    # Winget AppList
    $appListPath = $null
    if ($ConfigContent.PSObject.Properties.Match('AppListPath').Count -gt 0) {
        $appListPath = $ConfigContent.AppListPath
    }
    if (-not [string]::IsNullOrWhiteSpace($appListPath)) {
        if (Test-Path -LiteralPath $appListPath) {
            WriteLog "SupplementalImport: Loading Winget AppList from $appListPath"
            try {
                $importedAppsData = Get-Content -Path $appListPath -Raw | ConvertFrom-Json -ErrorAction Stop
                if ($null -ne $importedAppsData -and $null -ne $importedAppsData.apps) {
                    $defaultArch = $State.Controls.cmbWindowsArch.SelectedItem
                    $appsBuffer = [System.Collections.Generic.List[object]]::new()
                    foreach ($appInfo in $importedAppsData.apps) {
                        $arch = if ($appInfo.source -eq 'msstore') { 'NA' } else {
                            if ($appInfo.PSObject.Properties['architecture']) { $appInfo.architecture } else { $defaultArch }
                        }
                        $appsBuffer.Add([PSCustomObject]@{
                                IsSelected               = $true
                                Name                     = $appInfo.name
                                Id                       = $appInfo.id
                                Version                  = ""
                                Source                   = $appInfo.source
                                Architecture             = $arch
                                AdditionalExitCodes      = if ($appInfo.PSObject.Properties['AdditionalExitCodes']) { $appInfo.AdditionalExitCodes } else { "" }
                                IgnoreNonZeroExitCodes   = if ($appInfo.PSObject.Properties['IgnoreNonZeroExitCodes']) { [bool]$appInfo.IgnoreNonZeroExitCodes } else { $false }
                                DownloadStatus           = ""
                            })
                    }
                    $State.Controls.lstWingetResults.ItemsSource = $appsBuffer.ToArray()
                    Request-ListViewColumnAutoResize -ListView $State.Controls.lstWingetResults
                    $loadedWinget = $true
                    if ($null -ne $State.Controls.wingetSearchPanel) {
                        $State.Controls.wingetSearchPanel.Visibility = 'Visible'
                    }
                    if ($null -ne $State.Controls.chkSelectAllWingetResults -and (Get-Command -Name Update-SelectAllHeaderCheckBoxState -ErrorAction SilentlyContinue)) {
                        Update-SelectAllHeaderCheckBoxState -ListView $State.Controls.lstWingetResults -HeaderCheckBox $State.Controls.chkSelectAllWingetResults
                    }
                    WriteLog "SupplementalImport: Winget list loaded with $($appsBuffer.Count) entries."
                }
                else {
                    WriteLog "SupplementalImport: Winget AppList missing 'apps' array."
                }
            }
            catch {
                WriteLog "SupplementalImport: Failed loading Winget AppList ($appListPath): $($_.Exception.Message)"
            }
        }
        else {
            WriteLog "SupplementalImport: Winget AppList file missing: $appListPath"
            $missing.Add("Winget AppList (AppListPath): $appListPath")
        }
    }
    else {
        WriteLog "SupplementalImport: AppListPath not defined in config."
    }

    # UserAppList (BYO)
    $userAppListPath = $null
    if ($ConfigContent.PSObject.Properties.Match('UserAppListPath').Count -gt 0) {
        $userAppListPath = $ConfigContent.UserAppListPath
    }
    if (-not [string]::IsNullOrWhiteSpace($userAppListPath)) {
        if (Test-Path -LiteralPath $userAppListPath) {
            WriteLog "SupplementalImport: Loading UserAppList from $userAppListPath"
            try {
                $applications = Get-Content -Path $userAppListPath -Raw | ConvertFrom-Json -ErrorAction Stop
                if ($applications) {
                    $listView = $State.Controls.lstApplications
                    $listView.Items.Clear()
                    $sortedApps = $applications | Sort-Object Priority
                    foreach ($app in $sortedApps) {
                        $ignoreNonZero = if ($app.PSObject.Properties['IgnoreNonZeroExitCodes']) { $app.IgnoreNonZeroExitCodes } else { $false }
                        $listView.Items.Add([PSCustomObject]@{
                                IsSelected             = $false
                                Priority               = $app.Priority
                                Name                   = $app.Name
                                CommandLine            = $app.CommandLine
                                Arguments              = if ($app.PSObject.Properties['Arguments']) { $app.Arguments } else { "" }
                                Source                 = $app.Source
                                AdditionalExitCodes    = if ($app.PSObject.Properties['AdditionalExitCodes']) { $app.AdditionalExitCodes } else { "" }
                                IgnoreNonZeroExitCodes = $ignoreNonZero
                                IgnoreExitCodes        = if ($ignoreNonZero) { "Yes" } else { "No" }
                                CopyStatus             = ""
                            })
                    }
                    if (Get-Command -Name Update-ListViewPriorities -ErrorAction SilentlyContinue) {
                        Update-ListViewPriorities -ListView $listView
                    }
                    if (Get-Command -Name Update-CopyButtonState -ErrorAction SilentlyContinue) {
                        Update-CopyButtonState -State $State
                    }
                    if (Get-Command -Name Update-BYOAppsActionButtonsState -ErrorAction SilentlyContinue) {
                        Update-BYOAppsActionButtonsState -State $State
                    }
                    $loadedBYO = $true
                    WriteLog "SupplementalImport: UserAppList loaded with $($listView.Items.Count) entries."
                }
                else {
                    WriteLog "SupplementalImport: UserAppList JSON empty."
                }
            }
            catch {
                WriteLog "SupplementalImport: Failed loading UserAppList ($userAppListPath): $($_.Exception.Message)"
            }
        }
        else {
            WriteLog "SupplementalImport: UserAppList file missing: $userAppListPath"
            $missing.Add("UserAppList (UserAppListPath): $userAppListPath")
        }
    }
    else {
        WriteLog "SupplementalImport: UserAppListPath not defined in config."
    }

    # Drivers JSON
    $driversJsonPath = $null
    if ($ConfigContent.PSObject.Properties.Match('DriversJsonPath').Count -gt 0) {
        $driversJsonPath = $ConfigContent.DriversJsonPath
    }
    if (-not [string]::IsNullOrWhiteSpace($driversJsonPath)) {
        if (Test-Path -LiteralPath $driversJsonPath) {
            WriteLog "SupplementalImport: Loading Drivers JSON from $driversJsonPath"
            try {
                $rawDrivers = Get-Content -Path $driversJsonPath -Raw | ConvertFrom-Json -ErrorAction Stop
                if ($rawDrivers -and $rawDrivers.PSObject.Properties.Count -gt 0) {
                    $State.Data.allDriverModels.Clear()
                    foreach ($makeProp in $rawDrivers.PSObject.Properties) {
                        $makeName = $makeProp.Name
                        $makeObject = $makeProp.Value
                        if ($null -eq $makeObject -or -not ($makeObject.PSObject.Properties['Models'])) { continue }
                        $models = $makeObject.Models
                        if ($models -and ($models -is [System.Collections.IEnumerable])) {
                            foreach ($modelEntry in $models) {
                                if ($null -eq $modelEntry -or -not ($modelEntry.PSObject.Properties['Name'])) { continue }
                                $modelName = $modelEntry.Name
                                if ([string]::IsNullOrWhiteSpace($modelName)) { continue }
                                $downloadStatus = if ($modelEntry.PSObject.Properties['DownloadStatus']) { $modelEntry.DownloadStatus } else { "" }
                                $linkValue = if ($modelEntry.PSObject.Properties['Link']) { $modelEntry.Link } else { $null }
                                $productName = if ($modelEntry.PSObject.Properties['ProductName']) { $modelEntry.ProductName } else { $null }
                                $machineType = if ($modelEntry.PSObject.Properties['MachineType']) { $modelEntry.MachineType } else { $null }
                                $systemId = if ($modelEntry.PSObject.Properties['SystemId']) { $modelEntry.SystemId } else { $null }
                                $idValue = if ($modelEntry.PSObject.Properties['Id']) { $modelEntry.Id } else { $null }
                                if ($null -eq $idValue -and -not [string]::IsNullOrWhiteSpace($systemId)) { $idValue = $systemId }
                                if ($null -eq $idValue -and -not [string]::IsNullOrWhiteSpace($machineType)) { $idValue = $machineType }
                                $displayModel = Get-ConfigDriverDisplayName -Make $makeName -StoredName $modelName -ProductName $productName -SystemId $systemId -MachineType $machineType
                                if ([string]::IsNullOrWhiteSpace($displayModel)) {
                                    $displayModel = $modelName
                                }
                                $driverObj = [PSCustomObject]@{
                                    IsSelected     = $true
                                    Make           = $makeName
                                    Model          = $displayModel
                                    DownloadStatus = $downloadStatus
                                    Link           = $linkValue
                                    ProductName    = $productName
                                    MachineType    = $machineType
                                    SystemId       = $systemId
                                    Id             = $idValue
                                }
                                $State.Data.allDriverModels.Add($driverObj)
                            }
                        }
                    }
                    $State.Controls.lstDriverModels.ItemsSource = $State.Data.allDriverModels
                    Request-ListViewColumnAutoResize -ListView $State.Controls.lstDriverModels
                    if (Get-Command -Name Update-SelectAllHeaderCheckBoxState -ErrorAction SilentlyContinue) {
                        $headerChk = $State.Controls.chkSelectAllDriverModels
                        if ($null -ne $headerChk) {
                            Update-SelectAllHeaderCheckBoxState -ListView $State.Controls.lstDriverModels -HeaderCheckBox $headerChk
                        }
                    }
                    if ($State.Data.allDriverModels.Count -gt 0) {
                        if ($null -ne $State.Controls.spModelFilterSection) { $State.Controls.spModelFilterSection.Visibility = 'Visible' }
                        if ($null -ne $State.Controls.lstDriverModels) { $State.Controls.lstDriverModels.Visibility = 'Visible' }
                        if ($null -ne $State.Controls.spDriverActionButtons) { $State.Controls.spDriverActionButtons.Visibility = 'Visible' }
                        try {
                            if ($State.Controls.cmbMake.SelectedIndex -lt 0 -and $State.Data.allDriverModels.Count -gt 0) {
                                $firstMake = ($State.Data.allDriverModels | Select-Object -First 1).Make
                                if (-not [string]::IsNullOrWhiteSpace($firstMake)) {
                                    $makeItem = $State.Controls.cmbMake.Items | Where-Object { $_ -eq $firstMake } | Select-Object -First 1
                                    if ($makeItem) { $State.Controls.cmbMake.SelectedItem = $makeItem }
                                }
                            }
                        }
                        catch {
                            WriteLog "SupplementalImport: Non-fatal error selecting first Make: $($_.Exception.Message)"
                        }
                    }
                    $loadedDrivers = $true
                    WriteLog "SupplementalImport: Loaded $($State.Data.allDriverModels.Count) driver models."
                }
                else {
                    WriteLog "SupplementalImport: Drivers JSON empty or structure unexpected."
                }
            }
            catch {
                WriteLog "SupplementalImport: Failed loading Drivers JSON ($driversJsonPath): $($_.Exception.Message)"
            }
        }
        else {
            WriteLog "SupplementalImport: Drivers JSON file missing: $driversJsonPath"
            $missing.Add("Drivers (DriversJsonPath): $driversJsonPath")
        }
    }
    else {
        WriteLog "SupplementalImport: DriversJsonPath not defined in config."
    }

    if ($loadedWinget -or $loadedBYO) {
        $State.Controls.chkInstallApps.IsChecked = $true
    }
    if ($loadedWinget) {
        $State.Controls.chkInstallWingetApps.IsChecked = $true
    }
    if ($loadedBYO) {
        $State.Controls.chkBringYourOwnApps.IsChecked = $true
    }
    if ($loadedDrivers) {
        $State.Controls.chkDownloadDrivers.IsChecked = $true
    }

    if (Get-Command -Name Update-ApplicationPanelVisibility -ErrorAction SilentlyContinue) {
        Update-ApplicationPanelVisibility -State $State -TriggeringControlName 'SupplementalImport'
    }
    if (Get-Command -Name Update-DriverDownloadPanelVisibility -ErrorAction SilentlyContinue) {
        Update-DriverDownloadPanelVisibility -State $State
    }
    if (Get-Command -Name Update-DriverCheckboxStates -ErrorAction SilentlyContinue) {
        Update-DriverCheckboxStates -State $State
    }
    if (Get-Command -Name Update-OfficePanelVisibility -ErrorAction SilentlyContinue) {
        Update-OfficePanelVisibility -State $State
    }
    if (Get-Command -Name Update-CopyButtonState -ErrorAction SilentlyContinue) {
        Update-CopyButtonState -State $State
    }

    # Updated message to clarify successful load and that missing helper files are optional if not yet created.
    if ($ShowWarnings -and $missing.Count -gt 0) {
        $msg = "Configuration file loaded successfully.`n`n" +
            "Optional helper file(s) referenced in the configuration were not found:`n" +
            ($missing | ForEach-Object { "- $_" } | Out-String) +
            "`nThese files are optional. They won't exist until you create Winget (AppList.json), User (UserAppList.json), or Driver (Drivers.json) manifests. You can create them later or ignore this message."
        [System.Windows.MessageBox]::Show($msg.TrimEnd(), "Configuration Loaded - Optional Files Missing", "OK", "Information") | Out-Null
    }

    WriteLog ("SupplementalImport: Complete. Winget={0} BYO={1} Drivers={2} Missing={3}" -f $loadedWinget, $loadedBYO, $loadedDrivers, $missing.Count)
}

Export-ModuleMember -Function *