windows-imaging/Image-Win11Pro.psm1

# Windows imaging commands for the JRE-Modules composite module.
# Importing this file is side-effect free; the interactive workflow starts only
# when JRE-ImageWin11Pro is invoked from an elevated PowerShell session.

function ConvertTo-CleanInputPath {
    param([AllowEmptyString()][string]$Path)

    if ([string]::IsNullOrWhiteSpace($Path)) { return "" }
    return $Path.Trim().Trim('"').Trim("'").TrimEnd('\')
}

function Test-Yes {
    param([AllowEmptyString()][string]$Answer)

    return $Answer -match '^(?i:y|yes)$'
}

function Test-IsAdministrator {
    $Identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $Principal = [Security.Principal.WindowsPrincipal]::new($Identity)
    return $Principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Initialize-ImageSession {
    param(
        [Parameter(Mandatory)]
        [string]$OutputDirectory,

        [Parameter(Mandatory)]
        [string]$TemporaryWorkRoot
    )

    $script:OutputDir = [IO.Path]::GetFullPath($OutputDirectory)
    $script:WorkRoot = [IO.Path]::GetFullPath($TemporaryWorkRoot)
    $script:SessionWorkDir = Join-Path $script:WorkRoot ([guid]::NewGuid().ToString("N"))
    $script:MountDir = Join-Path $script:SessionWorkDir "Mount"
    $script:ScriptMountedIso = $null
}

function Initialize-WorkDirectories {
    if (-not (Test-Path -LiteralPath $MountDir)) {
        New-Item -ItemType Directory -Force -Path $MountDir | Out-Null
    }
}

function Resolve-WimPath {
    param(
        [Parameter(Mandatory)]
        [string]$InputPath
    )

    $ResolvedInput = ConvertTo-CleanInputPath $InputPath
    if (-not (Test-Path -LiteralPath $ResolvedInput)) {
        throw "The input path does not exist: $ResolvedInput"
    }

    $Item = Get-Item -LiteralPath $ResolvedInput
    if (-not $Item.PSIsContainer) {
        if ($Item.Extension -ine ".wim") {
            throw "Expected a .wim file, but received: $ResolvedInput"
        }
        return $Item.FullName
    }

    $Candidates = @(
        (Join-Path $Item.FullName "sources\install.wim"),
        (Join-Path $Item.FullName "install.wim")
    )
    $WimPath = $Candidates | Where-Object { Test-Path -LiteralPath $_ } | Select-Object -First 1
    if (-not $WimPath) {
        throw "Could not find sources\install.wim or install.wim under: $ResolvedInput"
    }

    return (Get-Item -LiteralPath $WimPath).FullName
}

function Resolve-InstallImagePath {
    param(
        [Parameter(Mandatory)]
        [string]$InputPath
    )

    $ResolvedInput = ConvertTo-CleanInputPath $InputPath
    if (-not (Test-Path -LiteralPath $ResolvedInput)) {
        throw "The image or media path does not exist: $ResolvedInput"
    }

    $Item = Get-Item -LiteralPath $ResolvedInput
    if (-not $Item.PSIsContainer) {
        if ($Item.Extension -notin @(".wim", ".esd")) {
            throw "Expected a .wim or .esd file, but received: $ResolvedInput"
        }
        return $Item.FullName
    }

    $Candidates = @(
        (Join-Path $Item.FullName "sources\install.wim"),
        (Join-Path $Item.FullName "sources\install.esd"),
        (Join-Path $Item.FullName "install.wim"),
        (Join-Path $Item.FullName "install.esd")
    )
    $ImagePath = $Candidates | Where-Object { Test-Path -LiteralPath $_ } | Select-Object -First 1
    if (-not $ImagePath) {
        throw "Could not find install.wim or install.esd under: $ResolvedInput"
    }

    return (Get-Item -LiteralPath $ImagePath).FullName
}

function Get-OptionalDriverPath {
    Write-Host "`n Optional: Extracted (.inf) Drivers Folder" -ForegroundColor DarkGray
    Write-Host " (Press ENTER to skip driver injection)" -ForegroundColor DarkGray
    $DriverPath = ConvertTo-CleanInputPath (Read-Host " Path to Drivers folder")
    if ($DriverPath -and -not (Test-Path -LiteralPath $DriverPath -PathType Container)) {
        throw "Driver folder not found at: $DriverPath"
    }
    return $DriverPath
}

function Copy-WindowsMedia {
    param(
        [Parameter(Mandatory)]
        [string]$SourceDrive,
        [Parameter(Mandatory)]
        [string]$Destination
    )

    if (-not (Test-Path -LiteralPath $SourceDrive -PathType Container)) {
        throw "Source drive or folder not found: $SourceDrive"
    }

    if (Test-Path -LiteralPath $Destination -PathType Container) {
        $ExistingItems = @(Get-ChildItem -LiteralPath $Destination -Force -ErrorAction Stop)
        if ($ExistingItems.Count -gt 0 -and
            (Test-Yes (Read-Host "`n Destination '$Destination' contains files. Clear it before copying? (y/n)"))) {
            Write-Host " Clearing existing destination contents..." -ForegroundColor Yellow
            $ExistingItems | Remove-Item -Recurse -Force -ErrorAction Stop
        }
    }
    else {
        New-Item -ItemType Directory -Force -Path $Destination | Out-Null
    }

    Write-Host "`n Copying Windows media contents to $Destination..." -ForegroundColor Yellow
    Copy-Item -Path (Join-Path $SourceDrive "*") -Destination $Destination -Recurse -Force

    Write-Host " Removing read-only flags..." -ForegroundColor DarkGray
    Get-ChildItem -LiteralPath $Destination -Recurse -File |
    ForEach-Object { $_.IsReadOnly = $false }
}

function Export-WindowsProImage {
    param(
        [Parameter(Mandatory)]
        [string]$SourceImage
    )

    Initialize-WorkDirectories
    Write-Host "`n Scanning Windows editions in image..." -ForegroundColor Yellow
    $ImageInfo = @(Get-WindowsImage -ImagePath $SourceImage)
    if ($ImageInfo.Count -eq 0) {
        throw "No Windows images were found in: $SourceImage"
    }

    $ProEdition = $ImageInfo |
    Where-Object { $_.ImageName -match '(?i)Windows 11 Pro(?:\s|$)' } |
    Select-Object -First 1

    if (-not $ProEdition) {
        Write-Host "`n Available Editions:" -ForegroundColor Cyan
        $ImageInfo | Format-Table ImageIndex, ImageName
        $ProIndex = Read-Host " Could not auto-detect Windows 11 Pro. Enter its image index"
        if ($ProIndex -notmatch '^\d+$' -or [int]$ProIndex -notin $ImageInfo.ImageIndex) {
            throw "Image index '$ProIndex' is not valid for this image."
        }
    }
    else {
        $ProIndex = $ProEdition.ImageIndex
        Write-Host " Detected Windows 11 Pro at index: $ProIndex" -ForegroundColor Green
    }

    if ($ImageInfo.Count -eq 1 -and [IO.Path]::GetExtension($SourceImage) -ieq ".wim") {
        Write-Host " [Smart Detect] Image is already a single-edition WIM; export skipped." -ForegroundColor Cyan
        return $SourceImage
    }

    $TargetWimPath = Join-Path (Split-Path -Parent $SourceImage) "install.wim"
    $TemporaryWim = Join-Path $SessionWorkDir "install.wim"

    Write-Host "`n Exporting Windows 11 Pro into a clean WIM..." -ForegroundColor Yellow
    & dism.exe /Export-Image /SourceImageFile:"$SourceImage" /SourceIndex:$ProIndex /DestinationImageFile:"$TemporaryWim" /Compress:max /CheckIntegrity
    if ($LASTEXITCODE -ne 0 -or -not (Test-Path -LiteralPath $TemporaryWim)) {
        throw "Failed to export the Pro edition with DISM (exit code $LASTEXITCODE)."
    }

    if ([IO.Path]::GetFullPath($SourceImage) -ine [IO.Path]::GetFullPath($TargetWimPath)) {
        Remove-Item -LiteralPath $SourceImage -Force
    }
    Remove-Item -LiteralPath $TargetWimPath -Force -ErrorAction SilentlyContinue
    Move-Item -LiteralPath $TemporaryWim -Destination $TargetWimPath
    return $TargetWimPath
}

function Add-WindowsDrivers {
    param(
        [Parameter(Mandatory)]
        [string]$WimPath,
        [Parameter(Mandatory)]
        [string]$DriverPath
    )

    if (-not (Test-Path -LiteralPath $WimPath -PathType Leaf)) {
        throw "WIM file not found: $WimPath"
    }
    if (-not (Test-Path -LiteralPath $DriverPath -PathType Container)) {
        throw "Driver folder not found: $DriverPath"
    }

    Initialize-WorkDirectories
    Write-Host "`n Mounting WIM and injecting drivers..." -ForegroundColor Yellow
    & dism.exe /Mount-Image /ImageFile:"$WimPath" /Index:1 /MountDir:"$MountDir"
    if ($LASTEXITCODE -ne 0) {
        throw "Failed to mount the WIM with DISM (exit code $LASTEXITCODE)."
    }

    & dism.exe /Image:"$MountDir" /Add-Driver /Driver:"$DriverPath" /Recurse
    if ($LASTEXITCODE -ne 0) {
        throw "Failed to add one or more drivers with DISM (exit code $LASTEXITCODE)."
    }

    Write-Host " Committing changes and unmounting WIM..." -ForegroundColor DarkGray
    & dism.exe /Unmount-Image /MountDir:"$MountDir" /Commit
    if ($LASTEXITCODE -ne 0) {
        throw "Failed to commit and unmount the WIM (exit code $LASTEXITCODE)."
    }
}

function Split-WindowsImage {
    param(
        [Parameter(Mandatory)]
        [string]$WimPath
    )

    $WimPath = Resolve-WimPath $WimPath
    $WimDirectory = Split-Path -Parent $WimPath
    $SwmPath = Join-Path $WimDirectory "install.swm"

    $ExistingParts = @(Get-ChildItem -LiteralPath $WimDirectory -Filter "install*.swm" -File)
    if ($ExistingParts.Count -gt 0) {
        Write-Host "`n Removing $($ExistingParts.Count) existing SWM part(s) to avoid mixing output sets..." -ForegroundColor DarkGray
        $ExistingParts | Remove-Item -Force
    }

    Write-Host "`n Splitting WIM into FAT32-safe 3800 MB SWM parts..." -ForegroundColor Yellow
    & dism.exe /Split-Image /ImageFile:"$WimPath" /SWMFile:"$SwmPath" /FileSize:3800 /CheckIntegrity
    if ($LASTEXITCODE -ne 0) {
        throw "Failed to split the WIM with DISM (exit code $LASTEXITCODE)."
    }

    $CreatedParts = @(Get-ChildItem -LiteralPath $WimDirectory -Filter "install*.swm" -File | Sort-Object Name)
    if ($CreatedParts.Count -eq 0 -or -not (Test-Path -LiteralPath $SwmPath)) {
        throw "DISM reported success, but install.swm was not created."
    }

    Write-Host " Created $($CreatedParts.Count) SWM part(s):" -ForegroundColor Green
    $CreatedParts | ForEach-Object {
        Write-Host (" {0} ({1:N0} MB)" -f $_.Name, ($_.Length / 1MB)) -ForegroundColor Gray
    }

    if (Test-Yes (Read-Host "`n SWM files verified. Remove the original WIM? (y/n)")) {
        Remove-Item -LiteralPath $WimPath -Force
        Write-Host " Removed original WIM: $WimPath" -ForegroundColor DarkGray
    }

    return $CreatedParts
}

function Select-SourceDrive {
    Write-Host "`n Select Windows 11 Source Image" -ForegroundColor Yellow
    Write-Host " ───────────────────────────────────" -ForegroundColor Gray
    Write-Host " 1. Path to a Windows 11 .ISO file (auto-mount)"
    Write-Host " 2. Drive letter or folder of mounted ISO/USB media"
    $SourceChoice = Read-Host "`n Choose option (1 or 2)"

    switch ($SourceChoice) {
        "1" {
            $IsoPath = ConvertTo-CleanInputPath (Read-Host "`n Enter full path to .ISO file")
            if (-not (Test-Path -LiteralPath $IsoPath -PathType Leaf)) {
                throw "ISO file not found: $IsoPath"
            }

            Write-Host "`n Mounting ISO image..." -ForegroundColor DarkGray
            $MountedDisk = Mount-DiskImage -ImagePath $IsoPath -PassThru
            $Volume = $MountedDisk | Get-Volume | Where-Object DriveLetter | Select-Object -First 1
            if (-not $Volume) {
                throw "The ISO mounted, but no drive letter was assigned."
            }

            $script:ScriptMountedIso = $IsoPath
            $IsoDrive = "$($Volume.DriveLetter):"
            Write-Host " Successfully mounted ISO to drive $IsoDrive" -ForegroundColor Green
            return $IsoDrive
        }
        "2" {
            $SourcePath = ConvertTo-CleanInputPath (Read-Host "`n Enter drive letter or media folder")
            if ($SourcePath -match '^[A-Za-z]:?$') {
                $SourcePath = $SourcePath.TrimEnd(':') + ":\"
            }
            if (-not (Test-Path -LiteralPath $SourcePath -PathType Container)) {
                throw "Source drive or folder not found: $SourcePath"
            }
            return $SourcePath
        }
        default {
            throw "Invalid source option '$SourceChoice'. Choose 1 or 2."
        }
    }
}

function New-BootableUsbDrive {
    param(
        [Parameter(Mandatory)]
        [string]$SourceFolder
    )

    $SourceFolder = ConvertTo-CleanInputPath $SourceFolder
    if (-not (Test-Path -LiteralPath $SourceFolder -PathType Container)) {
        throw "Windows media folder not found: $SourceFolder"
    }
    $SourceFolder = (Get-Item -LiteralPath $SourceFolder).FullName

    $RequiredFiles = @(
        "EFI\BOOT\BOOTX64.EFI",
        "sources\boot.wim"
    )
    foreach ($RelativePath in $RequiredFiles) {
        if (-not (Test-Path -LiteralPath (Join-Path $SourceFolder $RelativePath) -PathType Leaf)) {
            throw "The media folder is not UEFI-bootable because '$RelativePath' is missing."
        }
    }

    $InstallImages = @(
        (Join-Path $SourceFolder "sources\install.wim"),
        (Join-Path $SourceFolder "sources\install.esd"),
        (Join-Path $SourceFolder "sources\install.swm")
    )
    if (-not ($InstallImages | Where-Object { Test-Path -LiteralPath $_ -PathType Leaf })) {
        throw "The media folder does not contain install.wim, install.esd, or install.swm."
    }

    $OversizedImages = @(Get-ChildItem -LiteralPath (Join-Path $SourceFolder "sources") -File |
        Where-Object { $_.Extension -in @(".wim", ".esd") -and $_.Length -ge 4GB })
    if ($OversizedImages.Count -gt 0) {
        $Names = $OversizedImages.Name -join ", "
        throw "FAT32 cannot store files of 4 GB or larger ($Names). Split install.wim into SWM files first."
    }

    $UsbVolumes = @(
        Get-Disk |
        Where-Object { $_.BusType -eq "USB" -and -not $_.IsBoot -and -not $_.IsSystem } |
        ForEach-Object {
            $Disk = $_
            Get-Partition -DiskNumber $Disk.Number -ErrorAction SilentlyContinue |
            Where-Object DriveLetter |
            ForEach-Object {
                $Volume = Get-Volume -DriveLetter $_.DriveLetter -ErrorAction SilentlyContinue
                if ($Volume) {
                    [pscustomobject]@{
                        DriveLetter = $_.DriveLetter
                        Label       = $Volume.FileSystemLabel
                        FileSystem  = $Volume.FileSystem
                        SizeGB      = [math]::Round($Volume.Size / 1GB, 1)
                        FreeGB      = [math]::Round($Volume.SizeRemaining / 1GB, 1)
                        Disk        = $Disk.FriendlyName
                    }
                }
            }
        }
    )
    if ($UsbVolumes.Count -eq 0) {
        throw "No USB volumes with drive letters were found. Insert and format a flash drive, then try again."
    }

    Write-Host "`n Available USB volumes:" -ForegroundColor Cyan
    $UsbVolumes | Format-Table DriveLetter, Label, FileSystem, SizeGB, FreeGB, Disk -AutoSize | Out-Host

    $DriveLetter = (Read-Host " Enter the FAT32 USB drive letter to receive the installation media").Trim().TrimEnd(":")
    $TargetVolume = $UsbVolumes |
    Where-Object { $_.DriveLetter -ieq $DriveLetter } |
    Select-Object -First 1
    if (-not $TargetVolume) {
        throw "Drive '$DriveLetter' is not one of the USB volumes listed above."
    }
    if ($TargetVolume.FileSystem -ne "FAT32") {
        throw "Drive $($TargetVolume.DriveLetter): uses $($TargetVolume.FileSystem), not FAT32. Format it as FAT32 first."
    }

    $Destination = "$($TargetVolume.DriveLetter):\"
    $SourceRoot = [IO.Path]::GetPathRoot($SourceFolder)
    if ($SourceRoot -and $SourceRoot -ieq $Destination) {
        throw "The source media and destination are on the same drive."
    }

    # Fresh FAT32 volumes often have System Volume Information / $RECYCLE.BIN.
    # Explorer hides these; ignore them when deciding if the drive "has files".
    $IgnoredRootNames = @('System Volume Information', '$RECYCLE.BIN', 'RECYCLER')
    $ExistingItems = @(
        Get-ChildItem -LiteralPath $Destination -Force -ErrorAction Stop |
        Where-Object { $_.Name -notin $IgnoredRootNames }
    )
    if ($ExistingItems.Count -gt 0) {
        Write-Host "`n [!] Drive $($TargetVolume.DriveLetter): already contains $($ExistingItems.Count) item(s)." -ForegroundColor Yellow
        $ExistingItems | Select-Object -First 10 | ForEach-Object {
            Write-Host (" - {0}" -f $_.Name) -ForegroundColor Gray
        }
        if ($ExistingItems.Count -gt 10) {
            Write-Host (" ... and {0} more" -f ($ExistingItems.Count - 10)) -ForegroundColor Gray
        }

        if (Test-Yes (Read-Host " Delete existing files on the drive before copying? (y/n)")) {
            Write-Host " Deleting existing files on $Destination..." -ForegroundColor Yellow
            foreach ($Item in $ExistingItems) {
                Remove-Item -LiteralPath $Item.FullName -Recurse -Force -ErrorAction Stop
            }
            Write-Host " Drive cleared." -ForegroundColor Green
        }
        elseif (-not (Test-Yes (Read-Host " Continue without deleting (matching names may be overwritten)? (y/n)"))) {
            throw "USB media copy cancelled."
        }
    }

    $SourceSize = (Get-ChildItem -LiteralPath $SourceFolder -Recurse -File |
        Measure-Object -Property Length -Sum).Sum
    $DestinationVolume = Get-Volume -DriveLetter $TargetVolume.DriveLetter
    if ($SourceSize -gt $DestinationVolume.SizeRemaining) {
        Write-Host " [!] The source is larger than the reported free space. The copy could fail." -ForegroundColor Yellow
    }

    # /E recurse (incl. empty dirs), /J unbuffered I/O for large files,
    # /MT multithreaded copy, /R /W limited retries.
    Write-Host " Copying Windows installation files to $Destination with robocopy..." -ForegroundColor Yellow
    & robocopy.exe $SourceFolder $Destination /E /COPY:DAT /DCOPY:DAT /J /MT:8 /R:2 /W:2 /XJ | Out-Host
    $RobocopyExitCode = $LASTEXITCODE
    if ($RobocopyExitCode -ge 8) {
        throw "Failed to copy Windows media to the USB drive (Robocopy exit code $RobocopyExitCode)."
    }

    foreach ($RelativePath in $RequiredFiles) {
        if (-not (Test-Path -LiteralPath (Join-Path $Destination $RelativePath) -PathType Leaf)) {
            throw "USB verification failed because '$RelativePath' was not copied."
        }
    }
    if (-not (Test-Path -LiteralPath (Join-Path $Destination "sources\install.wim")) -and
        -not (Test-Path -LiteralPath (Join-Path $Destination "sources\install.esd")) -and
        -not (Test-Path -LiteralPath (Join-Path $Destination "sources\install.swm"))) {
        throw "USB verification failed because no Windows install image was copied."
    }

    Write-Host " Bootable UEFI USB verified successfully." -ForegroundColor Green
    return $Destination
}

function Get-OscdimgPath {
    $Command = Get-Command oscdimg.exe -ErrorAction SilentlyContinue
    if ($Command) {
        return $Command.Source
    }

    $KnownPaths = @(
        "${env:ProgramFiles(x86)}\Windows Kits\10\Assessment and Deployment Kit\Deployment Tools\amd64\Oscdimg\oscdimg.exe",
        "${env:ProgramFiles(x86)}\Windows Kits\10\Assessment and Deployment Kit\Deployment Tools\x86\Oscdimg\oscdimg.exe"
    )
    $OscdimgPath = $KnownPaths | Where-Object { Test-Path -LiteralPath $_ -PathType Leaf } | Select-Object -First 1
    if (-not $OscdimgPath) {
        throw "Oscdimg.exe was not found. Install the Deployment Tools component of the Windows ADK."
    }

    return $OscdimgPath
}

function New-BootableIso {
    param(
        [Parameter(Mandatory)]
        [string]$SourceFolder,
        [Parameter(Mandatory)]
        [string]$OutputPath
    )

    $SourceFolder = ConvertTo-CleanInputPath $SourceFolder
    if (-not (Test-Path -LiteralPath $SourceFolder -PathType Container)) {
        throw "Windows media folder not found: $SourceFolder"
    }
    $SourceFolder = (Get-Item -LiteralPath $SourceFolder).FullName

    $RequiredFiles = @(
        "boot\etfsboot.com",
        "efi\microsoft\boot\efisys.bin",
        "EFI\BOOT\BOOTX64.EFI",
        "sources\boot.wim"
    )
    foreach ($RelativePath in $RequiredFiles) {
        if (-not (Test-Path -LiteralPath (Join-Path $SourceFolder $RelativePath) -PathType Leaf)) {
            throw "Cannot create a bootable ISO because '$RelativePath' is missing."
        }
    }

    $InstallImages = @(
        (Join-Path $SourceFolder "sources\install.wim"),
        (Join-Path $SourceFolder "sources\install.esd"),
        (Join-Path $SourceFolder "sources\install.swm")
    )
    if (-not ($InstallImages | Where-Object { Test-Path -LiteralPath $_ -PathType Leaf })) {
        throw "The media folder does not contain install.wim, install.esd, or install.swm."
    }

    $OutputPath = ConvertTo-CleanInputPath $OutputPath
    if ([IO.Path]::GetExtension($OutputPath) -ine ".iso") {
        $OutputPath += ".iso"
    }
    $OutputPath = [IO.Path]::GetFullPath($OutputPath)
    if ($OutputPath.StartsWith($SourceFolder.TrimEnd('\') + "\", [StringComparison]::OrdinalIgnoreCase)) {
        throw "The ISO output cannot be located inside the source media folder."
    }

    $OutputDirectory = Split-Path -Parent $OutputPath
    if (-not (Test-Path -LiteralPath $OutputDirectory -PathType Container)) {
        New-Item -ItemType Directory -Force -Path $OutputDirectory | Out-Null
    }
    if (Test-Path -LiteralPath $OutputPath) {
        if (-not (Test-Yes (Read-Host " '$OutputPath' already exists. Replace it? (y/n)"))) {
            throw "ISO creation cancelled."
        }
        Remove-Item -LiteralPath $OutputPath -Force
    }

    $OscdimgPath = Get-OscdimgPath
    $BiosBootImage = Join-Path $SourceFolder "boot\etfsboot.com"
    $UefiBootImage = Join-Path $SourceFolder "efi\microsoft\boot\efisys.bin"
    $BootData = "-bootdata:2#p0,e,b$BiosBootImage#pEF,e,b$UefiBootImage"

    Write-Host "`n Creating BIOS/UEFI bootable ISO..." -ForegroundColor Yellow
    & $OscdimgPath -m -o -u2 -udfver102 -lCUSTOMWIN11 $BootData $SourceFolder $OutputPath | Out-Host
    if ($LASTEXITCODE -ne 0) {
        throw "Oscdimg failed to create the ISO (exit code $LASTEXITCODE)."
    }
    if (-not (Test-Path -LiteralPath $OutputPath -PathType Leaf) -or
        (Get-Item -LiteralPath $OutputPath).Length -eq 0) {
        throw "Oscdimg reported success, but the ISO was not created."
    }

    $IsoSize = (Get-Item -LiteralPath $OutputPath).Length
    Write-Host (" Bootable ISO created: {0} ({1:N2} GB)" -f $OutputPath, ($IsoSize / 1GB)) -ForegroundColor Green
    return $OutputPath
}

function Read-NinjaOneLocationConfig {
    param(
        [Parameter(Mandatory)]
        [string]$ConfigPath
    )

    $ResolvedPath = ConvertTo-CleanInputPath $ConfigPath
    if (-not (Test-Path -LiteralPath $ResolvedPath -PathType Leaf)) {
        throw "NinjaOne configuration file not found: $ResolvedPath"
    }

    try {
        $Config = Get-Content -LiteralPath $ResolvedPath -Raw -ErrorAction Stop |
            ConvertFrom-Json -ErrorAction Stop
    }
    catch {
        throw "Unable to read NinjaOne configuration '$ResolvedPath': $($_.Exception.Message)"
    }

    if ($null -eq $Config -or $Config.PSObject.Properties.Name -notcontains 'Locations') {
        throw "NinjaOne configuration must contain a 'Locations' array."
    }

    $Locations = @($Config.Locations)
    if ($Locations.Count -eq 0) {
        throw "NinjaOne configuration must contain at least one location."
    }

    $ValidatedLocations = foreach ($Location in $Locations) {
        $Name = [string]$Location.Name
        $Token = [string]$Location.Token

        if ([string]::IsNullOrWhiteSpace($Name)) {
            throw "Every NinjaOne location must have a non-empty Name."
        }
        if ([string]::IsNullOrWhiteSpace($Token)) {
            throw "NinjaOne location '$Name' must have a non-empty Token."
        }

        [pscustomobject]@{
            Name  = $Name.Trim()
            Token = $Token.Trim()
        }
    }

    $DuplicateNames = @(
        $ValidatedLocations |
            Group-Object -Property Name |
            Where-Object Count -gt 1
    )
    if ($DuplicateNames.Count -gt 0) {
        throw "NinjaOne location names must be unique. Duplicate: $($DuplicateNames[0].Name)"
    }

    return @($ValidatedLocations | Sort-Object -Property Name)
}

function New-NinjaOneInstallerSource {
    param(
        [Parameter(Mandatory)]
        [object[]]$Locations
    )

    $LocationJson = ConvertTo-Json -InputObject @($Locations) -Depth 4 -Compress
    $LocationData = [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($LocationJson))

    $Source = @'
# NinjaOne Agent Interactive Installer
$ErrorActionPreference = 'Stop'
$ProgressPreference = 'SilentlyContinue'

$isAdmin = ([Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole(
    [Security.Principal.WindowsBuiltInRole]::Administrator
)
if (-not $isAdmin) {
    throw 'This installer must be run as Administrator.'
}

$locationJson = [Text.Encoding]::UTF8.GetString(
    [Convert]::FromBase64String('__LOCATION_DATA__')
)
$locations = @($locationJson | ConvertFrom-Json | Sort-Object -Property Name)

Write-Host '==================================================' -ForegroundColor Cyan
Write-Host ' NinjaOne Agent Installer - J Ranck ' -ForegroundColor Cyan
Write-Host '==================================================' -ForegroundColor Cyan
Write-Host "Select the location for this deployment:`n" -ForegroundColor Yellow

for ($i = 0; $i -lt $locations.Count; $i++) {
    Write-Host ('[{0,2}] {1}' -f ($i + 1), $locations[$i].Name)
}

do {
    $selection = Read-Host "`nEnter selection number (1-$($locations.Count))"
    $selectedIndex = if ($selection -match '^\d+$') { [int]$selection - 1 } else { -1 }
} while ($selectedIndex -lt 0 -or $selectedIndex -ge $locations.Count)

$selected = $locations[$selectedIndex]
$token = $selected.Token
Write-Host "`n[+] Selected Location: $($selected.Name)" -ForegroundColor Green

$msiPath = Join-Path $env:TEMP 'NinjaOneAgent-x86.msi'
try {
    Write-Host "`n[1/2] Downloading NinjaOne Agent MSI..." -ForegroundColor Yellow
    Invoke-WebRequest -Uri 'https://jranck.rmmservice.com/ws/api/v2/generic-installer/NinjaOneAgent-x86.msi' -OutFile $msiPath -UseBasicParsing

    Write-Host '[2/2] Installing NinjaOne Agent...' -ForegroundColor Yellow
    $msiArgs = "/i `"$msiPath`" TOKENID=`"$token`" /qn"
    $process = Start-Process -FilePath 'msiexec.exe' -ArgumentList $msiArgs -Wait -PassThru
    if ($process.ExitCode -ne 0) {
        throw "NinjaOne installation failed with MSI exit code $($process.ExitCode)."
    }

    Write-Host "`n[+] Installation completed successfully." -ForegroundColor Green
}
finally {
    Remove-Item -LiteralPath $msiPath -Force -ErrorAction SilentlyContinue
}

Write-Host "`nPress any key to exit..."
$null = $Host.UI.RawUI.ReadKey('NoEcho,IncludeKeyDown')
'@


    return $Source.Replace('__LOCATION_DATA__', $LocationData)
}

<#
.SYNOPSIS
    Creates the interactive NinjaOne installer executable used with prepared Windows media.
.DESCRIPTION
    Reads NinjaOne location names and tokens from an external JSON configuration, generates
    a temporary installer script, and packages it with PS2EXE. PS2EXE embeds the generated
    script as extractable plaintext, so the resulting executable must be handled as a secret.
.PARAMETER ConfigPath
    Path to a JSON file containing a Locations array with Name and Token properties.
.PARAMETER OutputPath
    Destination path for Install-NinjaOne.exe. Defaults to the current directory.
.PARAMETER Force
    Replaces an existing executable at OutputPath.
.EXAMPLE
    JRE-NewNinjaOneInstaller -ConfigPath 'C:\Secure\NinjaOne-Locations.json'
.EXAMPLE
    JRE-NewNinjaOneInstaller -ConfigPath 'C:\Secure\NinjaOne-Locations.json' -OutputPath 'D:\Media\Install\Install-NinjaOne.exe' -Force
#>

function JRE-NewNinjaOneInstaller {
    [CmdletBinding()]
    param(
        [Parameter(Mandatory)]
        [ValidateNotNullOrEmpty()]
        [string]$ConfigPath,

        [Parameter()]
        [ValidateNotNullOrEmpty()]
        [string]$OutputPath = (Join-Path (Get-Location).ProviderPath 'Install-NinjaOne.exe'),

        [Parameter()]
        [switch]$Force
    )

    $Compiler = Get-Command -Name Invoke-ps2exe -ErrorAction SilentlyContinue
    if (-not $Compiler) {
        throw "PS2EXE is required to create the NinjaOne installer. Run 'Install-Module -Name ps2exe -Scope CurrentUser', then retry."
    }

    $Locations = Read-NinjaOneLocationConfig -ConfigPath $ConfigPath
    $ResolvedOutputPath = [IO.Path]::GetFullPath((ConvertTo-CleanInputPath $OutputPath))
    if ([IO.Path]::GetExtension($ResolvedOutputPath) -ine '.exe') {
        throw "NinjaOne installer output path must end in .exe: $ResolvedOutputPath"
    }

    $OutputDirectory = Split-Path -Parent $ResolvedOutputPath
    if (-not (Test-Path -LiteralPath $OutputDirectory -PathType Container)) {
        New-Item -ItemType Directory -Path $OutputDirectory -Force | Out-Null
    }

    if (Test-Path -LiteralPath $ResolvedOutputPath) {
        if (-not $Force) {
            throw "NinjaOne installer already exists: $ResolvedOutputPath. Use -Force to replace it."
        }
        Remove-Item -LiteralPath $ResolvedOutputPath -Force
    }

    $BuildDirectory = Join-Path ([IO.Path]::GetTempPath()) ("JRE-NinjaOne-{0}" -f ([guid]::NewGuid().ToString('N')))
    $SourcePath = Join-Path $BuildDirectory 'Install-NinjaOne.ps1'

    try {
        New-Item -ItemType Directory -Path $BuildDirectory -Force | Out-Null
        New-NinjaOneInstallerSource -Locations $Locations |
            Set-Content -LiteralPath $SourcePath -Encoding UTF8 -Force

        $CompilerParameters = @{
            inputFile   = $SourcePath
            outputFile  = $ResolvedOutputPath
            requireAdmin = $true
            title       = 'NinjaOne Agent Installer'
            description = 'J Ranck NinjaOne Agent Deployment Tool'
            company     = 'J-Ranck-Electric'
            product     = 'NinjaOne Agent Installer'
            supportOS   = $true
        }
        & $Compiler @CompilerParameters

        if (-not (Test-Path -LiteralPath $ResolvedOutputPath -PathType Leaf) -or
            (Get-Item -LiteralPath $ResolvedOutputPath).Length -eq 0) {
            throw "PS2EXE did not create the NinjaOne installer at: $ResolvedOutputPath"
        }

        Write-Host " NinjaOne installer created: $ResolvedOutputPath" -ForegroundColor Green
        return (Get-Item -LiteralPath $ResolvedOutputPath)
    }
    finally {
        Remove-Item -LiteralPath $BuildDirectory -Recurse -Force -ErrorAction SilentlyContinue
    }
}

function Show-StageMenu {
    Write-Host " Choose where to start:" -ForegroundColor Yellow
    Write-Host " 1. Full workflow (copy media, isolate Pro, drivers, optional ISO)"
    Write-Host " 2. Process existing media folder or WIM/ESD (isolate Pro, drivers, optional split)"
    Write-Host " 3. Inject drivers into an existing WIM"
    Write-Host " 4. Create a bootable ISO from a prepared media folder"
    Write-Host " 5. Split an existing WIM into SWM files"
    Write-Host " 6. Copy installation media to an existing FAT32 USB drive"
    return Read-Host "`n Choose option (1-6)"
}

function JRE-ImageWin11Help {
    [CmdletBinding()]
    param()

    Write-Host "`n=== Windows 11 Pro Image Customization ===" -ForegroundColor Cyan
    Write-Host "`nRun " -NoNewline
    Write-Host "JRE-ImageWin11Pro" -ForegroundColor Magenta -NoNewline
    Write-Host " as Administrator to start the interactive workflow." -ForegroundColor Cyan
    Write-Host "Default destination folder: " -NoNewline
    Write-Host "C:\CustomWin11_Output" -ForegroundColor Green
    Write-Host "Default temporary work folder: " -NoNewline
    Write-Host "C:\Win11_TempWork" -ForegroundColor Green

    Write-Host "`nStage menu:" -ForegroundColor Yellow
    Write-Host " 1. Full workflow (copy media, isolate Pro, drivers, optional ISO)"
    Write-Host " 2. Process existing media folder or WIM/ESD (isolate Pro, drivers, optional split)"
    Write-Host " 3. Inject drivers into an existing WIM"
    Write-Host " 4. Create a bootable ISO from a prepared media folder"
    Write-Host " 5. Split an existing WIM into SWM files"
    Write-Host " 6. Copy installation media to an existing FAT32 USB drive"
    Write-Host "`nNinjaOne installer:" -ForegroundColor Yellow
    Write-Host " JRE-NewNinjaOneInstaller -ConfigPath <secure-json-path>"
    Write-Host " The full workflow can place the EXE in the media's Install folder automatically."
    Write-Host ""
}

<#
.SYNOPSIS
    Creates or resumes a customized, single-edition Windows 11 Pro installation image.
.DESCRIPTION
    Each image-processing stage can run independently. This allows an existing media
    folder or install.wim to be reused for driver injection or SWM splitting. The
    command requires an elevated PowerShell session and prompts for workflow choices.
.PARAMETER OutputDirectory
    Folder used for the prepared Windows installation media.
.PARAMETER WorkRoot
    Parent folder used for temporary mount and image-processing files. A unique
    session folder is created beneath it and removed when processing finishes.
.EXAMPLE
    JRE-ImageWin11Pro
.EXAMPLE
    JRE-ImageWin11Pro -OutputDirectory 'D:\WindowsMedia' -WorkRoot 'D:\ImageWork'
#>

function JRE-ImageWin11Pro {
    [CmdletBinding()]
    param(
        [Parameter()]
        [ValidateNotNullOrEmpty()]
        [string]$OutputDirectory = "C:\CustomWin11_Output",

        [Parameter()]
        [ValidateNotNullOrEmpty()]
        [string]$WorkRoot = "C:\Win11_TempWork"
    )

    if (-not (Test-IsAdministrator)) {
        Write-Host "This command requires Administrator privileges. Would you like to launch PowerShell as Administrator? (y/n):" -ForegroundColor Yellow
        $response = Read-Host
        if ($response -eq "y") {
            $currentDirectory = (Get-Location).ProviderPath

            if (Get-Command wt.exe -ErrorAction SilentlyContinue) {
                Start-Process wt.exe -Verb RunAs `
                    -WorkingDirectory $currentDirectory `
                    -ArgumentList @('-d', "`"$currentDirectory`"", 'pwsh.exe', '-NoExit')
            }
            else {
                Start-Process pwsh.exe -Verb RunAs `
                    -WorkingDirectory $currentDirectory `
                    -ArgumentList '-NoExit'
            }
            exit
        }
        else {
            throw "Please run this command in PowerShell as Administrator."
        }
    }

    Initialize-ImageSession -OutputDirectory $OutputDirectory -TemporaryWorkRoot $WorkRoot

    Clear-Host
    Write-Host "┌────────────────────────────────────────────────────────┐" -ForegroundColor Cyan
    Write-Host "│ Windows 11 Pro Image Customization Tool │" -ForegroundColor Cyan
    Write-Host "└────────────────────────────────────────────────────────┘" -ForegroundColor Cyan
    Write-Host ""
    Write-Host " [Notice] Default destination folder: $OutputDir" -ForegroundColor Green
    Write-Host ""

    try {
        $StageChoice = Show-StageMenu

        Write-Host "`n ────────────────────────────────────────────────────────" -ForegroundColor Gray
        Write-Host " Starting Image Processing..." -ForegroundColor Cyan
        Write-Host " ────────────────────────────────────────────────────────" -ForegroundColor Gray

        switch ($StageChoice) {
            "1" {
                $SourceDrive = Select-SourceDrive
                $DriverDir = Get-OptionalDriverPath
                Copy-WindowsMedia -SourceDrive $SourceDrive -Destination $OutputDir
                $SourceImage = Resolve-InstallImagePath $OutputDir
                $WorkingWim = Export-WindowsProImage -SourceImage $SourceImage

                # A full ISO workflow uses the intact WIM. Remove split-image parts
                # that may remain in the output folder from an earlier run.
                Get-ChildItem -LiteralPath (Join-Path $OutputDir "sources") -Filter "install*.swm" -File -ErrorAction SilentlyContinue |
                Remove-Item -Force

                if ($DriverDir) {
                    Add-WindowsDrivers -WimPath $WorkingWim -DriverPath $DriverDir
                }
                else {
                    Write-Host "`n Skipping driver injection (no drivers specified)." -ForegroundColor DarkGray
                }

                if (Test-Yes (Read-Host "`n Create and add the NinjaOne installer to this media? (y/n)")) {
                    # output example json
                    Write-Host "Example JSON:" -ForegroundColor Yellow
                    Write-Host @'
{
    "Locations": [
        {
        "Name": "Example Location",
        "Token": "replace-with-the-ninjaone-install-token"
        },
        {
        "Name": "Example Location 2",
        "Token": "replace-with-the-ninjaone-install-token-2"
        }
    ]
}
'@
 -ForegroundColor Yellow

                    $NinjaOneConfigPath = ConvertTo-CleanInputPath (Read-Host " Path to NinjaOne location/token JSON config")
                    if (-not $NinjaOneConfigPath) {
                        throw "A NinjaOne configuration path is required to create the installer."
                    }

                    $NinjaOneOutputPath = Join-Path $OutputDir 'Install\Install-NinjaOne.exe'
                    JRE-NewNinjaOneInstaller -ConfigPath $NinjaOneConfigPath -OutputPath $NinjaOneOutputPath -Force |
                        Out-Null
                }
                else {
                    Write-Host " Skipping NinjaOne installer creation." -ForegroundColor DarkGray
                }

                $ResultLocation = $OutputDir
                Write-Host "`n Prepared media folder: $OutputDir" -ForegroundColor Green
                if (Test-Yes (Read-Host " Create a bootable ISO now? (y/n)")) {
                    Write-Host "`n Add autounattend.xml or any other custom files to $OutputDir" -ForegroundColor Red
                    Read-Host " Press ENTER after you have finished adding files"

                    $IsoPath = ConvertTo-CleanInputPath (Read-Host " Enter ISO output path [C:\CustomWin11.iso]")
                    if (-not $IsoPath) {
                        $IsoPath = "C:\CustomWin11.iso"
                    }
                    $ResultLocation = New-BootableIso -SourceFolder $OutputDir -OutputPath $IsoPath
                }
                else {
                    Write-Host " ISO creation skipped. Add your files and use option 6 when ready." -ForegroundColor DarkGray
                }
            }
            "2" {
                $ExistingInput = ConvertTo-CleanInputPath (Read-Host "`n Enter existing media folder or install.wim/install.esd path")
                $SourceImage = Resolve-InstallImagePath $ExistingInput
                $DriverDir = Get-OptionalDriverPath
                $WorkingWim = Export-WindowsProImage -SourceImage $SourceImage

                if ($DriverDir) {
                    Add-WindowsDrivers -WimPath $WorkingWim -DriverPath $DriverDir
                }
                else {
                    Write-Host "`n Skipping driver injection (no drivers specified)." -ForegroundColor DarkGray
                }

                if (Test-Yes (Read-Host "`n Split the WIM into SWM files? (y/n)")) {
                    Split-WindowsImage -WimPath $WorkingWim | Out-Null
                }
                else {
                    Write-Host " Skipping WIM splitting." -ForegroundColor DarkGray
                }
                $ResultLocation = Split-Path -Parent $WorkingWim
            }
            "3" {
                $ExistingInput = ConvertTo-CleanInputPath (Read-Host "`n Enter install.wim path or media folder")
                $WorkingWim = Resolve-WimPath $ExistingInput
                $DriverDir = ConvertTo-CleanInputPath (Read-Host " Enter extracted (.inf) drivers folder")
                if (-not $DriverDir) {
                    throw "A drivers folder is required for the driver-injection stage."
                }
                Add-WindowsDrivers -WimPath $WorkingWim -DriverPath $DriverDir

                if (Test-Yes (Read-Host "`n Split the updated WIM into SWM files now? (y/n)")) {
                    Split-WindowsImage -WimPath $WorkingWim | Out-Null
                }
                $ResultLocation = Split-Path -Parent $WorkingWim
            }
            "4" {
                $MediaFolder = ConvertTo-CleanInputPath (Read-Host "`n Enter prepared Windows media folder [$OutputDir]")
                if (-not $MediaFolder) {
                    $MediaFolder = $OutputDir
                }
                $IsoPath = ConvertTo-CleanInputPath (Read-Host " Enter ISO output path [C:\CustomWin11.iso]")
                if (-not $IsoPath) {
                    $IsoPath = "C:\CustomWin11.iso"
                }
                $ResultLocation = New-BootableIso -SourceFolder $MediaFolder -OutputPath $IsoPath
            }
            "5" {
                $ExistingInput = ConvertTo-CleanInputPath (Read-Host "`n Enter install.wim path or media folder")
                $WorkingWim = Resolve-WimPath $ExistingInput
                Split-WindowsImage -WimPath $WorkingWim | Out-Null
                $ResultLocation = Split-Path -Parent $WorkingWim
            }
            "6" {
                $MediaFolder = ConvertTo-CleanInputPath (Read-Host "`n Enter prepared Windows media folder [$OutputDir]")
                if (-not $MediaFolder) {
                    $MediaFolder = $OutputDir
                }
                $ResultLocation = New-BootableUsbDrive -SourceFolder $MediaFolder
            }
            default {
                throw "Invalid stage option '$StageChoice'. Choose a number from 1 through 6."
            }
        }

        Write-Host "`n┌────────────────────────────────────────────────────────┐" -ForegroundColor Green
        Write-Host "│ SUCCESS! │" -ForegroundColor Green
        Write-Host "└────────────────────────────────────────────────────────┘" -ForegroundColor Green
        Write-Host " Result Location: $ResultLocation" -ForegroundColor Green
        if ($StageChoice -eq "1" -and $ResultLocation -eq $OutputDir) {
            Write-Host "`n [!] NEXT STEPS:" -ForegroundColor Yellow
            Write-Host " You can copy autounattend.xml and supporting files into" -ForegroundColor Gray
            Write-Host " '$OutputDir', then run option 4 to create an ISO or option 6 to copy it to USB." -ForegroundColor Gray
        }
        Write-Host ""
    }
    catch {
        Write-Host "`n [!] An error occurred during processing: $($_.Exception.Message)" -ForegroundColor Red
        throw
    }
    finally {
        if (Test-Path -LiteralPath $MountDir) {
            $MountedCheck = Get-WindowsImage -Mounted |
            Where-Object { $_.Path -eq $MountDir }
            if ($MountedCheck) {
                Write-Host " Discarding the uncommitted image mount..." -ForegroundColor DarkGray
                & dism.exe /Unmount-Image /MountDir:"$MountDir" /Discard | Out-Null
            }
        }

        if ($ScriptMountedIso) {
            Dismount-DiskImage -ImagePath $ScriptMountedIso -ErrorAction SilentlyContinue | Out-Null
        }

        if (Test-Path -LiteralPath $SessionWorkDir) {
            Write-Host " Cleaning up this run's temporary files..." -ForegroundColor DarkGray
            Remove-Item -LiteralPath $SessionWorkDir -Recurse -Force -ErrorAction SilentlyContinue
        }
    }
}

Export-ModuleMember -Function 'JRE-ImageWin11Help', 'JRE-ImageWin11Pro', 'JRE-NewNinjaOneInstaller'