MDT新脚本 - 丿记性不太好丶

选硬盘
```
<#
.SYNOPSIS
MDT Disk Selection Script - Fix multi-disk error, use OSDDiskIndex variable
#>

# Ensure STA mode (Required for WPF)
if ([System.Threading.Thread]::CurrentThread.GetApartmentState() -ne "STA") {
    [System.Threading.Thread]::CurrentThread.SetApartmentState("STA")
}

# Load WPF Assemblies
Add-Type -AssemblyName PresentationFramework
Add-Type -AssemblyName System.Drawing

# ==============================================================================
# 1. Create WPF Window and Controls (Keep original)
# ==============================================================================
$window = New-Object System.Windows.Window
$window.Title = "Select Target Disk for OS Installation"
$window.Width = 800
$window.Height = 450
$window.WindowStartupLocation = "CenterScreen"
$window.ResizeMode = "NoResize"
$window.Topmost = $true
$window.Background = [System.Windows.Media.Brushes]::White
$window.WindowStyle = [System.Windows.WindowStyle]::None

$mainGrid = New-Object System.Windows.Controls.Grid
$mainGrid.Margin = "20"
$mainGrid.RowDefinitions.Add((New-Object System.Windows.Controls.RowDefinition -Property @{ Height = [System.Windows.GridLength]::Auto }))
$mainGrid.RowDefinitions.Add((New-Object System.Windows.Controls.RowDefinition -Property @{ Height = [System.Windows.GridLength]"*" }))
$mainGrid.RowDefinitions.Add((New-Object System.Windows.Controls.RowDefinition -Property @{ Height = [System.Windows.GridLength]::Auto }))
$mainGrid.ColumnDefinitions.Add((New-Object System.Windows.Controls.ColumnDefinition -Property @{ Width = [System.Windows.GridLength]"*" }))
$mainGrid.ColumnDefinitions.Add((New-Object System.Windows.Controls.ColumnDefinition -Property @{ Width = [System.Windows.GridLength]::Auto }))
$mainGrid.ColumnDefinitions.Add((New-Object System.Windows.Controls.ColumnDefinition -Property @{ Width = [System.Windows.GridLength]::Auto }))
$mainGrid.ColumnDefinitions.Add((New-Object System.Windows.Controls.ColumnDefinition -Property @{ Width = [System.Windows.GridLength]"*" }))
$window.Content = $mainGrid

# Title
$titleLabel = New-Object System.Windows.Controls.Label
$titleLabel.Content = "Please select the disk to install the operating system:"
$titleLabel.FontSize = 16
$titleLabel.FontWeight = "Bold"
$titleLabel.Margin = "0,0,0,15"
[System.Windows.Controls.Grid]::SetRow($titleLabel, 0)
[System.Windows.Controls.Grid]::SetColumnSpan($titleLabel, 4)
$mainGrid.Children.Add($titleLabel)

# Disk List Box
$diskListBox = New-Object System.Windows.Controls.ListBox
$diskListBox.Width = 750
$diskListBox.Height = 280
$diskListBox.Margin = "0,0,0,15"
$diskListBox.SetValue([System.Windows.Controls.ScrollViewer]::VerticalScrollBarVisibilityProperty, [System.Windows.Controls.ScrollBarVisibility]::Auto)
$diskListBox.FontSize = 13
[System.Windows.Controls.Grid]::SetRow($diskListBox, 1)
[System.Windows.Controls.Grid]::SetColumnSpan($diskListBox, 4)
$mainGrid.Children.Add($diskListBox)

# Buttons
$confirmButton = New-Object System.Windows.Controls.Button
$confirmButton.Content = "Confirm"
$confirmButton.FontSize = 14
$confirmButton.Width = 180
$confirmButton.Height = 40
$confirmButton.Margin = "0,0,10,0"
$confirmButton.IsDefault = $true
[System.Windows.Controls.Grid]::SetRow($confirmButton, 2)
[System.Windows.Controls.Grid]::SetColumn($confirmButton, 1)
$mainGrid.Children.Add($confirmButton)

$exitButton = New-Object System.Windows.Controls.Button
$exitButton.Content = "Exit MDT"
$exitButton.FontSize = 14
$exitButton.Width = 180
$exitButton.Height = 40
$exitButton.Margin = "10,0,0,0"
[System.Windows.Controls.Grid]::SetRow($exitButton, 2)
[System.Windows.Controls.Grid]::SetColumn($exitButton, 2)
$mainGrid.Children.Add($exitButton)

# ==============================================================================
# 2. Core Fix: Multi-disk support (Store full disk objects to avoid index mismatch)
# ==============================================================================
$Global:DiskObjects = @()  # Store full disk objects instead of only numbers

# Force WPF Message Pump Initialization
$dispatcherFrame = New-Object System.Windows.Threading.DispatcherFrame
[System.Windows.Threading.Dispatcher]::CurrentDispatcher.BeginInvoke([System.Windows.Threading.DispatcherPriority]::Normal, ([Action]{
    $dispatcherFrame.Continue = $false
}))
[System.Windows.Threading.Dispatcher]::PushFrame($dispatcherFrame)

if (Get-Module -ListAvailable -Name Storage) {
    try {
        # Get disks and filter (exclude USB + Online + Has size)
        $allDisks = Get-Disk -ErrorAction Stop | Where-Object { 
            $_.OperationalStatus -eq 'Online' -and $_.Size -gt 0 -and $_.BusType -ne 'USB' 
        } | Sort-Object Number  # Sort by number to ensure order consistency
        
        if ($allDisks.Count -eq 0) { throw "No physical disks found." }

# Populate listbox + Store full disk objects
        foreach ($disk in $allDisks) {
            $diskSizeGB = [math]::Round($disk.Size / 1GB, 2)
            $displayText = "Disk $($disk.Number): $($disk.Model) | Size: $diskSizeGB GB | Style: $($disk.PartitionStyle)"
            
            $listItem = New-Object System.Windows.Controls.ListBoxItem
            $listItem.Content = $displayText
            $listItem.Tag = $disk  # Store full disk object to avoid index mismatch
            $diskListBox.Items.Add($listItem)
            $Global:DiskObjects += $disk  # Sync storage for double protection
        }

# Auto-select first disk
        if ($diskListBox.Items.Count -gt 0) { $diskListBox.SelectedIndex = 0 }

} catch {
        $titleLabel.Content = "ERROR: $_"
        $titleLabel.Foreground = [System.Windows.Media.Brushes]::Red
        $confirmButton.IsEnabled = $false
    }
} else {
    $titleLabel.Content = "ERROR: Storage module not available."
    $titleLabel.Foreground = [System.Windows.Media.Brushes]::Red
    $confirmButton.IsEnabled = $false
}

# ==============================================================================
# 3. Button Events (Use OSDDiskIndex variable, MDT native recognition)
# ==============================================================================
$confirmButton.Add_Click({
    $selectedItem = $diskListBox.SelectedItem
    if ($selectedItem -and $selectedItem.Tag) {
        $selectedDisk = $selectedItem.Tag  # Get disk object directly from Tag, no index mismatch
        $selectedDiskNumber = $selectedDisk.Number
        $selectedDiskModel = $selectedDisk.Model

# Confirmation dialog
        $message = "You are about to format Disk $selectedDiskNumber ($selectedDiskModel).`n`nWARNING: All data will be erased.`n`nAre you sure?"
        $result = [System.Windows.MessageBox]::Show($message, "Confirm Disk Format", [System.Windows.MessageBoxButton]::YesNo, [System.Windows.MessageBoxImage]::Warning)
        
        if ($result -eq [System.Windows.MessageBoxResult]::Yes) {
            try {
                $tsEnv = New-Object -ComObject Microsoft.SMS.TSEnvironment -ErrorAction Stop
                $tsEnv.Value("OSDDiskIndex") = $selectedDiskNumber  # Key: Use OSDDiskIndex variable
                Write-Host "OSDDiskIndex set to $selectedDiskNumber"
                $window.Close()
            } catch {
                [System.Windows.MessageBox]::Show("Failed to set MDT variable: $_", "Error", [System.Windows.MessageBoxButton]::OK, [System.Windows.MessageBoxImage]::Error)
            }
        }
    } else {
        [System.Windows.MessageBox]::Show("Please select a disk first.", "Selection Required", [System.Windows.MessageBoxButton]::OK, [System.Windows.MessageBoxImage]::Information)
    }
})

$exitButton.Add_Click({
    # Enhanced exit prompt
    $message = "Are you sure you want to exit the MDT deployment wizard?`n`nThis will cancel the current OS installation process."
    $result = [System.Windows.MessageBox]::Show($message, "Exit Deployment", [System.Windows.MessageBoxButton]::YesNo, [System.Windows.MessageBoxImage]::Question)
    if ($result -eq [System.Windows.MessageBoxResult]::Yes) {
        try {
            $tsEnv = New-Object -ComObject Microsoft.SMS.TSEnvironment
            $tsEnv.Value("TSExitCode") = 1
            [System.Windows.MessageBox]::Show("Deployment cancelled successfully.", "Process Terminated", [System.Windows.MessageBoxButton]::OK, [System.Windows.MessageBoxImage]::Information)
            [Environment]::Exit(1)
        } catch {
            [System.Windows.MessageBox]::Show("Failed to exit gracefully: $_", "Exit Error", [System.Windows.MessageBoxButton]::OK, [System.Windows.MessageBoxImage]::Error)
        }
    }
})

# Display Window
$window.ShowDialog() | Out-Null

# Cleanup
if ($tsEnv) { [System.Runtime.Interopservices.Marshal]::ReleaseComObject($tsEnv) | Out-Null }
```

MDT选分区，不过滤U盘，都显示，测试通过
```
<#
.SYNOPSIS
MDT Partition Selection Script.
#>

# Ensure STA mode (Required for WPF)
if ([System.Threading.Thread]::CurrentThread.GetApartmentState() -ne "STA") {
    [System.Threading.Thread]::CurrentThread.SetApartmentState("STA")
}

# Load WPF Assemblies
Add-Type -AssemblyName PresentationFramework
Add-Type -AssemblyName System.Drawing

# ==============================================================================
# 1. Create WPF Window and Controls
# ==============================================================================
$window = New-Object System.Windows.Window
$window.Title = "Select Installation Partition"
$window.Width = 700
$window.Height = 500
$window.WindowStartupLocation = "CenterScreen"
$window.ResizeMode = "NoResize"
$window.Topmost = $true
$window.Background = [System.Windows.Media.Brushes]::White

$mainGrid.ColumnDefinitions.Add((New-Object System.Windows.Controls.ColumnDefinition -Property @{ Width = [System.Windows.GridLength]"*" }))
$mainGrid.ColumnDefinitions.Add((New-Object System.Windows.Controls.ColumnDefinition -Property @{ Width = [System.Windows.GridLength]::Auto }))
$mainGrid.ColumnDefinitions.Add((New-Object System.Windows.Controls.ColumnDefinition -Property @{ Width = [System.Windows.GridLength]::Auto }))
$mainGrid.ColumnDefinitions.Add((New-Object System.Windows.Controls.ColumnDefinition -Property @{ Width = [System.Windows.GridLength]"*" }))

$window.Content = $mainGrid
$window.WindowStyle = [System.Windows.WindowStyle]::None
# Title
$titleLabel = New-Object System.Windows.Controls.Label
$titleLabel.Content = "Please select a partition for OS installation:"
$titleLabel.FontSize = 16
$titleLabel.FontWeight = "Bold"
$titleLabel.Margin = "0,0,0,10"
[System.Windows.Controls.Grid]::SetRow($titleLabel, 0)
[System.Windows.Controls.Grid]::SetColumnSpan($titleLabel, 4)
$mainGrid.Children.Add($titleLabel)

# Partition List DataGrid
$dataGrid = New-Object System.Windows.Controls.DataGrid
$dataGrid.SelectionMode = "Single"
$dataGrid.SelectionUnit = "FullRow"
$dataGrid.AutoGenerateColumns = $false
$dataGrid.CanUserAddRows = $false
$dataGrid.CanUserDeleteRows = $false
$dataGrid.IsReadOnly = $true
$dataGrid.Margin = "0,10,0,10"
$dataGrid.VerticalScrollBarVisibility = "Auto"
[System.Windows.Controls.Grid]::SetRow($dataGrid, 1)
[System.Windows.Controls.Grid]::SetColumnSpan($dataGrid, 4)
$mainGrid.Children.Add($dataGrid)

# Define DataGrid Columns
$columnDriveLetter = New-Object System.Windows.Controls.DataGridTextColumn
$columnDriveLetter.Header = "Drive"
$columnDriveLetter.Binding = New-Object System.Windows.Data.Binding("DriveLetter")
$columnDriveLetter.Width = 60
$dataGrid.Columns.Add($columnDriveLetter)

$columnDiskNumber = New-Object System.Windows.Controls.DataGridTextColumn
$columnDiskNumber.Header = "Disk"
$columnDiskNumber.Binding = New-Object System.Windows.Data.Binding("DiskNumber")
$columnDiskNumber.Width = 60
$dataGrid.Columns.Add($columnDiskNumber)

$columnSize = New-Object System.Windows.Controls.DataGridTextColumn
$columnSize.Header = "Size (GB)"
$columnSize.Binding = New-Object System.Windows.Data.Binding("SizeGB")
$columnSize.Width = 80
$dataGrid.Columns.Add($columnSize)

$columnFileSystem = New-Object System.Windows.Controls.DataGridTextColumn
$columnFileSystem.Header = "File Sys"
$columnFileSystem.Binding = New-Object System.Windows.Data.Binding("FileSystem")
$columnFileSystem.Width = 80
$dataGrid.Columns.Add($columnFileSystem)

$columnLabel = New-Object System.Windows.Controls.DataGridTextColumn
$columnLabel.Header = "Label"
$columnLabel.Binding = New-Object System.Windows.Data.Binding("Label")
$columnLabel.Width = "*"
$dataGrid.Columns.Add($columnLabel)

$columnhealthstatus = New-Object System.Windows.Controls.DataGridTextColumn
$columnhealthstatus.Header = "HealthStatus"
$columnhealthstatus.Binding = New-Object System.Windows.Data.Binding("HealthStatus")
$columnhealthstatus.Width = "130"
$dataGrid.Columns.Add($columnhealthstatus)

# Confirm Button
$confirmButton = New-Object System.Windows.Controls.Button
$confirmButton.Content = "Confirm"
$confirmButton.FontSize = 14
$confirmButton.Width = 200
$confirmButton.Height = 40
$confirmButton.Margin = "10,0,10,0"
$confirmButton.IsDefault = $true
[System.Windows.Controls.Grid]::SetRow($confirmButton, 2)
[System.Windows.Controls.Grid]::SetColumn($confirmButton, 1)
$mainGrid.Children.Add($confirmButton)

# Exit Button
$exitButton = New-Object System.Windows.Controls.Button
$exitButton.Content = "Exit MDT"
$exitButton.FontSize = 14
$exitButton.Width = 200
$exitButton.Height = 40
$exitButton.Margin = "10,0,10,0"
[System.Windows.Controls.Grid]::SetRow($exitButton, 2)
[System.Windows.Controls.Grid]::SetColumn($exitButton, 2)
$mainGrid.Children.Add($exitButton)

# ==============================================================================
# 2. Core Logic: Initialize UI and Fetch Partition Data
# ==============================================================================

# Force WPF Message Pump Initialization (Critical for PE Environment)
$dispatcherFrame = New-Object System.Windows.Threading.DispatcherFrame
[System.Windows.Threading.Dispatcher]::CurrentDispatcher.BeginInvoke([System.Windows.Threading.DispatcherPriority]::Normal, ([Action]{
    $dispatcherFrame.Continue = $false
}))
[System.Windows.Threading.Dispatcher]::PushFrame($dispatcherFrame)

# Fetch and Filter Partition Data
$partitions = @()
if (Get-Module -ListAvailable -Name Storage) {
    $disks = Get-Disk -ErrorAction SilentlyContinue
    foreach ($disk in $disks) {
        $partitionsOnDisk = $disk | Get-Partition -ErrorAction SilentlyContinue
        foreach ($partition in $partitionsOnDisk) {
            # Filter: Basic data partitions with drive letter and size > 1GB
            if ($partition.DriveLetter -and $partition.Size -gt 1GB) {
                $volume = $partition | Get-Volume -ErrorAction SilentlyContinue
                $partitions += [PSCustomObject]@{
                    DriveLetter     = "$($partition.DriveLetter):"
                    DiskNumber      = $disk.Number
                    SizeGB          = [math]::Round($partition.Size / 1GB, 2)
                    FileSystem      = if ($volume) { $volume.FileSystem } else { "Unknown" }
                    Label           = if ($volume) { $volume.FileSystemLabel } else { "" }
                    HealthStatus    = if ($volume) { $volume.HealthStatus.ToString() } else { "Unknown" }
                }
            }
        }
    }
}

# Bind Data to DataGrid
if ($partitions.Count -gt 0) {
    $dataGrid.ItemsSource = $partitions
    $dataGrid.SelectedIndex = 0 # Auto-select first partition

} else {
    $titleLabel.Content = "Error: No valid installation partitions found!"
    $titleLabel.Foreground = [System.Windows.Media.Brushes]::Red
    $confirmButton.IsEnabled = $false
}

# ==============================================================================
# 3. Button Click Event and Window Display
# ==============================================================================

$confirmButton.Add_Click({
 
    if ($dataGrid.SelectedItem -ne $null) {
        $selectedDrive = $dataGrid.SelectedItem.DriveLetter
        
      
        if (-not $selectedDrive) {
            [System.Windows.MessageBox]::Show(
                "The selected item does not have a valid drive letter. Please select another drive.",
                "Invalid Selection",
                [System.Windows.MessageBoxButton]::OK,
                [System.Windows.MessageBoxImage]::Warning
            )
            return
        }
        
     
        $message = "You are about to install the operating system to drive $selectedDrive.`n`n" +
                   "WARNING: This will erase all data on drive $selectedDrive.`n`n" +
                   "Are you sure you want to continue?"
        
        $result = [System.Windows.MessageBox]::Show(
            $message, 
            "Confirm Installation Drive", 
            [System.Windows.MessageBoxButton]::YesNo, 
            [System.Windows.MessageBoxImage]::Warning
        )
        
      
        if ($result -eq [System.Windows.MessageBoxResult]::Yes) {
            try {
             
                $tsEnv = New-Object -ComObject Microsoft.SMS.TSEnvironment
                $tsEnv.Value("OSDInstallDrive") = $selectedDrive
                Write-Host "Successfully selected installation drive: $selectedDrive"
                
                
                $window.Close()
            } catch {
                [System.Windows.MessageBox]::Show(
                    "Failed to set MDT variable!`n`nError: $_",
                    "Error",
                    [System.Windows.MessageBoxButton]::OK,
                    [System.Windows.MessageBoxImage]::Error
                )
            }
        }
      
    } else {
        
        [System.Windows.MessageBox]::Show(
            "Please select a drive to install the operating system.",
            "No Drive Selected",
            [System.Windows.MessageBoxButton]::OK,
            [System.Windows.MessageBoxImage]::Information
        )
    }
})

$exitButton.Add_Click({
    $result = [System.Windows.MessageBox]::Show("Are you sure you want to exit MDT?`nThis will cancel the deployment process.", "Exit Confirmation", [System.Windows.MessageBoxButton]::YesNo, [System.Windows.MessageBoxImage]::Warning)
    if ($result -eq [System.Windows.MessageBoxResult]::Yes) {
        try {
            # Terminate MDT deployment
            $tsEnv = New-Object -ComObject Microsoft.SMS.TSEnvironment
            $tsEnv.Value("TSExitCode") = 1
            $tsEnv.Value("TSAction") = "Exit"
            Write-Host "Exiting MDT deployment..."
            # Force exit the script and terminate MDT
            [Environment]::Exit(1)
        } catch {
            [System.Windows.MessageBox]::Show("Failed to exit MDT!`n`nError: $_", "Error", [System.Windows.MessageBoxButton]::OK, [System.Windows.MessageBoxImage]::Error)
        }
    }
})

# Display the Window
$window.ShowDialog() | Out-Null

# Cleanup COM Object
if ($tsEnv) {
    [System.Runtime.Interopservices.Marshal]::ReleaseComObject($tsEnv) | Out-Null
}
```

选分区，过滤U盘（条件严格，未测试）

```
<#
.SYNOPSIS
MDT Partition Selection Script (Full Version with all columns).
#>

# Ensure STA mode (Required for WPF)
if ([System.Threading.Thread]::CurrentThread.GetApartmentState() -ne "STA") {
    [System.Threading.Thread]::CurrentThread.SetApartmentState("STA")
}

# Load WPF Assemblies
Add-Type -AssemblyName PresentationFramework
Add-Type -AssemblyName System.Drawing

# ==============================================================================
# 1. Create WPF Window and Controls
# ==============================================================================
$window = New-Object System.Windows.Window
$window.Title = "Select Installation Partition"
$window.Width = 800
$window.Height = 500
$window.WindowStartupLocation = "CenterScreen"
$window.ResizeMode = "NoResize"
$window.Topmost = $true
$window.Background = [System.Windows.Media.Brushes]::White

# Title
$titleLabel = New-Object System.Windows.Controls.Label
$titleLabel.Content = "Please select a partition for OS installation:"
$titleLabel.FontSize = 16
$titleLabel.FontWeight = "Bold"
$titleLabel.Margin = "0,0,0,10"
[System.Windows.Controls.Grid]::SetRow($titleLabel, 0)
$mainGrid.Children.Add($titleLabel)

$columnPartitionNumber = New-Object System.Windows.Controls.DataGridTextColumn
$columnPartitionNumber.Header = "Partition"
$columnPartitionNumber.Binding = New-Object System.Windows.Data.Binding("PartitionNumber")
$columnPartitionNumber.Width = 80
$dataGrid.Columns.Add($columnPartitionNumber)

$columnHealth = New-Object System.Windows.Controls.DataGridTextColumn
$columnHealth.Header = "Health Status"
$columnHealth.Binding = New-Object System.Windows.Data.Binding("HealthStatus")
$columnHealth.Width = "*"
$dataGrid.Columns.Add($columnHealth)

# Confirm Button
$confirmButton = New-Object System.Windows.Controls.Button
$confirmButton.Content = "Confirm Selection"
$confirmButton.FontSize = 14
$confirmButton.Width = 200
$confirmButton.Height = 40
$confirmButton.Margin = "0,20,0,0"
$confirmButton.IsDefault = $true
[System.Windows.Controls.Grid]::SetRow($confirmButton, 2)
$mainGrid.Children.Add($confirmButton)

# Fetch and Filter Partition Data (Loose Filtering)
$partitions = @()
if (Get-Module -ListAvailable -Name Storage) {
    $disks = Get-Disk -ErrorAction SilentlyContinue
    foreach ($disk in $disks) {
        # Skip removable disks to avoid USB drives
        if ($disk.BusType -eq 'USB') { continue }
        
        $partitionsOnDisk = $disk | Get-Partition -ErrorAction SilentlyContinue
        foreach ($partition in $partitionsOnDisk) {
            # Very loose filter: only skip partitions with no size
            if ($partition.Size -gt 0) {
                $volume = $partition | Get-Volume -ErrorAction SilentlyContinue
                
                $partitions += [PSCustomObject]@{
                    DriveLetter     = if ($partition.DriveLetter) { "$($partition.DriveLetter):" } else { "N/A" }
                    DiskNumber      = $disk.Number
                    PartitionNumber = $partition.PartitionNumber
                    SizeGB          = [math]::Round($partition.Size / 1GB, 2)
                    FileSystem      = if ($volume) { $volume.FileSystem } else { "Unknown" }
                    Label           = if ($volume) { $volume.FileSystemLabel } else { "" }
                    HealthStatus    = if ($volume) { $volume.HealthStatus.ToString() } else { "Unknown" }
                }
            }
        }
    }
}

# Bind Data to DataGrid
if ($partitions.Count -gt 0) {
    $dataGrid.ItemsSource = $partitions
    # Auto-select the first partition with a drive letter, if available
    $firstValidPartitionIndex = $partitions.FindIndex({$args[0].DriveLetter -ne "N/A"})
    if ($firstValidPartitionIndex -ne -1) {
        $dataGrid.SelectedIndex = $firstValidPartitionIndex
    } else {
        $dataGrid.SelectedIndex = 0 # Fallback: select first partition
    }
} else {
    $titleLabel.Content = "Error: No valid disks or partitions found!"
    $titleLabel.Foreground = [System.Windows.Media.Brushes]::Red
    $confirmButton.IsEnabled = $false
}

$confirmButton.Add_Click({
    if ($dataGrid.SelectedItem -ne $null) {
        $selectedDrive = $dataGrid.SelectedItem.DriveLetter
        # Only proceed if a drive letter is selected
        if ($selectedDrive -ne "N/A") {
            try {
                # Set MDT variable 'OSDInstallDrive'
                $tsEnv = New-Object -ComObject Microsoft.SMS.TSEnvironment
                $tsEnv.Value("OSDInstallDrive") = $selectedDrive
                Write-Host "Successfully selected installation drive: $selectedDrive"
                $window.Close()
            } catch {
                [System.Windows.MessageBox]::Show("Failed to set MDT variable!`n`nError: $_", "Error", [System.Windows.MessageBoxButton]::OK, [System.Windows.MessageBoxImage]::Error)
            }
        } else {
            [System.Windows.MessageBox]::Show("Please select a partition with a drive letter (e.g., C:).", "Invalid Selection", [System.Windows.MessageBoxButton]::OK, [System.Windows.MessageBoxImage]::Warning)
        }
    }
})

# Display the Window
$window.ShowDialog() | Out-Null

# Cleanup COM Object
if ($tsEnv) {
    [System.Runtime.Interopservices.Marshal]::ReleaseComObject($tsEnv) | Out-Null
}
```