# Tlamatini Auto-Updater
# =====================================================================
# Runs OUTSIDE the install directory (a copy is placed under
# %LOCALAPPDATA%\Tlamatini\updater by self_update.py) so it can safely
# replace the whole install — including python\, jre\, git\ and
# Tlamatini.exe — without locking the files it is overwriting.
#
# Sequence:
#   1. Validate install/staging boundaries and every staged file's receipt.
#   2. Close the running Tlamatini (and its child processes).
#   3. Verify a WAL-aware database backup, then rename agents -> agents_backup.
#   4. Delete the old install EXCEPT the preserved user-data set.
#   5. Move the new build in EXCEPT the preserved set (so the user's
#      config.json / DB / etc. are never overwritten).
#   6. Relaunch the new Tlamatini.
#
# It is launched by agent/self_update.py — do not run it by hand.
# =====================================================================

param(
    [Parameter(Mandatory = $true)][string]$InstallDir,
    [Parameter(Mandatory = $true)][string]$StagingDir,
    [Parameter(Mandatory = $true)][int]$ParentPid,
    [string]$RelaunchExe = "",
    [string]$LogPath = ""
)

$ErrorActionPreference = "Stop"
try { $Host.UI.RawUI.WindowTitle = "Tlamatini Updater" } catch {}

# Top-level names in the install root that must NEVER be deleted or
# overwritten. 'agents' is intentionally NOT here: it is renamed to
# 'agents_backup' (step 3) and then replaced by the new version (step 5).
# Keep the installed removal client until the incoming, verified client can replace it.

#
# The comment sits ABOVE the array on purpose. agent/test_preserved_user_state.py
# parses this block with a quoted-string regex, so an apostrophe INSIDE it - the
# old comment read "the user's uninstaller" - opened a phantom string that ate
# the 'Uninstaller.exe' entry, and the guard reported it missing from a file that
# was perfectly correct. Keep apostrophes out of this block.
$Preserve = @(
    'config.json', 'external_mcps.json', 'contacts.json', 'data.keys', 'DB', 'application', 'applications', 'content_generated',
    'Temp', 'context_files', 'doc_generated', 'documentation', 'Templates',
    'Uninstaller.exe'
)

function Write-Log {
    param([string]$Message, [string]$Color = "Gray")
    $line = "[{0}] {1}" -f (Get-Date -Format "HH:mm:ss"), $Message
    try { Write-Host $line -ForegroundColor $Color } catch { Write-Host $line }
    if ($LogPath) { try { Add-Content -Path $LogPath -Value $line -Encoding UTF8 } catch {} }
}

function Invoke-WithRetry {
    # Windows can hold file handles for a moment after a process dies, so
    # delete/rename/move operations are retried instead of failing outright.
    param([scriptblock]$Action, [int]$Attempts = 15, [int]$DelayMs = 600)
    for ($i = 1; $i -le $Attempts; $i++) {
        try { & $Action; return }
        catch {
            if ($i -eq $Attempts) { throw }
            Start-Sleep -Milliseconds $DelayMs
        }
    }
}

function Test-Preserved {
    param([string]$Name)
    foreach ($p in $Preserve) { if ($Name -ieq $p) { return $true } }
    return $false
}

function Install-UpdateUninstaller {
    param([string]$InstallDir, [string]$StagingDir, [string]$UpdateRoot)
    $incomingRemoval = Join-Path $StagingDir "Uninstaller.exe"
    if (Test-Path -LiteralPath $incomingRemoval -PathType Leaf) {
        $installedRemoval = Join-Path $InstallDir "Uninstaller.exe"
        Assert-NoReparseAncestors $installedRemoval
        if (Test-Path -LiteralPath $installedRemoval) {
            $previousRemoval = Join-Path $UpdateRoot ("previous-uninstaller-" + [guid]::NewGuid() + ".exe")
            Invoke-WithRetry { [IO.File]::Replace($incomingRemoval, $installedRemoval, $previousRemoval) }
        } else {
            Invoke-WithRetry { [IO.File]::Move($incomingRemoval, $installedRemoval) }
        }
    }
}

function Assert-NoReparseAncestors {
    param([string]$Path)
    $cursor = [IO.Path]::GetFullPath($Path)
    while ($cursor) {
        if (Test-Path -LiteralPath $cursor) {
            $item = Get-Item -LiteralPath $cursor -Force
            if ($item.Attributes -band [IO.FileAttributes]::ReparsePoint) {
                throw "Updater path traverses a reparse point: '$cursor'. Aborting."
            }
        }
        $cursor = [IO.Path]::GetDirectoryName($cursor)
    }
}

Write-Host ""
Write-Host "  ==================================================" -ForegroundColor Cyan
Write-Host "         TLAMATINI  --  APPLYING UPDATE" -ForegroundColor Cyan
Write-Host "  ==================================================" -ForegroundColor Cyan
Write-Host ""
Write-Log "Install dir : $InstallDir"
Write-Log "Staging dir : $StagingDir"
Write-Log "Parent PID  : $ParentPid"
Write-Host ""

try {
    # 1) Validate the staged new version BEFORE touching anything.
    $InstallDir = [IO.Path]::GetFullPath($InstallDir).TrimEnd('\', '/')
    $StagingDir = [IO.Path]::GetFullPath($StagingDir).TrimEnd('\', '/')
    $updateRoot = Join-Path $InstallDir "Temp\_update"
    $expectedStaging = Join-Path $updateRoot "staging"
    if ($StagingDir -ine $expectedStaging -or
        $InstallDir -ieq [IO.Path]::GetPathRoot($InstallDir).TrimEnd('\', '/') -or
        $InstallDir -ieq $env:USERPROFILE -or $InstallDir -ieq $env:WINDIR) {
        throw "Unsafe install/staging boundary. Expected staging strictly at InstallDir\Temp\_update\staging."
    }
    Assert-NoReparseAncestors $InstallDir
    Assert-NoReparseAncestors $StagingDir
    if (-not (Test-Path -LiteralPath (Join-Path $InstallDir "Tlamatini.exe") -PathType Leaf)) {
        throw "Install directory does not contain the existing Tlamatini.exe. Aborting."
    }
    if ($ParentPid -le 0 -or $ParentPid -eq $PID) { throw "Invalid application PID." }
    $applicationProcess = Get-CimInstance Win32_Process -Filter "ProcessId=$ParentPid" -ErrorAction Stop
    $installedExe = Join-Path $InstallDir "Tlamatini.exe"
    if ($applicationProcess -and $applicationProcess.ExecutablePath -ine $installedExe) {
        throw "The supplied PID does not belong to this installation. Nothing was replaced."
    }
    $newExe = Join-Path $StagingDir "Tlamatini.exe"
    if (-not (Test-Path -LiteralPath $newExe)) {
        throw "Staged update is invalid -- Tlamatini.exe not found at '$newExe'. Aborting; nothing was changed."
    }
    if (-not (Test-Path -LiteralPath $InstallDir)) {
        throw "Install directory '$InstallDir' does not exist. Aborting."
    }
    # Database preservation must be available before closing the running app.
    # Prefer the staged helper/runtime so the swap uses the new release's code.
    $backupPython = Join-Path $StagingDir "python\python.exe"
    $backupHelper = Join-Path $StagingDir "sqlite_copy.py"
    $integrityHelper = Join-Path $StagingDir "build_runtime_assets.py"
    foreach ($required in @($backupPython, $backupHelper, $integrityHelper)) {
        if (-not (Test-Path -LiteralPath $required -PathType Leaf)) {
            throw "Required staged updater asset missing: '$required'. Nothing was replaced."
        }
    }
    Write-Log "Verifying every staged file before closing Tlamatini..."
    # -B prevents the verifier from adding .pyc files to the tree it is checking;
    # -S avoids unrelated site-package startup hooks. Both helpers use stdlib only.
    & $backupPython -I -B -S $integrityHelper --verify-tree $StagingDir
    if ($LASTEXITCODE -ne 0) { throw "Staged release integrity check failed. Nothing was replaced." }
    $userDb = Join-Path $InstallDir "_internal\db.sqlite3"
    if ((Test-Path -LiteralPath $userDb -PathType Leaf) -and
        (-not (Test-Path -LiteralPath $backupPython -PathType Leaf) -or
         -not (Test-Path -LiteralPath $backupHelper -PathType Leaf))) {
        throw "Database backup helper or carried Python is missing. Aborting before closing Tlamatini."
    }
    Write-Log "Staged build and database-backup assets validated." "Green"

    # 2) Close the running Tlamatini and its whole process tree -- but NEVER
    #    this updater process. The updater is launched by self_update.py as a
    #    child of $ParentPid (the Tlamatini.exe). A blind
    #    `taskkill /T /PID $ParentPid` walks the recorded parent/child tree --
    #    which CREATE_BREAKAWAY_FROM_JOB does NOT change -- so it would also
    #    kill US, and the update would abort right here after "Closing
    #    Tlamatini..." (the exact symptom of the 1.19.5 -> 1.20.0 failure).
    #    So we enumerate the descendant tree ourselves and skip our own $PID.
    Write-Log "Closing Tlamatini (PID $ParentPid) and its child processes..."
    Start-Sleep -Seconds 2
    try {
        $procs   = Get-CimInstance Win32_Process -ErrorAction Stop
        $parentNow = $procs | Where-Object { [int]$_.ProcessId -eq $ParentPid }
        if ($parentNow -and (!$applicationProcess -or
            $parentNow.CreationDate -ne $applicationProcess.CreationDate -or
            $parentNow.ExecutablePath -ine $installedExe)) {
            throw "The application PID was reused. Update stopped before terminating any process."
        }
        $toKill  = New-Object System.Collections.Generic.List[int]
        $queue   = New-Object System.Collections.Generic.Queue[int]
        $seen    = New-Object System.Collections.Generic.HashSet[int]
        [void]$queue.Enqueue([int]$ParentPid)
        [void]$seen.Add([int]$ParentPid)
        while ($queue.Count -gt 0) {
            $cur = $queue.Dequeue()
            foreach ($p in $procs) {
                if ([int]$p.ParentProcessId -eq $cur) {
                    $cid = [int]$p.ProcessId
                    if ($cid -eq $PID) { continue }          # never our own updater
                    if ($seen.Add($cid)) {
                        [void]$queue.Enqueue($cid)
                        $toKill.Add($cid)                    # children, collected first
                    }
                }
            }
        }
        foreach ($k in $toKill) {
            try { & taskkill.exe /PID $k /F 2>&1 | Out-Null } catch {}
        }
        try { & taskkill.exe /PID $ParentPid /F 2>&1 | Out-Null } catch {}
        Write-Log ("Killed parent + {0} descendant process(es), kept updater (PID {1})." -f $toKill.Count, $PID)
    }
    catch {
        # Never fall back to /T: it can kill this updater, leaving a half-run swap.
        throw "Could not safely enumerate the application process tree; application files were not replaced."
    }
    # Let Windows release the file handles (python\, git\, jre\, exe).
    Start-Sleep -Seconds 3
    Write-Log "Tlamatini closed." "Green"

    # 3a) Preserve the user's DATABASE across the update. The live db.sqlite3
    #     lives INSIDE _internal\ (PyInstaller _MEIPASS), which step 4 deletes,
    #     so the top-level $Preserve set cannot protect it. Use SQLite's backup
    #     API to include committed WAL pages in one verified file, staged into
    #     the preserved DB\ToLoad folder and flag a post-update migrate: on the
    #     next launch manage.py's _apply_pending_db_swap restores it over the
    #     freshly shipped DB, then _run_post_update_migrate_if_flagged applies
    #     the new migrations (new agents / tools / demo prompts) to the user's
    #     data -- so chat history + custom toggles are KEPT, not wiped.
    if (Test-Path -LiteralPath $userDb) {
        $toLoadDir = Join-Path $InstallDir "DB\ToLoad"
        if (-not (Test-Path -LiteralPath $toLoadDir)) {
            Invoke-WithRetry { New-Item -ItemType Directory -Path $toLoadDir -Force | Out-Null }
        }
        $backupTarget = Join-Path $toLoadDir "db.sqlite3"
        & $backupPython -I -B -S $backupHelper $userDb $backupTarget
        if ($LASTEXITCODE -ne 0) {
            throw "Database backup failed. Application files have not been replaced; the original database is intact."
        }
        Set-Content -LiteralPath (Join-Path $InstallDir "DB\post_update_migrate.flag") -Value (Get-Date -Format o) -Encoding UTF8
        Write-Log "Preserved your database including WAL commits -> DB\ToLoad and flagged a post-update migrate." "Green"
    }
    else {
        Write-Log "No existing database at _internal\db.sqlite3 (fresh install?) -- skipping DB preserve." "Yellow"
    }

    # 3b) Back up the old agents directory: agents -> agents_backup.
    $agentsDir = Join-Path $InstallDir "agents"
    $agentsBak = Join-Path $InstallDir "agents_backup"
    if (Test-Path -LiteralPath $agentsDir) {
        if (Test-Path -LiteralPath $agentsBak) {
            Write-Log "Removing previous agents_backup..."
            Invoke-WithRetry { Remove-Item -LiteralPath $agentsBak -Recurse -Force }
        }
        Write-Log "Renaming agents -> agents_backup ..."
        Invoke-WithRetry { Rename-Item -LiteralPath $agentsDir -NewName "agents_backup" }
        Write-Log "Old agents backed up to 'agents_backup'." "Green"
    }
    else {
        Write-Log "No existing 'agents' directory to back up." "Yellow"
    }

    # 3c) Preserve the Blue-hat toolkit's EVIDENCE across the update.
    #     security\ is application code: the new release MUST replace the
    #     defender/whitelist scripts. But security\security_logs\ is the
    #     operator's own evidence -- alerts.log, monitor.log and the visible
    #     asset-test proof -- and it lives INSIDE security\, so like the
    #     database it cannot be protected by the top-level $Preserve set.
    #     A blind delete-and-replace destroys incident history at exactly the
    #     moment it matters. Stash it under the preserved Temp\ (same volume,
    #     so this is a rename, not a copy) and restore it in step 5b.
    #     Fail CLOSED before deletion if evidence cannot be moved safely.
    $securityLogs = Join-Path $InstallDir "security\security_logs"
    $logsCarryover = Join-Path $InstallDir ("Temp\_security_logs_carryover_" + [Guid]::NewGuid().ToString("N"))
    try {
        if (Test-Path -LiteralPath $securityLogs) {
            $carryParent = Split-Path -Parent $logsCarryover
            if (-not (Test-Path -LiteralPath $carryParent)) {
                New-Item -ItemType Directory -Path $carryParent -Force | Out-Null
            }
            Invoke-WithRetry { Move-Item -LiteralPath $securityLogs -Destination $logsCarryover -Force }
            Write-Log "Preserved your security evidence -> $logsCarryover." "Green"
        }
        else {
            Write-Log "No security\security_logs to preserve (toolkit never run?)." "Yellow"
        }
    }
    catch {
        throw "Could not preserve security evidence; application deletion aborted: $($_.Exception.Message)"
    }

    # Fix literal NTFS cache aliases while the verified staged Python is still in place.
    $recoveryHelper = Join-Path $StagingDir "setup_recovery.py"
    & $backupPython -I -B -S $recoveryHelper --clean-update-browser-caches $InstallDir
    if ($LASTEXITCODE -ne 0) { throw "Could not clear a legacy browser cache. Update stopped." }

    # Older clients stage only the inner ZIP. Replace its receipted uninstaller
    # atomically; when a compatible older payload lacks it, keep the installed copy.
    Install-UpdateUninstaller -InstallDir $InstallDir -StagingDir $StagingDir -UpdateRoot $updateRoot

    # 4) Delete the old install -- everything except the preserved set and
    #    the agents_backup we just created.
    Write-Log "Removing old application files (keeping your data)..."
    Get-ChildItem -LiteralPath $InstallDir -Force | ForEach-Object {
        $name = $_.Name
        if ($name -ieq "agents_backup") { return }
        if (Test-Preserved $name) { Write-Log "  keep    $name"; return }
        $full = $_.FullName
        Write-Log "  remove  $name"
        Invoke-WithRetry { Remove-Item -LiteralPath $full -Recurse -Force }
    }

    # 5) Move the new version in -- everything except the preserved set, so
    #    the user's config.json / DB / runtime dirs are never overwritten.
    Write-Log "Installing the new version..."
    Get-ChildItem -LiteralPath $StagingDir -Force | ForEach-Object {
        $name = $_.Name
        $src = $_.FullName
        $dest = Join-Path $InstallDir $name
        if ((Test-Preserved $name) -and (Test-Path -LiteralPath $dest)) {
            Write-Log "  skip    $name (kept yours)"; return
        }
        if (Test-Path -LiteralPath $dest) {
            Invoke-WithRetry { Remove-Item -LiteralPath $dest -Recurse -Force }
        }
        Write-Log "  install $name"
        Invoke-WithRetry { Move-Item -LiteralPath $src -Destination $dest -Force }
    }

    # 5b) Put the Blue-hat evidence back inside the NEW security\ tree.
    #     The new release ships security\ without security_logs\ (build.py
    #     ignores it on purpose), so this simply moves the operator's history
    #     back where the defender expects to append to it. FAIL-OPEN: on any
    #     error the unique stash is LEFT in Temp\ rather than
    #     deleted, so the evidence still exists and can be recovered by hand.
    try {
        if (Test-Path -LiteralPath $logsCarryover) {
            $securityDir = Join-Path $InstallDir "security"
            if (-not (Test-Path -LiteralPath $securityDir)) {
                New-Item -ItemType Directory -Path $securityDir -Force | Out-Null
            }
            if (Test-Path -LiteralPath $securityLogs) {
                throw "New evidence already exists at '$securityLogs'; retaining both copies."
            }
            Invoke-WithRetry { Move-Item -LiteralPath $logsCarryover -Destination $securityLogs -Force }
            Write-Log "Restored your security evidence -> security\security_logs." "Green"
        }
    }
    catch {
        Write-Log "WARN: security evidence kept at '$logsCarryover' -- move it back by hand: $($_.Exception.Message)" "Yellow"
    }

    # Repair owned flow registrations after files have been replaced. Explicitly
    # unregistered types and Windows default-app choices remain untouched.
    $flowRegistration = Join-Path $InstallDir 'flow_file_associations.ps1'
    if (Test-Path -LiteralPath $flowRegistration -PathType Leaf) {
        try {
            & $flowRegistration -Action Repair -InstallDir $InstallDir
            Write-Log 'Flow-file registration repair completed.'
        } catch { Write-Log "Flow-file registration repair needs attention: $_" 'Yellow' }
    }

    # 6) Clean up the staging area (best effort).
    try {
        # Exact validated directory from step 1, never a caller-derived parent.
        if ($updateRoot -and (Test-Path -LiteralPath $updateRoot)) {
            Remove-Item -LiteralPath $updateRoot -Recurse -Force -ErrorAction SilentlyContinue
        }
    }
    catch {}

    Write-Host ""
    Write-Log "UPDATE COMPLETE." "Green"

    # 7) Relaunch Tlamatini.
    if (-not $RelaunchExe) { $RelaunchExe = Join-Path $InstallDir "Tlamatini.exe" }
    if (Test-Path -LiteralPath $RelaunchExe) {
        Write-Log "Starting the new Tlamatini..." "Cyan"
        Start-Process -FilePath $RelaunchExe -WorkingDirectory $InstallDir -WindowStyle Hidden
    }
    else {
        Write-Log "Tlamatini.exe not found after update -- please start it manually." "Yellow"
    }
    Start-Sleep -Seconds 4
    exit 0
}
catch {
    Write-Host ""
    Write-Log "UPDATE FAILED: $($_.Exception.Message)" "Red"
    Write-Log "Your previous agents are preserved in 'agents_backup'." "Yellow"
    if ($LogPath) { Write-Log "Full log: $LogPath" "Yellow" }
    Write-Host ""
    try { Read-Host "Press Enter to close" } catch {}
    exit 1
}