<#
.SYNOPSIS
    Enables all Windows features required for WSL 2 on Windows 11 and installs the latest Ubuntu distribution.

.DESCRIPTION
    This script must be run as Administrator. It will:
      1. Enable the "Microsoft-Windows-Subsystem-Linux" and "VirtualMachinePlatform" optional features.
      2. Set WSL 2 as the default WSL version.
      3. Install/update the WSL kernel components.
      4. Install the latest Ubuntu distribution from the Microsoft Store catalog (via `wsl --install -d Ubuntu`).

    A restart may be required after enabling the Windows features for the first time. If a restart is
    needed, the script will detect it, prompt you, and re-run automatically after reboot (via a scheduled
    task) unless you decline.

.NOTES
    Run this script from an elevated PowerShell prompt:
        powershell -ExecutionPolicy Bypass -File .\Install-WSL2Ubuntu.ps1
#>

[CmdletBinding()]
param(
    [string]$Distribution = "Ubuntu"
)

$ErrorActionPreference = "Stop"

function Test-IsAdmin {
    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = New-Object Security.Principal.WindowsPrincipal($identity)
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

function Write-Step {
    param([string]$Message)
    Write-Host "`n==> $Message" -ForegroundColor Cyan
}

if (-not (Test-IsAdmin)) {
    Write-Error "This script must be run as Administrator. Right-click PowerShell and choose 'Run as Administrator', then re-run this script."
    exit 1
}

# --- Step 1: Enable required Windows optional features ---
Write-Step "Enabling 'Windows Subsystem for Linux' feature"
$wslFeature = Enable-WindowsOptionalFeature -Online -FeatureName Microsoft-Windows-Subsystem-Linux -NoRestart

Write-Step "Enabling 'Virtual Machine Platform' feature"
$vmpFeature = Enable-WindowsOptionalFeature -Online -FeatureName VirtualMachinePlatform -NoRestart

$restartNeeded = $wslFeature.RestartNeeded -or $vmpFeature.RestartNeeded

if ($restartNeeded) {
    Write-Host "`nA restart is required to finish enabling the WSL 2 prerequisites." -ForegroundColor Yellow
    $answer = Read-Host "Restart now and continue installation automatically after reboot? (Y/N)"

    if ($answer -match '^[Yy]') {
        Write-Step "Scheduling this script to resume after reboot"

        $scriptPath = $MyInvocation.MyCommand.Path
        $taskName = "Resume-WSL2Install"

        $action = New-ScheduledTaskAction -Execute "powershell.exe" `
            -Argument "-NoProfile -ExecutionPolicy Bypass -File `"$scriptPath`" -Distribution `"$Distribution`""
        $trigger = New-ScheduledTaskTrigger -AtLogOn
        $principal = New-ScheduledTaskPrincipal -UserId "$env:USERDOMAIN\$env:USERNAME" -RunLevel Highest

        Register-ScheduledTask -TaskName $taskName -Action $action -Trigger $trigger -Principal $principal -Force | Out-Null

        Write-Host "Restarting in 10 seconds... (Ctrl+C to cancel)" -ForegroundColor Yellow
        Start-Sleep -Seconds 10
        Restart-Computer -Force
        exit 0
    }
    else {
        Write-Host "Please restart your computer manually, then re-run this script to continue." -ForegroundColor Yellow
        exit 0
    }
}

# If we get here after a reboot via the scheduled task, clean it up.
$existingTask = Get-ScheduledTask -TaskName "Resume-WSL2Install" -ErrorAction SilentlyContinue
if ($existingTask) {
    Unregister-ScheduledTask -TaskName "Resume-WSL2Install" -Confirm:$false
}

# --- Step 2: Update WSL and set default version to 2 ---
Write-Step "Updating the WSL kernel components"
wsl --update

Write-Step "Setting WSL default version to 2"
wsl --set-default-version 2

# --- Step 3: Install the latest Ubuntu distribution ---
Write-Step "Checking for an existing '$Distribution' installation"
$installedDistros = (wsl --list --quiet 2>$null) -replace "`0", ""

if ($installedDistros -match [regex]::Escape($Distribution)) {
    Write-Host "'$Distribution' is already installed. Skipping installation." -ForegroundColor Green
}
else {
    Write-Step "Installing '$Distribution' (latest version from Microsoft Store)"
    wsl --install -d $Distribution --no-launch
}

Write-Step "Done"
Write-Host "WSL 2 features are enabled and '$Distribution' has been installed." -ForegroundColor Green
Write-Host "Launch it by running: wsl -d $Distribution" -ForegroundColor Green
Write-Host "On first launch you'll be asked to create a Unix username and password." -ForegroundColor Green