# Enable debugging if needed
# Set-PSDebug -Trace 1

param (
    [ValidatePattern('^[c-zC-Z]$')]
    [string]$ScratchDisk
)

# Determine scratch disk location
if (-not $ScratchDisk) {
    $ScratchDisk = $PSScriptRoot -replace '[\\]+$', ''
} else {
    $ScratchDisk = "$ScratchDisk:"
}

Write-Output "Scratch disk set to $ScratchDisk"

# Ensure script is running with proper execution policy
if ((Get-ExecutionPolicy) -eq 'Restricted') {
    Write-Host "Your current PowerShell Execution Policy is 'Restricted'. This prevents scripts from running."
    $response = Read-Host "Do you want to change it to 'RemoteSigned'? (yes/no)"
    if ($response -eq 'yes') {
        Set-ExecutionPolicy RemoteSigned -Scope CurrentUser -Confirm:$false
    } else {
        Write-Host "The script cannot proceed without changing the execution policy. Exiting..."
        exit
    }
}

# Ensure script is running as administrator
$adminRole = [System.Security.Principal.WindowsBuiltInRole]::Administrator
$myWindowsPrincipal = New-Object System.Security.Principal.WindowsPrincipal([System.Security.Principal.WindowsIdentity]::GetCurrent())

if (-not $myWindowsPrincipal.IsInRole($adminRole)) {
    Write-Host "Restarting the script as administrator..."
    $newProcess = New-Object System.Diagnostics.ProcessStartInfo "PowerShell"
    $newProcess.Arguments = "-File `"$($MyInvocation.MyCommand.Definition)`""
    $newProcess.Verb = "runas"
    [System.Diagnostics.Process]::Start($newProcess)
    exit
}

# Start logging
Start-Transcript -Path "$ScratchDisk\tiny11.log"

# Set window title
$Host.UI.RawUI.WindowTitle = "Tiny11 Image Creator"
Clear-Host
Write-Host "Welcome to Tiny11 Image Creator! Release: 05-06-24"

# Ensure necessary directories exist
New-Item -ItemType Directory -Force -Path "$ScratchDisk\tiny11\sources" | Out-Null

# Prompt for Windows 11 installation media drive letter
do {
    $DriveLetter = Read-Host "Enter the drive letter of the Windows 11 installation media"
    if ($DriveLetter -match '^[c-zC-Z]$') {
        $DriveLetter = "$DriveLetter:"
        Write-Output "Drive letter set to $DriveLetter"
    } else {
        Write-Output "Invalid drive letter. Enter a letter between C and Z."
    }
} while ($DriveLetter -notmatch '^[c-zC-Z]:$')

# Validate Windows installation files
if (-not (Test-Path "$DriveLetter\sources\boot.wim") -or -not (Test-Path "$DriveLetter\sources\install.wim")) {
    if (Test-Path "$DriveLetter\sources\install.esd") {
        Write-Host "install.esd found. Converting to install.wim..."
        Get-WindowsImage -ImagePath "$DriveLetter\sources\install.esd"
        $index = Read-Host "Enter the image index"
        Write-Host "Converting install.esd to install.wim..."
        Export-WindowsImage -SourceImagePath "$DriveLetter\sources\install.esd" `
                            -SourceIndex $index `
                            -DestinationImagePath "$ScratchDisk\tiny11\sources\install.wim" `
                            -CompressionType Maximum -CheckIntegrity
    } else {
        Write-Host "Windows installation files not found. Please provide the correct drive letter."
        exit
    }
}

# Copy Windows installation files
Write-Host "Copying Windows installation files..."
Copy-Item -Path "$DriveLetter\*" -Destination "$ScratchDisk\tiny11" -Recurse -Force | Out-Null
Remove-Item "$ScratchDisk\tiny11\sources\install.esd" -ErrorAction SilentlyContinue
Write-Host "Copy complete!"

Start-Sleep -Seconds 2
Clear-Host
Write-Host "Retrieving image information..."
Get-WindowsImage -ImagePath "$ScratchDisk\tiny11\sources\install.wim"

# Prompt for image index
$index = Read-Host "Enter the image index"

# Mount Windows image
Write-Host "Mounting Windows image... This may take a while."
New-Item -ItemType Directory -Force -Path "$ScratchDisk\scratchdir" | Out-Null
Mount-WindowsImage -ImagePath "$ScratchDisk\tiny11\sources\install.wim" -Index $index -Path "$ScratchDisk\scratchdir"

# Retrieve system UI language
$imageIntl = & dism /English /Get-Intl "/Image:$($ScratchDisk)\scratchdir"
if ($imageIntl -match 'Default system UI language : ([a-zA-Z]{2}-[a-zA-Z]{2})') {
    Write-Host "Default system UI language: $($matches[1])"
} else {
    Write-Host "System UI language could not be determined."
}

# Retrieve architecture
$imageInfo = & dism /English /Get-WimInfo "/wimFile:$($ScratchDisk)\tiny11\sources\install.wim" "/index:$index"
if ($imageInfo -match "Architecture : (\S+)") {
    $architecture = $matches[1] -replace "x64", "amd64"
    Write-Host "Architecture: $architecture"
} else {
    Write-Host "Architecture information not found."
}

# Remove bloatware applications
Write-Host "Removing preinstalled apps..."
$appsToRemove = @(
    'Clipchamp.Clipchamp_', 'Microsoft.BingNews_', 'Microsoft.BingWeather_', 'Microsoft.GamingApp_',
    'Microsoft.GetHelp_', 'Microsoft.Getstarted_', 'Microsoft.MicrosoftOfficeHub_', 'Microsoft.MicrosoftSolitaireCollection_',
    'Microsoft.People_', 'Microsoft.PowerAutomateDesktop_', 'Microsoft.Todos_', 'Microsoft.WindowsAlarms_',
    'microsoft.windowscommunicationsapps_', 'Microsoft.WindowsFeedbackHub_', 'Microsoft.WindowsMaps_', 'Microsoft.WindowsSoundRecorder_',
    'Microsoft.Xbox.TCUI_', 'Microsoft.XboxGamingOverlay_', 'Microsoft.XboxGameOverlay_', 'Microsoft.XboxSpeechToTextOverlay_',
    'Microsoft.YourPhone_', 'Microsoft.ZuneMusic_', 'Microsoft.ZuneVideo_', 'MicrosoftCorporationII.MicrosoftFamily_',
    'MicrosoftCorporationII.QuickAssist_', 'MicrosoftTeams_', 'Microsoft.549981C3F5F10_'
)

$installedPackages = & dism /English "/image:$($ScratchDisk)\scratchdir" /Get-ProvisionedAppxPackages |
    ForEach-Object { if ($_ -match 'PackageName : (.*)') { $matches[1] } }

foreach ($package in $appsToRemove) {
    if ($installedPackages -match "$package*") {
        & dism /English "/image:$($ScratchDisk)\scratchdir" /Remove-ProvisionedAppxPackage "/PackageName:$matches[0]"
    }
}

# Remove Microsoft Edge
Write-Host "Removing Microsoft Edge..."
$edgeFolders = @(
    "Program Files (x86)\Microsoft\Edge",
    "Program Files (x86)\Microsoft\EdgeUpdate",
    "Program Files (x86)\Microsoft\EdgeCore"
)

foreach ($folder in $edgeFolders) {
    Remove-Item -Path "$ScratchDisk\scratchdir\$folder" -Recurse -Force -ErrorAction SilentlyContinue
}

# Remove OneDrive
Write-Host "Removing OneDrive..."
$oneDriveSetup = "$ScratchDisk\scratchdir\Windows\System32\OneDriveSetup.exe"
if (Test-Path $oneDriveSetup) {
    & takeown /f $oneDriveSetup | Out-Null
    & icacls $oneDriveSetup /grant "$($adminGroup.Value):(F)" /T /C | Out-Null
    Remove-Item -Path $oneDriveSetup -Force | Out-Null
}

Write-Host "Optimizations complete!"
Start-Sleep -Seconds 2
Clear-Host
Write-Host "Tiny11 Image preparation finished!"