# AI Stack Windows service updater
# Run in elevated PowerShell:
# powershell -NoProfile -ExecutionPolicy Bypass -File .\update.ps1

param(
    [switch]$SkipRestart
)

$ErrorActionPreference = "Stop"

$BASE_DIR = $PSScriptRoot
$WINSW_EXE = Join-Path $BASE_DIR "WinSW.exe"
$LITELLM_XML = Join-Path $BASE_DIR "litellm\litellm-service.xml"
$WEBUI_XML = Join-Path $BASE_DIR "open_webui\openwebui-service.xml"
$ENV_FILE = Join-Path $BASE_DIR ".env"
$LITELLM_VENV = Join-Path $BASE_DIR "litellm\venv"
$WEBUI_VENV = Join-Path $BASE_DIR "open_webui\venv"

$LOG_DIR = Join-Path $BASE_DIR "logs\install"
$UPDATE_LOG = Join-Path $LOG_DIR ("update-{0}.log" -f (Get-Date -Format "yyyyMMdd-HHmmss"))

function Fail {
    param([string]$Message)
    Write-Host "[ERROR] $Message" -ForegroundColor Red
    throw $Message
}

function Step {
    param([string]$Message)
    Write-Host "[STEP] $Message" -ForegroundColor Cyan
}

function Diag {
    param([string]$Message)
    Write-Host "[DIAG] $Message" -ForegroundColor DarkGray
}

function Ensure-LogDir {
    if (-not (Test-Path $LOG_DIR)) {
        New-Item -ItemType Directory -Path $LOG_DIR -Force | Out-Null
    }
}

function Get-DotEnvValue {
    param([string]$Name)

    if (-not (Test-Path $ENV_FILE)) {
        return $null
    }

    foreach ($line in Get-Content $ENV_FILE) {
        $trimmed = $line.Trim()
        if ([string]::IsNullOrWhiteSpace($trimmed)) { continue }
        if ($trimmed.StartsWith("#")) { continue }
        if ($trimmed -like "$Name=*") {
            return $trimmed.Substring($Name.Length + 1).Trim()
        }
    }

    return $null
}

function Ensure-GeminiApiKey {
    $apiKey = $env:GEMINI_API_KEY
    if ([string]::IsNullOrWhiteSpace($apiKey)) {
        $apiKey = Get-DotEnvValue -Name "GEMINI_API_KEY"
    }

    if ([string]::IsNullOrWhiteSpace($apiKey)) {
        Fail "GEMINI_API_KEY not found. Set a system environment variable or add GEMINI_API_KEY=... to .env"
    }

    [Environment]::SetEnvironmentVariable("GEMINI_API_KEY", $apiKey, "Machine")
    $env:GEMINI_API_KEY = $apiKey
    Diag "GEMINI_API_KEY synced to machine environment for Windows services"
}

function Check-Admin {
    $identity = [System.Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = New-Object System.Security.Principal.WindowsPrincipal($identity)
    if (-not $principal.IsInRole([System.Security.Principal.WindowsBuiltInRole]::Administrator)) {
        Fail "Please run PowerShell as Administrator."
    }
}

function Invoke-External {
    param(
        [Parameter(Mandatory = $true)]
        [string]$File,
        [string[]]$Arguments = @(),
        [Parameter(Mandatory = $true)]
        [string]$ActionLabel
    )

    $display = if ($Arguments.Count -gt 0) { " " + ($Arguments -join " ") } else { "" }
    Step "$ActionLabel -> $File$display"

    try {
        & $File @Arguments
        if ($LASTEXITCODE -ne $null -and $LASTEXITCODE -ne 0) {
            Fail "$ActionLabel failed, exit code: $LASTEXITCODE"
        }
    } catch {
        Write-Host "[ERROR] file: $File" -ForegroundColor Red
        Write-Host "[ERROR] args: $($Arguments -join ' ')" -ForegroundColor Red
        Write-Host "[ERROR] detail: $($_.Exception.Message)" -ForegroundColor Red
        throw
    }
}

function Check-Prerequisites {
    if (-not (Test-Path $WINSW_EXE)) {
        Fail "WinSW.exe not found in project root. Download WinSW-x64.exe and rename to WinSW.exe."
    }

    foreach ($path in @($LITELLM_XML, $WEBUI_XML)) {
        if (-not (Test-Path $path)) {
            Fail "Required file not found: $path"
        }
    }

    foreach ($venvPath in @($LITELLM_VENV, $WEBUI_VENV)) {
        $pythonExe = Join-Path $venvPath "Scripts\python.exe"
        if (-not (Test-Path $pythonExe)) {
            Fail "Virtual environment not found: $venvPath. Run install.ps1 first."
        }
    }
}

function Stop-ServicesIfPresent {
    foreach ($svc in @("ai-stack-openwebui", "ai-stack-litellm")) {
        $service = Get-Service -Name $svc -ErrorAction SilentlyContinue
        if ($service -and $service.Status -eq "Running") {
            Step "Stop service $svc"
            Stop-Service $svc -Force
            Start-Sleep -Seconds 2
        }
    }
}

function Upgrade-VenvPackages {
    param(
        [string]$PythonExe,
        [string[]]$Packages,
        [string]$Name
    )

    Invoke-External -File $PythonExe -Arguments @("-m", "pip", "install", "--upgrade", "pip", "setuptools", "wheel") -ActionLabel "Upgrade pip toolchain ($Name)"
    $upgradeArgs = @("-m", "pip", "install", "--upgrade") + $Packages
    Invoke-External -File $PythonExe -Arguments $upgradeArgs -ActionLabel "Upgrade packages ($Name)"
}

function Refresh-ServiceBinaries {
    $litellmWinSW = Join-Path $BASE_DIR "litellm\WinSW.exe"
    $webuiWinSW = Join-Path $BASE_DIR "open_webui\WinSW.exe"

    Step "Sync WinSW.exe into service folders"
    Copy-Item $WINSW_EXE $litellmWinSW -Force
    Copy-Item $WINSW_EXE $webuiWinSW -Force
}

function Start-ServicesIfNeeded {
    if ($SkipRestart) {
        Diag "SkipRestart is set, services remain stopped after package upgrade."
        return
    }

    Step "Start LiteLLM service"
    Start-Service ai-stack-litellm
    Start-Sleep -Seconds 3

    Step "Start Open WebUI service"
    Start-Service ai-stack-openwebui
}

function Update-All {
    Write-Host ""
    Write-Host "=== Update AI Stack Windows Services ===" -ForegroundColor Cyan
    Write-Host ""
    Write-Host "This update preserves config files, Open WebUI data, and existing venv directories." -ForegroundColor Yellow
    Write-Host "It only upgrades packages and refreshes service runtime files." -ForegroundColor Yellow
    Write-Host ""

    Check-Admin
    Check-Prerequisites
    Ensure-GeminiApiKey
    Stop-ServicesIfPresent
    Refresh-ServiceBinaries

    $litellmPython = Join-Path $LITELLM_VENV "Scripts\python.exe"
    $webuiPython = Join-Path $WEBUI_VENV "Scripts\python.exe"

    Upgrade-VenvPackages -PythonExe $litellmPython -Packages @("email-validator", "litellm[proxy]") -Name "LiteLLM"
    Upgrade-VenvPackages -PythonExe $webuiPython -Packages @("open-webui") -Name "Open WebUI"

    Start-ServicesIfNeeded

    Write-Host ""
    Write-Host "[OK] Update completed" -ForegroundColor Green
    Write-Host "Configuration and data were preserved." -ForegroundColor Green
    if ($SkipRestart) {
        Write-Host "Services were not restarted because -SkipRestart was specified." -ForegroundColor Yellow
    } else {
        Write-Host "Open WebUI: http://127.0.0.1:3000"
        Write-Host "LiteLLM API: http://127.0.0.1:4000"
    }
    Write-Host ""
}

Ensure-LogDir
Write-Host "Update log file: $UPDATE_LOG" -ForegroundColor DarkGray
Start-Transcript -Path $UPDATE_LOG -Append | Out-Null

try {
    Update-All
} catch {
    Write-Host ""
    Write-Host "Update failed. Send this log file:" -ForegroundColor Red
    Write-Host "  $UPDATE_LOG" -ForegroundColor Red
    Write-Host "Error: $($_.Exception.Message)" -ForegroundColor Red
    throw
} finally {
    Stop-Transcript | Out-Null
}