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

param(
    [switch]$Uninstall,
    [switch]$Start,
    [switch]$Stop,
    [switch]$Status,
    [string]$VenvDir,
    [switch]$SkipVenvBootstrap
)

$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"

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

# Venv paths for each service (separate venv for independence)
$LITELLM_VENV = Join-Path $BASE_DIR "litellm\venv"
$WEBUI_VENV = Join-Path $BASE_DIR "open_webui\venv"

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-InstallLogDir {
    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 Test-PEExecutable {
    param(
        [string]$Path,
        [string]$Name
    )

    if (-not (Test-Path $Path)) {
        Fail "$Name not found: $Path"
    }

    $file = Get-Item $Path
    $bytes = [System.IO.File]::ReadAllBytes($Path)
    $header = if ($bytes.Length -ge 2) { "{0:X2} {1:X2}" -f $bytes[0], $bytes[1] } else { "N/A" }

    Diag "$Name path: $Path"
    Diag "$Name size: $($file.Length) bytes"
    Diag "$Name header: $header"

    if ($bytes.Length -lt 2 -or $bytes[0] -ne 0x4D -or $bytes[1] -ne 0x5A) {
        Fail "$Name is not a valid Windows executable (missing MZ header)."
    }
}

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-WinSW {
    if (-not (Test-Path $WINSW_EXE)) {
        Fail "WinSW.exe not found in project root. Download WinSW-x64.exe and rename to WinSW.exe."
    }

    $file = Get-Item $WINSW_EXE
    if ($file.Length -lt 100KB) {
        Fail "WinSW.exe file size looks invalid ($($file.Length) bytes)."
    }

    Test-PEExecutable -Path $WINSW_EXE -Name "WinSW.exe"

    Step "Validate WinSW executable"
    Invoke-External -File $WINSW_EXE -Arguments @("--version") -ActionLabel "Check WinSW version"
}

function Resolve-VenvDir {
    param([string]$VenvPath)
    $pythonExe = Join-Path $VenvPath "Scripts\python.exe"
    return (Test-Path $pythonExe)
}

function New-VenvIfNeeded {
    param(
        [Parameter(Mandatory = $true)]
        [string]$VenvPath,
        [Parameter(Mandatory = $true)]
        [string]$ServiceName
    )

    $pythonExe = Join-Path $VenvPath "Scripts\python.exe"
    if (Test-Path $pythonExe) {
        return $VenvPath
    }

    if ($SkipVenvBootstrap) {
        Fail "No virtual environment found at $VenvPath and -SkipVenvBootstrap is set."
    }

    Step "Create virtual environment for $ServiceName : $VenvPath"

    $created = $false

    try {
        & py -3.11 -m venv $VenvPath
        if ((Test-Path $pythonExe) -and $LASTEXITCODE -eq 0) { $created = $true }
    } catch {}

    if (-not $created) {
        try {
            & py -m venv $VenvPath
            if ((Test-Path $pythonExe) -and $LASTEXITCODE -eq 0) { $created = $true }
        } catch {}
    }

    if (-not $created) {
        try {
            & python -m venv $VenvPath
            if ((Test-Path $pythonExe) -and $LASTEXITCODE -eq 0) { $created = $true }
        } catch {}
    }

    if (-not $created) {
        Fail "Failed to create virtual environment at $VenvPath. Ensure Python 3.11 is installed."
    }

    Write-Host "[OK] virtual environment created: $VenvPath" -ForegroundColor Green
    return $VenvPath
}

function Install-LiteLLMDependencies {
    param([string]$VenvPath)

    $venvPython = Join-Path $VenvPath "Scripts\python.exe"
    $litellmExe = Join-Path $VenvPath "Scripts\litellm.exe"

    if (Test-Path $litellmExe) {
        Diag "LiteLLM already installed in $VenvPath"
        return
    }

    if ($SkipVenvBootstrap) {
        Fail "Missing litellm in venv and -SkipVenvBootstrap is set."
    }

    Step "Install LiteLLM dependencies"
    Invoke-External -File $venvPython -Arguments @("-m", "pip", "install", "--upgrade", "pip", "setuptools", "wheel") -ActionLabel "Upgrade pip toolchain (LiteLLM)"

    # Install litellm[proxy] with email-validator for pydantic
    Invoke-External -File $venvPython -Arguments @("-m", "pip", "install", "email-validator") -ActionLabel "Install email-validator"
    Invoke-External -File $venvPython -Arguments @("-m", "pip", "install", "litellm[proxy]") -ActionLabel "Install litellm[proxy]"

    if (-not (Test-Path $litellmExe)) {
        Fail "LiteLLM installed but litellm.exe not found at $litellmExe"
    }
}

function Install-OpenWebUIDependencies {
    param([string]$VenvPath)

    $venvPython = Join-Path $VenvPath "Scripts\python.exe"
    $webuiExe = Join-Path $VenvPath "Scripts\open-webui.exe"

    if (Test-Path $webuiExe) {
        Diag "Open WebUI already installed in $VenvPath"
        return
    }

    if ($SkipVenvBootstrap) {
        Fail "Missing open-webui in venv and -SkipVenvBootstrap is set."
    }

    Step "Install Open WebUI dependencies"
    Invoke-External -File $venvPython -Arguments @("-m", "pip", "install", "--upgrade", "pip", "setuptools", "wheel") -ActionLabel "Upgrade pip toolchain (Open WebUI)"
    Invoke-External -File $venvPython -Arguments @("-m", "pip", "install", "open-webui") -ActionLabel "Install open-webui"

    if (-not (Test-Path $webuiExe)) {
        Fail "Open WebUI installed but open-webui.exe not found at $webuiExe"
    }
}

function Check-Venvs {
    Step "Check/create LiteLLM venv"
    $script:LiteLLMVenv = New-VenvIfNeeded -VenvPath $LITELLM_VENV -ServiceName "LiteLLM"
    Install-LiteLLMDependencies -VenvPath $script:LiteLLMVenv

    Step "Check/create Open WebUI venv"
    $script:WebUIVenv = New-VenvIfNeeded -VenvPath $WEBUI_VENV -ServiceName "Open WebUI"
    Install-OpenWebUIDependencies -VenvPath $script:WebUIVenv

    $litellmExe = Join-Path $script:LiteLLMVenv "Scripts\litellm.exe"
    $webuiExe = Join-Path $script:WebUIVenv "Scripts\open-webui.exe"

    if (-not (Test-Path $litellmExe)) { Fail "litellm.exe not found: $litellmExe" }
    if (-not (Test-Path $webuiExe)) { Fail "open-webui.exe not found: $webuiExe" }

    Diag "LiteLLM venv: $script:LiteLLMVenv"
    Diag "Open WebUI venv: $script:WebUIVenv"
}

function Ensure-DataDirs {
    $dirs = @(
        (Join-Path $BASE_DIR "logs\litellm"),
        (Join-Path $BASE_DIR "logs\openwebui"),
        (Join-Path $BASE_DIR ".data\open-webui")
    )

    foreach ($d in $dirs) {
        if (-not (Test-Path $d)) {
            New-Item -ItemType Directory -Path $d -Force | Out-Null
            Diag "Created directory: $d"
        }
    }
}

function Install-Services {
    Write-Host ""
    Write-Host "=== Install AI Stack Windows Services ===" -ForegroundColor Cyan
    Write-Host ""
    Write-Host "This install preserves existing config and data directories." -ForegroundColor Yellow
    Write-Host "It only creates missing venvs, installs packages, and registers services." -ForegroundColor Yellow
    Write-Host ""

    Check-Admin
    Check-WinSW
    Ensure-GeminiApiKey
    Check-Venvs
    Ensure-DataDirs

    $litellmWinSW = Join-Path $BASE_DIR "litellm\WinSW.exe"
    $webuiWinSW = Join-Path $BASE_DIR "open_webui\WinSW.exe"

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

    Test-PEExecutable -Path $litellmWinSW -Name "litellm\\WinSW.exe"
    Test-PEExecutable -Path $webuiWinSW -Name "open_webui\\WinSW.exe"

    Invoke-External -File $litellmWinSW -Arguments @("--version") -ActionLabel "Validate litellm WinSW"
    Invoke-External -File $webuiWinSW -Arguments @("--version") -ActionLabel "Validate open_webui WinSW"

    Invoke-External -File $litellmWinSW -Arguments @("install", $LITELLM_XML) -ActionLabel "Install LiteLLM service"
    Invoke-External -File $webuiWinSW -Arguments @("install", $WEBUI_XML) -ActionLabel "Install Open WebUI service"

    Step "Start services"
    Start-Service ai-stack-litellm
    Start-Sleep -Seconds 3
    Start-Service ai-stack-openwebui

    Write-Host ""
    Write-Host "[OK] Installation completed" -ForegroundColor Green
    Write-Host "Open WebUI: http://127.0.0.1:3000"
    Write-Host "LiteLLM API: http://127.0.0.1:4000"
    Write-Host "Configuration and persistent data were preserved." -ForegroundColor Green
    Write-Host ""
}

function Uninstall-Services {
    Write-Host ""
    Write-Host "=== Uninstall AI Stack Windows Services ===" -ForegroundColor Yellow
    Write-Host ""

    Check-Admin

    $litellmWinSW = Join-Path $BASE_DIR "litellm\WinSW.exe"
    $webuiWinSW = Join-Path $BASE_DIR "open_webui\WinSW.exe"

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

    if (Test-Path $webuiWinSW) {
        Invoke-External -File $webuiWinSW -Arguments @("uninstall", $WEBUI_XML) -ActionLabel "Uninstall Open WebUI service"
    }

    if (Test-Path $litellmWinSW) {
        Invoke-External -File $litellmWinSW -Arguments @("uninstall", $LITELLM_XML) -ActionLabel "Uninstall LiteLLM service"
    }

    Write-Host ""
    Write-Host "[OK] Uninstall completed" -ForegroundColor Green
    Write-Host ""
}

function Start-Services {
    Check-Admin
    Step "Start LiteLLM"
    Start-Service ai-stack-litellm
    Start-Sleep -Seconds 3
    Step "Start Open WebUI"
    Start-Service ai-stack-openwebui
}

function Stop-Services {
    Check-Admin
    Step "Stop Open WebUI"
    Stop-Service ai-stack-openwebui -Force -ErrorAction SilentlyContinue
    Step "Stop LiteLLM"
    Stop-Service ai-stack-litellm -Force -ErrorAction SilentlyContinue
}

function Show-Status {
    Write-Host ""
    Write-Host "=== AI Stack Service Status ===" -ForegroundColor Cyan
    Write-Host ""

    foreach ($svc in @("ai-stack-litellm", "ai-stack-openwebui")) {
        $s = Get-Service -Name $svc -ErrorAction SilentlyContinue
        if ($s) {
            $color = if ($s.Status -eq "Running") { "Green" } else { "Red" }
            Write-Host "  $($s.DisplayName): $($s.Status)" -ForegroundColor $color
        } else {
            Write-Host "  ${svc}: not installed" -ForegroundColor DarkGray
        }
    }

    Write-Host ""
}

Ensure-InstallLogDir
Write-Host "Install log file: $INSTALL_LOG" -ForegroundColor DarkGray
Start-Transcript -Path $INSTALL_LOG -Append | Out-Null

try {
    if ($Uninstall) {
        Uninstall-Services
    } elseif ($Start) {
        Start-Services
    } elseif ($Stop) {
        Stop-Services
    } elseif ($Status) {
        Show-Status
    } else {
        Install-Services
    }
} catch {
    Write-Host ""
    Write-Host "Install failed. Send this log file:" -ForegroundColor Red
    Write-Host "  $INSTALL_LOG" -ForegroundColor Red
    Write-Host "Error: $($_.Exception.Message)" -ForegroundColor Red
    throw
} finally {
    Stop-Transcript | Out-Null
}