# update.ps1 - Update PRE (Personal Reasoning Engine) on Windows
#
# Detects whether PRE was installed via git clone or zip download,
# compares local version against the latest release on GitHub,
# and updates accordingly. User data (~/.pre/) is never touched.
#
# Usage:
#   Right-click update.ps1 -> "Run with PowerShell"
#   powershell -ExecutionPolicy Bypass -File update.ps1
#   powershell -File update.ps1 -Yes    # Non-interactive (accept all defaults)

param(
    [switch]$Yes
)

$ErrorActionPreference = "Stop"

# Trap unexpected termination so the window stays open
trap {
    Write-Host ""
    Write-Host "  ERROR: $_" -ForegroundColor Red
    Write-Host "  At: $($_.InvocationInfo.ScriptName):$($_.InvocationInfo.ScriptLineNumber)" -ForegroundColor Red
    Write-Host ""
    Write-Host "  Press Enter to exit..." -ForegroundColor Yellow
    $null = Read-Host
    exit 1
}

# -- Helpers ----------------------------------------------------------------

function Step($msg) { Write-Host "`n=== $msg ===`n" -ForegroundColor Cyan }
function Ok($msg) { Write-Host "  $msg" -ForegroundColor Green }
function Warn($msg) { Write-Host "  $msg" -ForegroundColor Yellow }
function Fail($msg) {
    Write-Host "`n  $msg" -ForegroundColor Red
    Write-Host "`n  Press Enter to exit..." -ForegroundColor Yellow
    $null = Read-Host
    exit 1
}

function Ask-YN {
    param([string]$Prompt, [string]$Default = "Y")
    if ($Yes) {
        return ($Default -eq "Y")
    }
    $reply = Read-Host "$Prompt"
    if ($Default -eq "Y") {
        return ($reply -notmatch '^[Nn]')
    }
    return ($reply -match '^[Yy]')
}

$REPO_URL = "https://github.com/sunkencity999/pre"
$RAW_URL = "https://raw.githubusercontent.com/sunkencity999/pre/main"
$REPO_DIR = Split-Path -Parent $MyInvocation.MyCommand.Path
$WEB_DIR = Join-Path $REPO_DIR "web"
$ENGINE_DIR = Join-Path $REPO_DIR "engine"
$PRE_PORT = if ($env:PRE_PORT) { $env:PRE_PORT } else { "7749" }

# -- Detect install type ----------------------------------------------------

Step "Checking installation"

$IS_GIT = $false
if (Test-Path (Join-Path $REPO_DIR ".git")) {
    $IS_GIT = $true
    Ok "Install type: git repository"
    # Check for local modifications
    $status = & git -C $REPO_DIR status --porcelain 2>$null
    if ($status) {
        Warn "Local modifications detected."
        & git -C $REPO_DIR status --short
        Write-Host ""
    }
} else {
    Ok "Install type: zip download (no git repository)"
}

# -- Get local version ------------------------------------------------------

$LOCAL_VERSION = "unknown"
$VERSION_FILE = Join-Path $REPO_DIR "VERSION"
if (Test-Path $VERSION_FILE) {
    $LOCAL_VERSION = (Get-Content $VERSION_FILE -Raw).Trim()
}
Ok "Local version: $LOCAL_VERSION"

# -- Get remote version -----------------------------------------------------

Step "Checking for updates"

$REMOTE_VERSION = ""
try {
    $response = Invoke-WebRequest -Uri "$RAW_URL/VERSION" -UseBasicParsing -TimeoutSec 10
    $REMOTE_VERSION = $response.Content.Trim()
} catch {
    # Silently continue
}

if (-not $REMOTE_VERSION) {
    Warn "Could not fetch remote version. Check your internet connection."
    if (-not (Ask-YN "  Continue anyway? [y/N]" "N")) {
        exit 0
    }
    $REMOTE_VERSION = "unknown"
}

Ok "Remote version: $REMOTE_VERSION"

# Compare versions
if (($LOCAL_VERSION -eq $REMOTE_VERSION) -and ($REMOTE_VERSION -ne "unknown")) {
    Ok "You're already on the latest version ($LOCAL_VERSION)."
    if (-not (Ask-YN "  Force update anyway? [y/N]" "N")) {
        Write-Host ""
        Write-Host "  No update needed."
        Write-Host ""
        Write-Host "  Press Enter to exit..." -ForegroundColor Yellow
        $null = Read-Host
        exit 0
    }
} else {
    if ($REMOTE_VERSION -ne "unknown") {
        Ok "Update available: $LOCAL_VERSION -> $REMOTE_VERSION"
    }
}

# -- Check if server is running ---------------------------------------------

$SERVER_WAS_RUNNING = $false
try {
    $null = Invoke-WebRequest -Uri "http://localhost:${PRE_PORT}/api/status" -UseBasicParsing -TimeoutSec 3
    $SERVER_WAS_RUNNING = $true
} catch {
    # Server not running
}

if ($SERVER_WAS_RUNNING) {
    Warn "PRE server is currently running."
    if (Ask-YN "  Stop it for the update? [Y/n]" "Y") {
        Write-Host "  Stopping server..."
        # Try graceful shutdown via the API
        try {
            $null = Invoke-WebRequest -Uri "http://localhost:${PRE_PORT}/api/shutdown" -UseBasicParsing -TimeoutSec 3
        } catch {}
        Start-Sleep -Seconds 1
        # Kill by port if still running
        try {
            $connections = Get-NetTCPConnection -LocalPort $PRE_PORT -State Listen -ErrorAction SilentlyContinue
            if ($connections) {
                $connections | ForEach-Object {
                    Stop-Process -Id $_.OwningProcess -Force -ErrorAction SilentlyContinue
                }
            }
        } catch {}
        Start-Sleep -Seconds 1
        Ok "Server stopped."
    } else {
        Warn "Updating while server is running - restart manually after."
    }
}

# -- Perform update ---------------------------------------------------------

Step "Updating PRE"

if ($IS_GIT) {
    # ---- Git-based update ----
    Write-Host "  Pulling latest changes from origin..."

    # Stash local changes if any
    $STASHED = $false
    $status = & git -C $REPO_DIR status --porcelain 2>$null
    if ($status) {
        Warn "Stashing local changes..."
        & git -C $REPO_DIR stash push -m "pre-update-$(Get-Date -Format 'yyyyMMdd-HHmmss')" 2>$null
        $STASHED = $true
    }

    # Pull
    $pullOutput = & git -C $REPO_DIR pull origin main 2>&1
    if ($LASTEXITCODE -ne 0) {
        Write-Host "  $pullOutput"
        if ($STASHED) {
            Warn "Restoring stashed changes..."
            & git -C $REPO_DIR stash pop 2>$null
        }
        Fail "Git pull failed. Resolve conflicts manually."
    }
    Write-Host "  $pullOutput"

    # Restore stashed changes
    if ($STASHED) {
        Write-Host ""
        Warn "Restoring stashed changes..."
        $popResult = & git -C $REPO_DIR stash pop 2>&1
        if ($LASTEXITCODE -eq 0) {
            Ok "Local changes restored."
        } else {
            Warn "Could not auto-restore local changes. Run 'git stash pop' manually."
        }
    }

    Ok "Git update complete."

} else {
    # ---- Zip-based update ----
    Write-Host "  Downloading latest version from GitHub..."

    $TEMP_DIR = Join-Path $env:TEMP "pre-update-$(Get-Date -Format 'yyyyMMdd-HHmmss')"
    New-Item -ItemType Directory -Path $TEMP_DIR -Force | Out-Null
    $ZIP_FILE = Join-Path $TEMP_DIR "pre-latest.zip"
    $EXTRACT_DIR = Join-Path $TEMP_DIR "pre-main"

    # Download
    try {
        Invoke-WebRequest -Uri "$REPO_URL/archive/refs/heads/main.zip" -OutFile $ZIP_FILE -UseBasicParsing
    } catch {
        Remove-Item -Recurse -Force $TEMP_DIR -ErrorAction SilentlyContinue
        Fail "Download failed. Check your internet connection."
    }
    Ok "Downloaded."

    # Extract
    Write-Host "  Extracting..."
    try {
        Expand-Archive -Path $ZIP_FILE -DestinationPath $TEMP_DIR -Force
    } catch {
        Remove-Item -Recurse -Force $TEMP_DIR -ErrorAction SilentlyContinue
        Fail "Extraction failed."
    }

    # The zip extracts to a directory like "pre-main/"
    if (-not (Test-Path $EXTRACT_DIR)) {
        # Try to find the extracted directory
        $found = Get-ChildItem -Path $TEMP_DIR -Directory -Filter "pre-*" | Select-Object -First 1
        if ($found) {
            $EXTRACT_DIR = $found.FullName
        }
    }

    if (-not (Test-Path $EXTRACT_DIR)) {
        Remove-Item -Recurse -Force $TEMP_DIR -ErrorAction SilentlyContinue
        Fail "Could not find extracted files."
    }

    Ok "Extracted to: $EXTRACT_DIR"

    # Copy files using robocopy (built into Windows, reliable recursive copy)
    Write-Host "  Updating files..."

    # Update engine/ (source code, Modelfile, scripts -- exclude compiled binaries)
    $srcEngine = Join-Path $EXTRACT_DIR "engine"
    if (Test-Path $srcEngine) {
        # /E = recursive, /IS = include same files, /IT = include tweaked files
        # /XF = exclude files, /NFL /NDL /NJH /NJS /NP = quiet output
        $roboArgs = @(
            "`"$srcEngine`"", "`"$ENGINE_DIR`"",
            "/E", "/IS", "/IT",
            "/XF", "*.o", "*.exe", "pre", "telegram",
            "/NFL", "/NDL", "/NJH", "/NJS", "/NP"
        )
        $roboResult = Start-Process -FilePath "robocopy" -ArgumentList $roboArgs -NoNewWindow -Wait -PassThru
        # robocopy exit codes: 0=no change, 1=files copied, 2=extra files, 3=both -- all OK
        # Exit codes >= 8 indicate errors
        if ($roboResult.ExitCode -ge 8) {
            Warn "Engine copy had issues (robocopy exit code $($roboResult.ExitCode))"
        } else {
            Ok "Engine files updated"
        }
    }

    # Update web/ (server, tools, frontend -- exclude node_modules)
    $srcWeb = Join-Path $EXTRACT_DIR "web"
    if (Test-Path $srcWeb) {
        $roboArgs = @(
            "`"$srcWeb`"", "`"$WEB_DIR`"",
            "/E", "/IS", "/IT",
            "/XD", "node_modules",
            "/NFL", "/NDL", "/NJH", "/NJS", "/NP"
        )
        $roboResult = Start-Process -FilePath "robocopy" -ArgumentList $roboArgs -NoNewWindow -Wait -PassThru
        if ($roboResult.ExitCode -ge 8) {
            Warn "Web copy had issues (robocopy exit code $($roboResult.ExitCode))"
        } else {
            Ok "Web files updated"
        }
    } else {
        Warn "web/ directory not found in downloaded archive"
        # List what was extracted for debugging
        Write-Host "    Contents of extracted archive:"
        Get-ChildItem -Path $EXTRACT_DIR -Name | ForEach-Object { Write-Host "      $_" }
    }

    # Update root files (scripts, README, VERSION, etc.)
    $rootFiles = @(
        "VERSION", "README.md", "CLAUDE.md", "install.sh", "install-linux.sh",
        "install.ps1", "install.cmd",
        "Launch PRE.command", "Launch PRE.cmd", "PRE Tray.cmd",
        "Install PRE.command", "Update PRE.command", "Update PRE.cmd",
        "update.sh", "update.ps1", "system.md", "benchmark.sh"
    )
    $rootCount = 0
    foreach ($f in $rootFiles) {
        $srcFile = Join-Path $EXTRACT_DIR $f
        if (Test-Path $srcFile) {
            Copy-Item -Path $srcFile -Destination (Join-Path $REPO_DIR $f) -Force
            $rootCount++
        }
    }
    Ok "$rootCount root files updated"

    # Clean up
    Remove-Item -Recurse -Force $TEMP_DIR -ErrorAction SilentlyContinue
    Ok "Files updated."
}

# -- Update dependencies ----------------------------------------------------

Step "Updating dependencies"

$packageJson = Join-Path $WEB_DIR "package.json"
if (Test-Path $packageJson) {
    Write-Host "  Running npm install..."
    Push-Location $WEB_DIR
    try {
        & npm install --silent 2>&1 | Select-Object -Last 3
        Ok "Dependencies updated."
    } finally {
        Pop-Location
    }
} else {
    Warn "web\package.json not found - skipping npm install."
}

# -- Check if Modelfile changed ---------------------------------------------

if ($IS_GIT) {
    try {
        $changed = & git -C $REPO_DIR diff HEAD~1 --name-only 2>$null | Select-String "Modelfile"
        if ($changed) {
            Warn "Modelfile has changed. You may want to recreate the custom model:"
            Write-Host "    ollama create pre-gemma4 -f engine\Modelfile"
        }
    } catch {}
}

# -- Post-update migration (1.0 -> 2.0+) ------------------------------------
# Ensures features added after the user's initial install are properly set up.
# Each step is idempotent -- safe to run on every update.

Step "Post-update migration"

$PRE_DATA = Join-Path $env:USERPROFILE ".pre"
$migrationCount = 0

# -- Create missing data directories --
$requiredDirs = @(
    (Join-Path $PRE_DATA "agent-checkpoints"),
    (Join-Path $PRE_DATA "models"),
    (Join-Path $PRE_DATA "training"),
    (Join-Path (Join-Path $PRE_DATA "training") "exports"),
    (Join-Path $PRE_DATA "custom_tools"),
    (Join-Path $PRE_DATA "workflows"),
    (Join-Path $PRE_DATA "rag")
)
foreach ($dir in $requiredDirs) {
    if (-not (Test-Path $dir)) {
        New-Item -ItemType Directory -Path $dir -Force | Out-Null
        $migrationCount++
    }
}
if ($migrationCount -gt 0) {
    Ok "Created $migrationCount missing data directories"
    $migrationCount = 0
}

# -- Fast tier model (pre-gemma4-fast) --
$FAST_MODEL = "pre-gemma4-fast"
$fastModelfile = Join-Path $ENGINE_DIR "Modelfile.fast"
if ((Get-Command ollama -ErrorAction SilentlyContinue) -and (Test-Path $fastModelfile)) {
    $modelList = & ollama list 2>&1
    if ($modelList -notmatch $FAST_MODEL) {
        Write-Host ""
        Write-Host "  The multi-model router's fast tier uses $FAST_MODEL (Gemma 4 E2B, ~3 GB)."
        Write-Host "  This model provides ~140 tok/s responses for simple queries."
        if (Ask-YN "  Create the fast tier model? [Y/n]" "Y") {
            Write-Host "  Pulling base model (gemma4:e2b)..." -ForegroundColor DarkGray
            & ollama pull gemma4:e2b 2>&1 | Out-Null
            $modelList2 = & ollama list 2>&1
            if ($modelList2 -match "gemma4:e2b") {
                & ollama create $FAST_MODEL -f $fastModelfile 2>&1 | Out-Null
                if ($LASTEXITCODE -eq 0) {
                    Ok "Fast tier model created ($FAST_MODEL)"
                } else {
                    Warn "Failed to create $FAST_MODEL - fast tier will fall back to main model"
                }
            } else {
                Warn "Failed to pull gemma4:e2b - fast tier will be unavailable"
            }
        } else {
            Write-Host "  Skipped. Create later with:"
            Write-Host "    ollama pull gemma4:e2b"
            Write-Host "    ollama create $FAST_MODEL -f engine\Modelfile.fast"
        }
    } else {
        Ok "Fast tier model already exists ($FAST_MODEL)"
    }
}

# -- Ollama environment variables (NUM_PARALLEL, MAX_LOADED_MODELS) --
# 2.0 uses NUM_PARALLEL=2 for parallel inference (user + background agents)
# and MAX_LOADED_MODELS=2 for main model + embedding model simultaneously.
# On Windows, Ollama reads env vars from system/user environment.
$currentParallel = [System.Environment]::GetEnvironmentVariable("OLLAMA_NUM_PARALLEL", "User")
$currentLoaded = [System.Environment]::GetEnvironmentVariable("OLLAMA_MAX_LOADED_MODELS", "User")
$needsEnvUpdate = ($currentParallel -ne "2") -or ($currentLoaded -ne "2")

if ($needsEnvUpdate) {
    Write-Host ""
    Write-Host "  PRE 2.0 uses OLLAMA_NUM_PARALLEL=2 (parallel user + background inference)"
    Write-Host "  and OLLAMA_MAX_LOADED_MODELS=2 (main model + embedding model)."
    if (Ask-YN "  Update Ollama environment variables? [Y/n]" "Y") {
        [System.Environment]::SetEnvironmentVariable("OLLAMA_NUM_PARALLEL", "2", "User")
        [System.Environment]::SetEnvironmentVariable("OLLAMA_MAX_LOADED_MODELS", "2", "User")
        # Also set for current process
        $env:OLLAMA_NUM_PARALLEL = "2"
        $env:OLLAMA_MAX_LOADED_MODELS = "2"
        Ok "Ollama environment updated (NUM_PARALLEL=2, MAX_LOADED_MODELS=2)"
        Warn "Restart Ollama for changes to take effect"
    }
} else {
    Ok "Ollama environment settings are current"
}

# -- Update Windows Startup shortcut if present --
$startupDir = Join-Path $env:APPDATA "Microsoft\Windows\Start Menu\Programs\Startup"
$startupVbs = Join-Path $startupDir "PRE Server.vbs"
if (Test-Path $startupVbs) {
    $vbsContent = Get-Content $startupVbs -Raw -ErrorAction SilentlyContinue
    if ($vbsContent -and ($vbsContent -notmatch "OLLAMA_NUM_PARALLEL")) {
        # Old-style startup script without env vars — update it to set them
        $envPrefix = @"
Dim shell : Set shell = CreateObject("WScript.Shell")
shell.Environment("Process")("OLLAMA_NUM_PARALLEL") = "2"
shell.Environment("Process")("OLLAMA_MAX_LOADED_MODELS") = "2"
shell.Environment("Process")("OLLAMA_KEEP_ALIVE") = "24h"

"@
        $newContent = $envPrefix + $vbsContent
        Set-Content -Path $startupVbs -Value $newContent -NoNewline
        Ok "Updated Startup script with Ollama environment variables"
    }
}

# -- Recalculate context window if VRAM-based sizing was missed --
# Older installs sized context from system RAM instead of VRAM, resulting
# in a context window too large for the GPU's KV cache (causes empty responses).
$contextFile = Join-Path $PRE_DATA "context"
if (Test-Path $contextFile) {
    $currentCtx = 0
    try { $currentCtx = [int](Get-Content $contextFile -Raw).Trim() } catch {}
    # Detect NVIDIA GPU VRAM
    $updateVramGB = 0
    try {
        $nvOut = & nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits 2>$null
        if ($nvOut -and $nvOut.Trim() -match '^\d+$') {
            $updateVramGB = [math]::Floor([int]$nvOut.Trim() / 1024)
        }
    } catch {}
    if ($updateVramGB -gt 0 -and $currentCtx -gt 0) {
        # Determine model size (q8_0=28GB, q4_K_M=15GB)
        $modelList = & ollama list 2>&1
        $updateModelSize = 15  # default to q4_K_M
        if ($modelList -match 'q8_0') { $updateModelSize = 28 }
        $vramHeadroom = $updateVramGB - $updateModelSize
        # Calculate what context SHOULD be based on VRAM
        if ($vramHeadroom -ge 68) { $correctCtx = 131072 }
        elseif ($vramHeadroom -ge 36) { $correctCtx = 65536 }
        elseif ($vramHeadroom -ge 20) { $correctCtx = 32768 }
        elseif ($vramHeadroom -ge 8) { $correctCtx = 16384 }
        elseif ($vramHeadroom -ge 4) { $correctCtx = 8192 }
        else { $correctCtx = 4096 }
        if ($currentCtx -gt $correctCtx) {
            $oldHuman = "$([math]::Floor($currentCtx / 1024))K"
            $newHuman = "$([math]::Floor($correctCtx / 1024))K"
            Warn "Context window was $oldHuman (sized from system RAM) but GPU has ${updateVramGB}GB VRAM."
            Warn "Reducing to $newHuman to fit in VRAM (${vramHeadroom}GB headroom after model)."
            Set-Content -Path $contextFile -Value $correctCtx -NoNewline
            Ok "Context window updated: $oldHuman -> $newHuman"
        }
    }
}

Ok "Post-update migration complete"

# -- Restart server if it was running ----------------------------------------

if ($SERVER_WAS_RUNNING) {
    Step "Restarting server"
    if (Ask-YN "  Restart PRE server? [Y/n]" "Y") {
        Push-Location $WEB_DIR
        try {
            Start-Process -FilePath "node" -ArgumentList "server.js" -WindowStyle Hidden -WorkingDirectory $WEB_DIR
            Start-Sleep -Seconds 2
            try {
                $null = Invoke-WebRequest -Uri "http://localhost:${PRE_PORT}/api/status" -UseBasicParsing -TimeoutSec 3
                Ok "Server restarted on port $PRE_PORT."
            } catch {
                Warn "Server may still be starting. Check http://localhost:${PRE_PORT}"
            }
        } finally {
            Pop-Location
        }
    }
}

# -- Done -------------------------------------------------------------------

Step "Update Complete!"

$NEW_VERSION = "unknown"
if (Test-Path (Join-Path $REPO_DIR "VERSION")) {
    $NEW_VERSION = (Get-Content (Join-Path $REPO_DIR "VERSION") -Raw).Trim()
}

Write-Host ""
Ok "PRE updated to version $NEW_VERSION"
Write-Host ""
Write-Host "  Changes in this update:"
if ($IS_GIT) {
    Write-Host "    git log --oneline -10"
} else {
    Write-Host "    See: $REPO_URL/commits/main"
}
Write-Host ""
Write-Host "  If you encounter issues after updating:"
Write-Host "    1. Restart the server: cd web && node server.js"
Write-Host "    2. Recreate the model: ollama create pre-gemma4 -f engine\Modelfile"
Write-Host "    3. Report issues: $REPO_URL/issues"
Write-Host ""

Write-Host "  Press Enter to exit..." -ForegroundColor Yellow
$null = Read-Host
