# TinyWindowsMaker.ps1 - Universal Windows Image Creator
# Automatically detects Windows version from ISO and runs the appropriate maker script
# or allows manual selection between Windows 10 and Windows 11

param (
    [ValidatePattern('^[c-zC-Z]:?$|^[a-zA-Z]:\\.*$')]
    [string]$ScratchDisk,
    [string]$windowsisopath,
    [string]$imageindex,
    [ValidateSet("10", "11", "auto")]
    [string]$WindowsVersion = "auto",
    [switch]$UseSetupTemplate
)

# Check if PowerShell execution is Restricted or AllSigned or Undefined
$needchange = @("AllSigned", "Restricted", "Undefined")
$curpolicy = Get-ExecutionPolicy
if ($curpolicy -in $needchange) {
    Write-Host "Your current PowerShell Execution Policy is set to $curpolicy, which prevents scripts from running. Do you want to change it to RemoteSigned? (yes/no)"
    $response = Read-Host
    if ($response -eq 'yes') {
        Set-ExecutionPolicy RemoteSigned -Scope Process -Confirm:$false
    }
    else {
        Write-Host "The script cannot be run without changing the execution policy. Exiting..."
        exit
    }
}

# Check and run the script as admin if required
$myWindowsID = [System.Security.Principal.WindowsIdentity]::GetCurrent()
$myWindowsPrincipal = new-object System.Security.Principal.WindowsPrincipal($myWindowsID)
$adminRole = [System.Security.Principal.WindowsBuiltInRole]::Administrator
if (! $myWindowsPrincipal.IsInRole($adminRole)) {
    Write-Host "Restarting TinyWindowsMaker as admin in a new window, you can close this one."
    $newProcess = new-object System.Diagnostics.ProcessStartInfo "PowerShell";
    $argString = "-NoProfile -ExecutionPolicy Bypass -File `"$PSCommandPath`""
    # Add additional parameters if they are set
    if ($ScratchDisk) {
        $argString += " -ScratchDisk `"$ScratchDisk`""
    }
    if ($windowsisopath) {
        $argString += " -windowsisopath `"$windowsisopath`""
    }
    if ($imageindex) {
        $argString += " -imageindex `"$imageindex`""
    }
    if ($WindowsVersion -ne "auto") {
        $argString += " -WindowsVersion `"$WindowsVersion`""
    }
    if ($UseSetupTemplate) { $argString += " -UseSetupTemplate" }
    $newProcess.Arguments = $argString;
    $newProcess.Verb = "runas";
    [System.Diagnostics.Process]::Start($newProcess);
    exit
}

$Host.UI.RawUI.WindowTitle = "TinyWindowsMaker - Universal Windows Image Creator"
Clear-Host

Write-Host "==================================="
Write-Host "     TinyWindowsMaker v1.0"
Write-Host "   Universal Windows Image Creator"
Write-Host "==================================="
Write-Host ""

# Function to detect Windows version from ISO
function Get-WindowsVersionFromISO {
    param([string]$IsoPath)
    
    try {
        Write-Host "Mounting ISO to analyze Windows version..."
        
        # Mount the ISO and get the drive letter
        $mountResult = Mount-DiskImage -ImagePath $IsoPath -PassThru
        $driveLetter = ($mountResult | Get-Volume).DriveLetter + ":"
        
        Write-Host "ISO mounted at drive $driveLetter"
        
        # Try to get version info from install.wim or install.esd
        $installWim = "$driveLetter\sources\install.wim"
        $installEsd = "$driveLetter\sources\install.esd"
        
        $imagePath = $null
        if (Test-Path $installWim) {
            $imagePath = $installWim
            Write-Host "Found install.wim, analyzing..."
        }
        elseif (Test-Path $installEsd) {
            $imagePath = $installEsd
            Write-Host "Found install.esd, analyzing..."
        }
        
        if ($imagePath) {
            $imageInfo = Get-WindowsImage -ImagePath $imagePath -Index 1
            $version = $imageInfo.Version
            $imageName = $imageInfo.ImageName
            
            Write-Host "Image Name: $imageName"
            Write-Host "Version: $version"
            
            # Parse version to determine Windows 10 or 11
            if ($version) {
                $versionParts = $version.Split('.')
                if ($versionParts.Count -ge 3) {
                    $buildNumber = [int]$versionParts[2]
                    Write-Host "Build Number: $buildNumber"
                    
                    # Check for unsupported Windows versions (older than Windows 10)
                    if ($buildNumber -lt 10240) {
                        $majorVersion = [int]$versionParts[0]
                        $minorVersion = [int]$versionParts[1]
                        
                        # Determine the Windows version name
                        $windowsVersionName = "Unknown"
                        if ($majorVersion -eq 6) {
                            switch ($minorVersion) {
                                3 { $windowsVersionName = "Windows 8.1" }
                                2 { $windowsVersionName = "Windows 8" }
                                1 { $windowsVersionName = "Windows 7" }
                                0 { $windowsVersionName = "Windows Vista" }
                            }
                        }
                        elseif ($majorVersion -eq 5) {
                            switch ($minorVersion) {
                                2 { $windowsVersionName = "Windows XP (64-bit) / Windows Server 2003" }
                                1 { $windowsVersionName = "Windows XP" }
                                0 { $windowsVersionName = "Windows 2000" }
                            }
                        }
                        elseif ($majorVersion -lt 5) {
                            $windowsVersionName = "Windows 98/ME or older"
                        }
                        
                        Write-Host ""
                        Write-Host "=========================================="
                        Write-Host "         UNSUPPORTED WINDOWS VERSION"
                        Write-Host "=========================================="
                        Write-Host "Detected: $windowsVersionName (Build $buildNumber)"
                        Write-Host ""
                        Write-Host "This tool only supports Windows 10 and Windows 11."
                        Write-Host "Windows versions older than Windows 10 are not supported."
                        Write-Host ""
                        Write-Host "Supported versions:"
                        Write-Host "- Windows 10 (Build 10240 and newer)"
                        Write-Host "- Windows 11 (Build 22000 and newer)"
                        Write-Host "=========================================="
                        
                        return @{
                            Version = "unsupported"
                            DriveLetter = $driveLetter
                            BuildNumber = $buildNumber
                            ImageName = $imageName
                            WindowsVersionName = $windowsVersionName
                        }
                    }
                    # Windows 11 starts from build 22000
                    elseif ($buildNumber -ge 22000) {
                        Write-Host "Detected Windows 11 based on build number"
                        return @{
                            Version = "11"
                            DriveLetter = $driveLetter
                            BuildNumber = $buildNumber
                            ImageName = $imageName
                        }
                    }
                    else {
                        Write-Host "Detected Windows 10 based on build number"
                        return @{
                            Version = "10"
                            DriveLetter = $driveLetter
                            BuildNumber = $buildNumber
                            ImageName = $imageName
                        }
                    }
                }
            }
            
            # Fallback: check image name for version hints
            $imageNameLower = $imageName.ToLower()
            if ($imageNameLower -match "11" -or $imageNameLower -match "eleven") {
                Write-Host "Detected Windows 11 based on image name"
                return @{
                    Version = "11"
                    DriveLetter = $driveLetter
                    BuildNumber = "Unknown"
                    ImageName = $imageName
                }
            }
            elseif ($imageNameLower -match "10" -or $imageNameLower -match "ten") {
                Write-Host "Detected Windows 10 based on image name"
                return @{
                    Version = "10"
                    DriveLetter = $driveLetter
                    BuildNumber = "Unknown"
                    ImageName = $imageName
                }
            }
        }
        
        return @{
            Version = $null
            DriveLetter = $driveLetter
            BuildNumber = "Unknown"
            ImageName = "Unknown"
        }
    }
    catch {
        Write-Host "Warning: Could not analyze ISO file - $($_.Exception.Message)"
        return $null
    }
}

# Function to get source path from user (ISO, WIM, ESD, or drive letter)
function Get-IsoPath {
    do {
        $srcPath = Read-Host "Enter path to Windows ISO/WIM/ESD file, or a mounted drive letter (e.g. C:\Win11.iso, D:)"

        if ([string]::IsNullOrEmpty($srcPath)) {
            Write-Host "Please enter a valid path."
            continue
        }

        # Drive letter shorthand
        if ($srcPath -match '^[c-zC-Z]:?$') { return $srcPath.TrimEnd('\') + ":" }

        if (-not (Test-Path $srcPath)) {
            Write-Host "Path not found: $srcPath"
            continue
        }

        $ext = [System.IO.Path]::GetExtension($srcPath).ToLower()
        if ($ext -notin @('.iso', '.wim', '.esd')) {
            Write-Host "Unsupported file type '$ext'. Accepted: .iso, .wim, .esd"
            continue
        }

        return $srcPath
    } while ($true)
}

# Helper: detect Windows version from a WIM or ESD file directly (no mount needed)
function Get-WindowsVersionFromWimEsd {
    param([string]$FilePath)
    try {
        $allIndexes = Get-WindowsImage -ImagePath $FilePath
        # WOR-format ESDs have >=4 indexes; OS edition is the last one.
        # Install-only WIMs/ESDs use index 1.
        $targetIdx = if ($allIndexes.Count -ge 4) { $allIndexes[-1].ImageIndex } else { 1 }
        $imgInfo   = Get-WindowsImage -ImagePath $FilePath -Index $targetIdx
        $ver       = $imgInfo.Version
        $name      = $imgInfo.ImageName

        if ($ver) {
            $parts = $ver.Split('.')
            if ($parts.Count -ge 3) {
                $build  = [int]$parts[2]
                $winVer = if ($build -lt 10240) { "unsupported" } elseif ($build -ge 22000) { "11" } else { "10" }
                return @{ Version = $winVer; BuildNumber = $build; ImageName = $name; DriveLetter = $FilePath }
            }
        }
        return @{ Version = $null; BuildNumber = "Unknown"; ImageName = $name; DriveLetter = $FilePath }
    }
    catch {
        Write-Host "Warning: Could not read WIM/ESD: $($_.Exception.Message)"
        return $null
    }
}

# Determine the source and mount it if needed
$isoPath = $null
$isoInfo = $null

# Normalize windowsisopath: accept D, D:, D:\ all as drive letter D:
if ($windowsisopath -match '^[a-zA-Z]:?\\?$') {
    $windowsisopath = $windowsisopath[0] + ':'
}

# Resolve input: parameter, or ask user
$sourcePath = if ($windowsisopath) { $windowsisopath } else { Get-IsoPath }

$ext = [System.IO.Path]::GetExtension($sourcePath).ToLower()

if ($sourcePath -match '^[c-zC-Z]:$') {
    # Already-mounted drive letter
    Write-Host "Drive letter provided: $sourcePath"
    $hasInstallWim = Test-Path "$sourcePath\sources\install.wim"
    $hasInstallEsd = Test-Path "$sourcePath\sources\install.esd"
    $hasBootWim    = Test-Path "$sourcePath\sources\boot.wim"

    # Multi-arch ISO: no sources\ at root, but x64\sources\ present — use x64 for detection
    $versionSrcBase = $sourcePath
    if (-not $hasInstallWim -and -not $hasInstallEsd -and (Test-Path "$sourcePath\x64\sources")) {
        Write-Host "Multi-architecture ISO detected. Using x64."
        $versionSrcBase = "$sourcePath\x64"
        $hasInstallWim  = Test-Path "$versionSrcBase\sources\install.wim"
        $hasInstallEsd  = Test-Path "$versionSrcBase\sources\install.esd"
        $hasBootWim     = Test-Path "$versionSrcBase\sources\boot.wim"
    }

    if (-not $hasInstallWim -and -not $hasInstallEsd) {
        Write-Host "Error: Windows installation files not found in $sourcePath\sources\"
        Write-Host "Expected install.wim or install.esd"
        Write-Host ""
        Write-Host "Root of ${sourcePath}:"
        Get-ChildItem "$sourcePath\" -ErrorAction SilentlyContinue | ForEach-Object { Write-Host "  $($_.Name)" }
        Write-Host ""
        Write-Host "Contents of ${sourcePath}\sources (if it exists):"
        Get-ChildItem "$sourcePath\sources\" -ErrorAction SilentlyContinue | ForEach-Object { Write-Host "  $($_.Name)" }
        Read-Host "Press Enter to exit"
        exit
    }

    if (-not $hasBootWim) {
        if ($hasInstallEsd -and -not $hasInstallWim) {
            Write-Host "Note: boot.wim not found. ESD-only media - maker script will build boot structure."
        } else {
            Write-Host "Error: boot.wim not found in $versionSrcBase\sources\"
            Read-Host "Press Enter to exit"
            exit
        }
    }

    $isoInfo = @{ Version = $null; DriveLetter = $sourcePath; BuildNumber = "Unknown"; ImageName = "Unknown" }

    if ($WindowsVersion -eq "auto") {
        try {
            $imgSrc = if (Test-Path "$versionSrcBase\sources\install.wim") { "$versionSrcBase\sources\install.wim" }
                      else { "$versionSrcBase\sources\install.esd" }
            $imgDetail = Get-WindowsImage -ImagePath $imgSrc -Index 1
            $parts = $imgDetail.Version.Split('.')
            if ($parts.Count -ge 3) {
                $build = [int]$parts[2]
                $isoInfo.Version     = if ($build -lt 10240) { "unsupported" } elseif ($build -ge 22000) { "11" } else { "10" }
                $isoInfo.BuildNumber = $build
                $isoInfo.ImageName   = $imgDetail.ImageName
            }
        }
        catch { Write-Host "Warning: Could not auto-detect version from drive." }
    }
}
elseif ($ext -in @('.wim', '.esd')) {
    # Direct WIM/ESD - pass the file path straight to the maker script
    if (-not (Test-Path $sourcePath)) {
        Write-Host "Error: File not found: $sourcePath"
        Read-Host "Press Enter to exit"
        exit
    }
    Write-Host "WIM/ESD source detected: $sourcePath"
    $isoInfo = Get-WindowsVersionFromWimEsd -FilePath (Resolve-Path $sourcePath).Path
    if (-not $isoInfo) {
        Write-Host "Error: Could not read WIM/ESD file."
        Read-Host "Press Enter to exit"
        exit
    }
}
else {
    # ISO file - mount it and detect version
    if (-not (Test-Path $sourcePath)) {
        Write-Host "Error: File not found: $sourcePath"
        Read-Host "Press Enter to exit"
        exit
    }
    $isoPath = $sourcePath
    $isoInfo = Get-WindowsVersionFromISO -IsoPath $isoPath
    if (-not $isoInfo) {
        Write-Host "Error: Could not mount or analyze the ISO."
        Read-Host "Press Enter to exit"
        exit
    }
}

# Determine which script to run
$scriptToRun = $null
$detectedVersion = $null

if ($WindowsVersion -eq "auto") {
    if ($isoInfo.Version -eq "unsupported") {
        Write-Host ""
        Write-Host "=========================================="
        Write-Host "         UNSUPPORTED WINDOWS VERSION"
        Write-Host "=========================================="
        Write-Host "Build: $($isoInfo.BuildNumber)"
        Write-Host "Image: $($isoInfo.ImageName)"
        Write-Host ""
        Write-Host "This tool only supports Windows 10 and Windows 11."
        Write-Host "Windows versions older than Windows 10 are not supported."
        Write-Host ""
        Write-Host "Supported versions:"
        Write-Host "- Windows 10 (Build 10240 and newer)"
        Write-Host "- Windows 11 (Build 22000 and newer)"
        Write-Host "=========================================="
        Write-Host ""
        Read-Host "Press Enter to exit"
        exit
    }
    elseif ($isoInfo.Version) {
        Write-Host "Detected Windows $($isoInfo.Version)"
        Write-Host "Build: $($isoInfo.BuildNumber)"
        Write-Host "Image: $($isoInfo.ImageName)"
        $detectedVersion = $isoInfo.Version
        $scriptToRun = "tiny$detectedVersion" + "maker.ps1"
    }
    else {
        Write-Host "Could not automatically detect Windows version."
        Write-Host ""
        Write-Host "Please select the Windows version:"
        Write-Host "1. Windows 10"
        Write-Host "2. Windows 11"
        
        do {
            $choice = Read-Host "Enter your choice (1 or 2)"
            switch ($choice) {
                "1" { 
                    $scriptToRun = "tiny10maker.ps1"
                    $detectedVersion = "10"
                }
                "2" { 
                    $scriptToRun = "tiny11maker.ps1"
                    $detectedVersion = "11"
                }
                default {
                    Write-Host "Invalid choice. Please enter 1 or 2."
                }
            }
        } while ($choice -notin @("1", "2"))
    }
}
else {
    Write-Host "Manual Windows version specified: Windows $WindowsVersion"
    $scriptToRun = "tiny$WindowsVersion" + "maker.ps1"
    $detectedVersion = $WindowsVersion
}

# Verify the target script exists
$scriptPath = Join-Path $PSScriptRoot $scriptToRun
if (-not (Test-Path $scriptPath)) {
    Write-Host "Error: Cannot find $scriptToRun in $PSScriptRoot"
    Write-Host "Please make sure the script file exists."
    Read-Host "Press Enter to exit"
    exit
}

Write-Host ""
Write-Host "===========================================" 
Write-Host "Starting Tiny Windows $detectedVersion creation process..."
Write-Host "Using script: $scriptToRun"
Write-Host "===========================================" 
Write-Host ""

# Build the argument string for the target script
$argumentList = @()
$argumentList += "-NoProfile"
$argumentList += "-ExecutionPolicy"
$argumentList += "Bypass"
$argumentList += "-File"
$argumentList += "`"$scriptPath`""

if ($ScratchDisk) {
    $argumentList += "-ScratchDisk"
    $argumentList += "`"$ScratchDisk`""
}
# Pass the source path - drive letter, ISO path, or WIM/ESD path.
# The maker scripts resolve all three formats.
$argumentList += "-windowsisopath"
$argumentList += "`"$($isoInfo.DriveLetter)`""

if ($imageindex) {
    $argumentList += "-imageindex"
    $argumentList += "`"$imageindex`""
}
if ($UseSetupTemplate) {
    $argumentList += "-UseSetupTemplate"
}

# Start the appropriate maker script
try {
    $process = Start-Process -FilePath "PowerShell" -ArgumentList $argumentList -Wait -PassThru -NoNewWindow
    
    if ($process.ExitCode -eq 0) {
        Write-Host ""
        Write-Host "===========================================" 
        Write-Host "Tiny Windows $detectedVersion creation completed successfully!"
        Write-Host "===========================================" 
    }
    else {
        Write-Host ""
        Write-Host "===========================================" 
        Write-Host "Tiny Windows $detectedVersion creation failed with exit code: $($process.ExitCode)"
        Write-Host "===========================================" 
    }
}
catch {
    Write-Host "Error running $scriptToRun`: $($_.Exception.Message)"
}
finally {
    # Unmount the ISO if we mounted it
    if ($isoPath) {
        try {
            Write-Host "Unmounting ISO..."
            Dismount-DiskImage -ImagePath $isoPath | Out-Null
            Write-Host "ISO unmounted successfully."
        }
        catch {
            Write-Host "Warning: Could not unmount ISO - $($_.Exception.Message)"
        }
    }
}

Write-Host ""
Read-Host "Press Enter to exit"