<#
.SYNOPSIS
    Upgrades an already-enrolled Togen Windows Display Agent in place.

.DESCRIPTION
    Staged-swap upgrade, mirroring the Vision agent's DVI-920/DVI-1007
    pattern exactly (same gotchas apply to this agent):

      - DVI-920 gotcha: an AV/Defender file lock on the RUNNING exe makes a
        naive stop-then-copy silently fail. Fix: stage the new files under a
        *.new name first (no lock conflict, different filename), then do the
        actual stop/rename/restart from a separate task.
      - DVI-1007 gotcha: a SYSTEM/task-hosted process is inside a Job Object
        that does not set JOB_OBJECT_LIMIT_BREAKAWAY_OK, so if THIS script
        ever runs from within the display agent's own process tree (e.g. a
        future remote "update agent" command), calling
        Stop-ScheduledTask/Stop-Process on the agent's own task from here
        could kill this script mid-swap instead of the target. Fix: this
        script NEVER stops TogenDisplayAgent directly. It registers a
        SEPARATE one-shot Scheduled Task (TogenDisplayAgentSwap, its own Job
        Object) that fires ~30s later and performs the actual stop / file
        swap / restart, then self-unregisters. upgrade.ps1 only stages files
        and registers that task.

    Run as Administrator. No enrollment token is required -- the machine
    is already enrolled; this only replaces the agent binary/source.

.PARAMETER Server
    Base URL of the Togen server. When omitted, auto-detected from the
    existing TogenDisplayAgent Scheduled Task's action arguments (same
    fallback the Vision install.ps1 uses).

.PARAMETER PreDownloadedZip
    Optional path to an already-downloaded togen-display-agent.zip. When
    supplied, the download step is skipped entirely.

.PARAMETER InstallDir
    Installation directory. Defaults to %ProgramFiles%\Togen\Display.

.PARAMETER TaskName
    The logon-triggered kiosk task name. Defaults to TogenDisplayAgent.

.EXAMPLE
    .\upgrade.ps1 -Server https://togen.example.com

.EXAMPLE
    # Pre-downloaded ZIP (e.g. pushed by an RMM tool), no network needed here
    .\upgrade.ps1 -PreDownloadedZip C:\Temp\togen-display-agent.zip
#>
[CmdletBinding()]
param(
    [string]$Server = '',
    [string]$PreDownloadedZip = '',
    [string]$InstallDir = (Join-Path $env:ProgramFiles 'Togen\Display'),
    [string]$TaskName   = 'TogenDisplayAgent'
)

$ErrorActionPreference = 'Stop'
$SwapTaskName = "${TaskName}Swap"
$LogDir  = Join-Path $env:ProgramData 'Togen\Display\Logs'
$LogFile = Join-Path $LogDir 'agent.log'
New-Item -ItemType Directory -Path $LogDir -Force -ErrorAction SilentlyContinue | Out-Null

function Write-Step { param([string]$Msg) Write-Host ">> $Msg" -ForegroundColor Cyan; Write-Log "[upgrade] $Msg" }
function Write-OK   { param([string]$Msg) Write-Host "   OK: $Msg" -ForegroundColor Green; Write-Log "[upgrade] OK: $Msg" }
function Write-Warn { param([string]$Msg) Write-Host "   WARN: $Msg" -ForegroundColor Yellow; Write-Log "[upgrade] WARN: $Msg" }
function Write-Log  {
    param([string]$Msg)
    try {
        "$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss') $Msg" | Add-Content -Path $LogFile -ErrorAction SilentlyContinue
    } catch { }
}

# Runtime elevation check (warn only, matches the Vision agent's DVI-850
# fix -- #Requires -RunAsAdministrator is evaluated at PARSE time and would
# silently abort the whole script before any diagnostics if launched from a
# non-elevated subprocess).
try {
    $isAdmin = ([Security.Principal.WindowsPrincipal][Security.Principal.WindowsIdentity]::GetCurrent()).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
    if (-not $isAdmin) { Write-Warn 'Not running as Administrator -- some steps may fail' }
    else { Write-OK 'Running as Administrator' }
} catch { Write-Warn "Could not check elevation: $_" }

# ---------------------------------------------------------------------------
# 0. Auto-detect server URL from the existing task, if not given.
# ---------------------------------------------------------------------------
if (-not $Server) {
    Write-Step 'Auto-detecting server URL from the existing Scheduled Task'
    $existingTask = Get-ScheduledTask -TaskName $TaskName -ErrorAction SilentlyContinue
    if ($existingTask) {
        $runScriptPath = Join-Path $InstallDir 'run_agent.ps1'
        if (Test-Path $runScriptPath) {
            $runScriptContent = Get-Content $runScriptPath -Raw
            if ($runScriptContent -match "--server',\s*'([^']+)'") { $Server = $Matches[1] }
        }
    }
    if (-not $Server) {
        Write-Host 'ERROR: Cannot detect server URL. Pass -Server <url> explicitly.' -ForegroundColor Red
        exit 1
    }
    Write-OK "Server: $Server"
}

# ---------------------------------------------------------------------------
# 1. Verify existing installation
# ---------------------------------------------------------------------------
Write-Step 'Checking existing installation'
$DestExe = Join-Path $InstallDir 'display_agent.exe'
if (-not (Test-Path $DestExe)) {
    Write-Host 'ERROR: No existing Display Agent found. Run install.ps1 for a fresh install.' -ForegroundColor Red
    exit 1
}
Write-OK "Found: $DestExe"

# ---------------------------------------------------------------------------
# 2. Agent bundle ZIP (pre-downloaded, or fetched from the server now).
# ---------------------------------------------------------------------------
Write-Step 'Locating agent bundle'
$Zip = Join-Path $env:TEMP 'togen-display-agent-upgrade.zip'
$ZipOk = $false

if ($PreDownloadedZip -and (Test-Path $PreDownloadedZip)) {
    $Zip = $PreDownloadedZip
    $ZipOk = $true
    Write-OK "Using pre-downloaded ZIP: $Zip"
} else {
    Write-Step 'Downloading latest agent bundle'
    $AgentZipUrl = "$($Server.TrimEnd('/'))/enroll/display-agent.zip"
    Remove-Item $Zip -Force -ErrorAction SilentlyContinue
    try { [Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12 } catch { }
    for ($dlTry = 1; $dlTry -le 3; $dlTry++) {
        if ($dlTry -gt 1) { Write-Warn "Download retry $dlTry/3..."; Start-Sleep -Seconds 15 }
        try {
            Invoke-WebRequest -Uri $AgentZipUrl -OutFile $Zip -UseBasicParsing -TimeoutSec 300
            if ((Get-Item $Zip -ErrorAction SilentlyContinue).Length -gt 0) { $ZipOk = $true; break }
        } catch {
            Write-Warn "Download attempt $dlTry failed: $_"
        }
    }
    if ($ZipOk) { Write-OK "Downloaded: $Zip" }
}
if (-not $ZipOk) {
    Write-Host 'ERROR: Could not obtain the agent bundle ZIP.' -ForegroundColor Red
    exit 1
}

# ---------------------------------------------------------------------------
# 3. Extract and stage new files as *.new (no lock conflict with the running
#    exe -- the DVI-920 fix).
# ---------------------------------------------------------------------------
Write-Step 'Staging new agent files'
$ExtractDir = Join-Path $env:TEMP "togen-display-agent-extract-$(Get-Random)"
Remove-Item -Recurse -Force $ExtractDir -ErrorAction SilentlyContinue
Expand-Archive -Path $Zip -DestinationPath $ExtractDir -Force

$StagedFiles = @()
foreach ($name in @('display_agent.exe', 'display_agent.py', 'winlayer.py', 'server_client.py')) {
    $src = Join-Path $ExtractDir $name
    if (Test-Path $src) {
        $newDest = Join-Path $InstallDir "$name.new"
        Copy-Item -Path $src -Destination $newDest -Force
        $StagedFiles += [PSCustomObject]@{ Name = $name; NewPath = $newDest; FinalPath = (Join-Path $InstallDir $name) }
        Write-OK "Staged: $name.new"
    }
}
if (-not ($StagedFiles | Where-Object { $_.Name -eq 'display_agent.exe' })) {
    Write-Host 'ERROR: display_agent.exe not found in the downloaded bundle.' -ForegroundColor Red
    exit 1
}

# Refresh the installed swap_agent.ps1 itself from the bundle (it may have
# been updated too) -- this is the SAME file the server-rendered
# /enroll/display-upgrade.ps1 route looks for at $InstallDir\swap_agent.ps1
# and prefers over its own minimal single-exe fallback (see swap_agent.ps1's
# header comment). Keeping one shared, installed copy means both upgrade
# paths get the same manifest-aware, multi-file, lock-retrying swap logic.
$SwapScriptSrc = Join-Path $ExtractDir 'swap_agent.ps1'
$SwapScript = Join-Path $InstallDir 'swap_agent.ps1'
if (Test-Path $SwapScriptSrc) {
    Copy-Item -Path $SwapScriptSrc -Destination $SwapScript -Force
    Write-OK 'Refreshed installed swap_agent.ps1'
} elseif (-not (Test-Path $SwapScript)) {
    Write-Host 'ERROR: swap_agent.ps1 not found in the bundle and none is installed.' -ForegroundColor Red
    exit 1
}
Remove-Item -Recurse -Force $ExtractDir -ErrorAction SilentlyContinue

# Serialize the staged-file manifest for swap_agent.ps1 to consume (avoids
# re-deriving it and keeps the swap script itself simple/reliable).
$ManifestPath = Join-Path $env:ProgramData 'Togen\Display\swap_manifest.json'
New-Item -ItemType Directory -Path (Split-Path $ManifestPath) -Force -ErrorAction SilentlyContinue | Out-Null
$StagedFiles | ConvertTo-Json | Set-Content -Path $ManifestPath -Encoding UTF8

# ---------------------------------------------------------------------------
# 4. Register the SEPARATE Swap task (own Job Object -- DVI-1007). It fires
#    once, ~20s from now, runs the installed swap_agent.ps1 (stop/rename/
#    restart against the manifest above), then self-unregisters. upgrade.ps1
#    (this script) never touches TogenDisplayAgent directly.
# ---------------------------------------------------------------------------
Write-Step "Registering swap task: $SwapTaskName"
Unregister-ScheduledTask -TaskName $SwapTaskName -Confirm:$false -ErrorAction SilentlyContinue

$PsExe = 'C:\Windows\System32\WindowsPowerShell\v1.0\powershell.exe'
if (-not (Test-Path $PsExe)) { $PsExe = 'powershell.exe' }
$SwapAction  = New-ScheduledTaskAction -Execute $PsExe `
    -Argument "-NonInteractive -WindowStyle Hidden -ExecutionPolicy Bypass -File `"$SwapScript`"" `
    -WorkingDirectory $InstallDir
$SwapTrigger = New-ScheduledTaskTrigger -Once -At (Get-Date).AddSeconds(20)
$SwapSettings = New-ScheduledTaskSettingsSet -ExecutionTimeLimit (New-TimeSpan -Minutes 5) -MultipleInstances IgnoreNew
# SYSTEM: the swap needs to Stop-ScheduledTask/Stop-Process/rename files in
# Program Files regardless of whether the kiosk user is currently logged on.
$SwapPrincipal = New-ScheduledTaskPrincipal -UserID 'NT AUTHORITY\SYSTEM' -LogonType ServiceAccount -RunLevel Highest

$swapRegistered = $false
try {
    Register-ScheduledTask -TaskName $SwapTaskName -Trigger $SwapTrigger -Action $SwapAction `
        -Settings $SwapSettings -Principal $SwapPrincipal -Force | Out-Null
    $swapRegistered = $true
    Write-OK "Swap task '$SwapTaskName' registered -- fires in ~20s and self-unregisters"
} catch {
    Write-Warn "Could not register the swap task: $_"
}

# ---------------------------------------------------------------------------
# 5. Fallback: if the swap task could not be registered at all, run
#    swap_agent.ps1 directly and synchronously right now (best effort,
#    without the separate-Job-Object protection). Documented tradeoff,
#    matches the Vision agent's own inline-fallback precedent.
# ---------------------------------------------------------------------------
if (-not $swapRegistered) {
    Write-Warn 'Falling back to running swap_agent.ps1 directly (no separate-task '
    Write-Warn 'protection -- if this script is ever invoked from the agent''s own '
    Write-Warn 'process tree this path risks the DVI-1007 self-kill gotcha; prefer '
    Write-Warn 'fixing task registration over relying on this path).'
    & $PsExe -NonInteractive -ExecutionPolicy Bypass -File $SwapScript
}

Write-Host ''
Write-Host 'Upgrade staged.' -ForegroundColor Green
Write-Host "  Watch $LogFile for:"
Write-Host '    [DisplayAgentSwap] swap start -> swap_complete -> agent_restarted ok'
