# infra/install.ps1: put a Windows machine to work for Unitfabric, through WSL2 (Plan U2). The site serves it; in an
# administrator PowerShell:
#   $env:UNITFABRIC_SUPPLY_TOKEN = '<uf_supply_...>'; & ([scriptblock]::Create((irm https://unitfabric-site.vercel.app/install.ps1))) -ControlUrl <url>
# The supply token comes from UNITFABRIC_SUPPLY_TOKEN or, left out, is asked for (without echoing it): never a parameter,
# so that no process's command line carries it. The line above, as typed, stays in PowerShell's history for the token's
# day: the token enrols only machines its org then confirms, and the org revokes it at once (the Machines page, or
# `unitfabric supply tokens revoke`); leave out the line's first statement to be asked for the token instead.
# Windows itself runs no jobs: a WSL2 distribution of its own does, Unitfabric (Ubuntu 24.04, imported fresh from Ubuntu's
# image; the user's own distributions are never used or changed), with the GPU through the Windows NVIDIA driver (nothing
# is installed for it inside Linux), no checkpointing (a lost lease's units restart, so only jobs whose units are short
# come to it), no Windows drive mounted in it and no Windows program callable from it (automount and interop off), and its
# agent under systemd. This, idempotently:
#   1. checks the machine (-Check stops here, changing nothing);
#   2. installs WSL if it is missing (Windows then restarts once: run this again afterwards), and imports the distribution;
#   3. sizes the WSL2 VM in %UserProfile%\.wslconfig, adding only what it lacks (memory: three quarters of this machine's,
#      or -MemoryGB; swap -SwapGB), after a copy of it as it was (.wslconfig.unitfabric-<time>), in UTF-8; WSL takes them
#      at the VM's next start. It never stops WSL or another distribution (no wsl --shutdown, and wsl --update only while
#      none runs): its own alone is restarted, with wsl --terminate;
#   4. runs infra/install.sh inside it, as root, with the token in its environment (WSLENV), never on a command line, and
#      out of every environment once that is done, however it ended;
#   5. registers the scheduled task "Unitfabric agent (WSL)", at startup and at logon, which starts the distribution and
#      keeps it running, so its systemd starts the agent with Windows.
# It never exits: the one-liner runs it as a script block in the administrator's own PowerShell, which an `exit` would
# close, and every message with it. It returns, its status in $LASTEXITCODE: 0 done; 1 a step failed; 2 usage or not an
# administrator; 3 restart Windows and run this again.
[CmdletBinding()]
param(
    [string]$ControlUrl = $env:UF_CONTROL_URL,
    [string]$Name = $env:COMPUTERNAME,
    [string]$Rate = "",
    [int]$MemoryGB = 0,
    [int]$SwapGB = 8,
    [string]$SiteUrl = "https://unitfabric-site.vercel.app",
    [switch]$Check
)
$ErrorActionPreference = "Stop"
$env:WSL_UTF8 = "1"                                  # wsl.exe answers in UTF-8, not UTF-16, so its lists can be read
$TaskName = "Unitfabric agent (WSL)"
$Distro = "Unitfabric"
$RootfsUrl = "https://cloud-images.ubuntu.com/wsl/releases/24.04/20240423/ubuntu-noble-wsl-amd64-wsl.rootfs.tar.gz"
$RootfsSha256 = "8251e27ffff381a4af5f41dcb94d867de3e0d9774a9241908ab34555d99315ea"     # as that release's SHA256SUMS has it
$WslConf = "[boot]`nsystemd=true`n[automount]`nenabled=false`nmountFsTab=false`n[interop]`nenabled=false`nappendWindowsPath=false`n[user]`ndefault=root`n"

function Fail([string]$Message, [int]$Code = 1) { Write-Host "install.ps1: $Message" -ForegroundColor Red; $global:LASTEXITCODE = $Code }
function Report([string]$Level, [string]$What, [string]$Detail) { Write-Host ("{0,-5} {1,-12} {2}" -f $Level, $What, $Detail) }

function Quiet([scriptblock]$Command) {
    # A native command run for its status alone ($LASTEXITCODE), its output dropped. Once its stderr is redirected, Windows
    # PowerShell 5.1 makes each line of it an error, which ErrorActionPreference Stop stops at: here they are only dropped.
    $eap = $ErrorActionPreference; $ErrorActionPreference = "Continue"
    try { & $Command *> $null } finally { $ErrorActionPreference = $eap }
}

function Lines([scriptblock]$Command) {
    # A native command's standard output, its stderr dropped as Quiet drops it.
    $eap = $ErrorActionPreference; $ErrorActionPreference = "Continue"
    try { & $Command 2>$null } finally { $ErrorActionPreference = $eap }
}

function Get-Distros([switch]$Running) {
    $out = if ($Running) { Lines { wsl.exe --list --running --quiet } } else { Lines { wsl.exe --list --quiet } }
    if ($LASTEXITCODE -ne 0) { return @() }          # none (wsl says so in words, and fails)
    return @($out | ForEach-Object { $_.Trim([char]0, ' ') } | Where-Object { $_ })
}

function Update-WslConfig([string[]]$Lines, [int]$MemoryGB, [int]$SwapGB) {
    # The lines of .wslconfig with what [wsl2] lacks of memory and swap added at its end (or the section added), the rest
    # as it was: a memory or a swap already there is the user's, for every distribution, and stays.
    $out = New-Object System.Collections.Generic.List[string]
    $want = [ordered]@{ memory = "$($MemoryGB)GB"; swap = "$($SwapGB)GB" }
    $inWsl2 = $false; $seen = $false; $have = @{}
    foreach ($line in $Lines) {
        if ($line -match '^\s*\[(.+)\]\s*$') {
            $section = $Matches[1].Trim()
            if ($inWsl2) { foreach ($k in $want.Keys) { if (-not $have[$k]) { $out.Add("$k=$($want[$k])") } } }
            $inWsl2 = $section -eq "wsl2"
            if ($inWsl2) { $seen = $true; $have = @{} }
            $out.Add($line); continue
        }
        if ($inWsl2 -and $line -match '^\s*(memory|swap)\s*=') { $have[$Matches[1].ToLower()] = $true }
        $out.Add($line)
    }
    if ($inWsl2) { foreach ($k in $want.Keys) { if (-not $have[$k]) { $out.Add("$k=$($want[$k])") } } }
    if (-not $seen) { $out.Add("[wsl2]"); foreach ($k in $want.Keys) { $out.Add("$k=$($want[$k])") } }
    return , $out.ToArray()
}

function Set-WslConfig([string]$Path, [int]$MemoryGB, [int]$SwapGB) {
    # .wslconfig with what it lacks added (Update-WslConfig): read in the encoding its byte order mark names, else UTF-8;
    # written in UTF-8 without one, after a copy of it as it was; left alone when it lacks nothing. Whether it changed.
    $lines = if (Test-Path -LiteralPath $Path) { [IO.File]::ReadAllLines($Path) } else { @() }
    $new = Update-WslConfig $lines $MemoryGB $SwapGB
    if (($new -join "`n") -ceq ($lines -join "`n")) { return $false }
    if (Test-Path -LiteralPath $Path) { Copy-Item -LiteralPath $Path -Destination "$Path.unitfabric-$(Get-Date -Format yyyyMMddHHmmss)" }
    [IO.File]::WriteAllLines($Path, [string[]]$new, (New-Object Text.UTF8Encoding $false))
    return $true
}

function Test-Wsl { Quiet { wsl.exe --status }; return $LASTEXITCODE -eq 0 }

function Invoke-Install {
    $admin = (New-Object Security.Principal.WindowsPrincipal([Security.Principal.WindowsIdentity]::GetCurrent())).IsInRole(
        [Security.Principal.WindowsBuiltInRole]::Administrator)
    if (-not $Check) {
        if (-not $admin) { Fail "run this in an administrator PowerShell" 2; return }
        if ($ControlUrl -notmatch '^https://[A-Za-z0-9.-]+(:[0-9]{1,5})?$') { Fail "-ControlUrl must be https://host[:port]" 2; return }
    }
    if ($SiteUrl -notmatch '^https://[A-Za-z0-9.-]+(:[0-9]{1,5})?$') { Fail "-SiteUrl must be https://host[:port]" 2; return }
    if ($Rate -and $Rate -notmatch '^[0-9]{1,4}(\.[0-9]{1,4})?$') { Fail "-Rate is USD an hour, at most four decimals" 2; return }
    $script:Name = ($Name -replace '[^A-Za-z0-9._ -]', '')
    if ($script:Name.Length -gt 64) { $script:Name = $script:Name.Substring(0, 64) }
    if (-not $script:Name) { $script:Name = "windows" }
    $token = $env:UNITFABRIC_SUPPLY_TOKEN            # out of this PowerShell's environment until install.sh is given it
    Remove-Item Env:UNITFABRIC_SUPPLY_TOKEN -ErrorAction SilentlyContinue
    if (-not $Check) {
        if (-not $token) {                           # asked for, not echoed, and never a parameter
            $secret = Read-Host -AsSecureString "Supply token (uf_supply_...)"
            $token = [Runtime.InteropServices.Marshal]::PtrToStringBSTR([Runtime.InteropServices.Marshal]::SecureStringToBSTR($secret))
        }
        if ($token -notmatch '^uf_supply_[0-9a-f]{40}$') { Fail "UNITFABRIC_SUPPLY_TOKEN must be a supply token, uf_supply_ and 40 hex digits" 2; return }
    }

    # --- What this machine is -----------------------------------------------------------------------------------------
    $build = [Environment]::OSVersion.Version.Build
    $memory = [math]::Floor((Get-CimInstance Win32_ComputerSystem).TotalPhysicalMemory / 1GB)
    $gpu = Get-CimInstance Win32_VideoController | Where-Object { $_.Name -like "*NVIDIA*" } | Select-Object -First 1
    Write-Host "Unitfabric: checking $env:COMPUTERNAME (Windows build $build)"
    if ($admin) { Report "ok" "admin" "an administrator PowerShell" } else { Report "warn" "admin" "not an administrator: the install needs one" }
    if ($build -ge 22000) { Report "ok" "windows" "Windows 11 (build $build)" }
    elseif ($build -ge 19044) { Report "ok" "windows" "Windows 10 21H2 or newer (build $build): WSL2 with GPU works" }
    else { Report "FAIL" "windows" "build ${build}: WSL2 with a GPU needs Windows 10 21H2 (19044) or Windows 11"; if (-not $Check) { $global:LASTEXITCODE = 2; return } }
    if ($gpu) { Report "ok" "gpu" "$($gpu.Name), driver $($gpu.DriverVersion): WSL2 reaches it through this driver" }
    else { Report "warn" "gpu" "no NVIDIA GPU: the machine runs CPU templates only" }
    Report "ok" "memory" "$memory GB"
    $wsl = Test-Wsl
    if ($wsl) { Report "ok" "wsl" "installed" } else { Report "warn" "wsl" "not installed: the install adds it, and Windows restarts once" }
    $present = $wsl -and ((Get-Distros) -contains $Distro)
    if ($present) { Report "ok" "distro" "$Distro, its own" } else { Report "warn" "distro" "$Distro not there yet: the install imports Ubuntu 24.04 as it" }
    if ($Check) {
        if ($present) { & wsl.exe -d $Distro -u root -- bash -c "curl -fsSL $SiteUrl/install.sh | UNITFABRIC_SITE_URL=$SiteUrl bash -s -- --check" }
        Write-Host "check done; nothing was changed"
        $global:LASTEXITCODE = 0
        return
    }

    # --- WSL and the distribution -------------------------------------------------------------------------------------
    if (-not $wsl) {
        & wsl.exe --install --no-distribution
        if ($LASTEXITCODE -ne 0) { Fail "wsl --install failed (above)"; return }
        Write-Host "WSL is installed. Restart Windows, then run this again." -ForegroundColor Yellow
        $global:LASTEXITCODE = 3
        return
    }
    $running = Get-Distros -Running
    if ($running.Count -eq 0) { Quiet { wsl.exe --update } }       # an update stops WSL: only while nothing runs in it
    else { Report "warn" "wsl" "not updated while $($running -join ', ') run: wsl --update when they may stop" }
    if (-not $present) {
        $dir = Join-Path $env:LOCALAPPDATA "Unitfabric\wsl"
        New-Item -ItemType Directory -Force -Path $dir | Out-Null
        $tar = Join-Path $dir "ubuntu-24.04-wsl.rootfs.tar.gz"
        Write-Host "downloading Ubuntu 24.04 for WSL (about 360 MB): $RootfsUrl"
        & "$env:SystemRoot\System32\curl.exe" -fsSL -o $tar $RootfsUrl
        if ($LASTEXITCODE -ne 0) { Fail "the download of $RootfsUrl failed (above)"; return }
        if ((Get-FileHash -LiteralPath $tar -Algorithm SHA256).Hash.ToLower() -ne $RootfsSha256) {
            Remove-Item -LiteralPath $tar
            Fail "$RootfsUrl is not the image its sha256 names: refused"; return
        }
        & wsl.exe --import $Distro (Join-Path $dir "disk") $tar --version 2    # WSL2 named here: the default is left alone
        $imported = $LASTEXITCODE
        Remove-Item -LiteralPath $tar
        if ($imported -ne 0) { Fail "wsl --import $Distro failed (above)"; return }
        & wsl.exe -d $Distro -u root -- bash -c "export DEBIAN_FRONTEND=noninteractive; apt-get update -qq && apt-get full-upgrade -y -qq >/dev/null"
        if ($LASTEXITCODE -ne 0) { Fail "$Distro could not be brought up to date (above)"; return }
    }

    # --- Its wsl.conf, and the VM's size ------------------------------------------------------------------------------
    $conf = [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($WslConf))
    & wsl.exe -d $Distro -u root -- sh -c "echo $conf | base64 -d > /etc/wsl.conf.new; if cmp -s /etc/wsl.conf.new /etc/wsl.conf; then rm /etc/wsl.conf.new; else mv /etc/wsl.conf.new /etc/wsl.conf; exit 10; fi"
    if ($LASTEXITCODE -eq 10) { Quiet { wsl.exe --terminate $Distro } }      # read again when it next starts
    if ($MemoryGB -le 0) { $MemoryGB = [math]::Max(4, [math]::Floor($memory * 3 / 4)) }
    $wslconfig = Join-Path $env:USERPROFILE ".wslconfig"
    if (Set-WslConfig $wslconfig $MemoryGB $SwapGB) {
        Report "ok" "wslconfig" "${wslconfig}: what it lacked of memory and swap added, its old copy kept beside it; WSL takes it at its next start"
    }

    # --- install.sh inside it -----------------------------------------------------------------------------------------
    $env:UF_CONTROL_URL = $ControlUrl
    $env:UNITFABRIC_SUPPLY_TOKEN = $token
    $env:WSLENV = (@("UNITFABRIC_SUPPLY_TOKEN/u", "UF_CONTROL_URL/u") + @($script:SavedWslenv -split ':' | Where-Object { $_ })) -join ':'
    $extra = if ($Rate) { "--rate $Rate" } else { "" }
    $status = 1
    try {
        for ($try = 1; $try -le 2; $try++) {             # curl without the token: install.sh alone reads it, then unsets it
            & wsl.exe -d $Distro -u root -- bash -c "env -u UNITFABRIC_SUPPLY_TOKEN curl -fsSL $SiteUrl/install.sh | UNITFABRIC_SITE_URL=$SiteUrl bash -s -- --name '$script:Name' $extra"
            $status = $LASTEXITCODE
            if ($status -eq 3 -and $try -eq 1) { Quiet { wsl.exe --terminate $Distro }; continue }      # systemd now on
            break
        }
    } finally {
        Clear-Secrets
    }
    if ($status -ne 0) { Fail "the install inside $Distro failed (above)"; return }

    # --- Start with Windows -------------------------------------------------------------------------------------------
    $user = "$env:USERDOMAIN\$env:USERNAME"
    # A raw command line for wsl.exe, which Windows splits on spaces with no regard for single quotes: no quotes in it.
    $action = New-ScheduledTaskAction -Execute "$env:SystemRoot\System32\wsl.exe" -Argument "-d $Distro -u root --exec /bin/sleep infinity"
    $triggers = @((New-ScheduledTaskTrigger -AtStartup), (New-ScheduledTaskTrigger -AtLogOn -User $user))
    $principal = New-ScheduledTaskPrincipal -UserId $user -LogonType S4U -RunLevel Highest      # without a logon, and no stored password
    $settings = New-ScheduledTaskSettingsSet -ExecutionTimeLimit ([TimeSpan]::Zero) -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries `
        -RestartCount 3 -RestartInterval (New-TimeSpan -Minutes 1) -MultipleInstances IgnoreNew
    Register-ScheduledTask -TaskName $TaskName -Action $action -Trigger $triggers -Principal $principal -Settings $settings -Force | Out-Null
    Start-ScheduledTask -TaskName $TaskName
    Write-Host "done: $script:Name works through WSL2 ($Distro), started with Windows by the task '$TaskName'. Manage it on the dashboard's Machines page."
    $global:LASTEXITCODE = 0
}

function Clear-Secrets {
    # The token out of this PowerShell's environment and WSLENV as they were: after install.sh, and however the install ended.
    Remove-Item Env:UNITFABRIC_SUPPLY_TOKEN, Env:UF_CONTROL_URL -ErrorAction SilentlyContinue
    $env:WSLENV = $script:SavedWslenv
}

$script:SavedWslenv = $env:WSLENV
try {
    Invoke-Install
} catch {                                            # a step that threw: 1, and the window stays
    Write-Host "install.ps1: $_" -ForegroundColor Red
    $global:LASTEXITCODE = 1
} finally {
    Clear-Secrets
}