<#
.SYNOPSIS
  Restart a LedGrab Python server (real or demo) reliably.

.DESCRIPTION
  Gracefully asks the running instance to shut down via its HTTP API, waits
  for the port to free, then launches a detached replacement and polls the
  port until it is actually accepting connections.

  The script is parameterised so it works for the real server (default:
  port 8080, module `ledgrab`), the demo server, and any future variant —
  no code edits required to point it somewhere else.

.PARAMETER Port
  TCP port the server binds. Used both to locate the running process and
  to poll startup readiness.

.PARAMETER Module
  Python `-m` module to launch. Also used as a substring match when
  identifying which python.exe processes belong to this server so we don't
  kill unrelated Python instances.

.PARAMETER ServerRoot
  Working directory for the server process. Defaults to the directory that
  contains this script.

.PARAMETER ConfigPath
  Path (relative to -ServerRoot or absolute) to the YAML config the running
  server is using. Used only to read the API key for the graceful-shutdown
  request. If empty or missing we skip graceful shutdown and force-kill.

.PARAMETER StartupTimeoutSec
  How long to poll for the new server to start accepting connections.

.PARAMETER ShutdownTimeoutSec
  How long to wait for the graceful-shutdown API call to cause the running
  process to exit before force-killing it.

.PARAMETER SkipBrowser
  Set LEDGRAB_RESTART=1 in the child env so the app doesn't open a browser
  tab on startup. On by default — pass -SkipBrowser:$false to allow it.

.PARAMETER Quiet
  Suppress progress messages; only emit warnings/errors.

.EXAMPLE
  # Restart the real server (default invocation)
  powershell -ExecutionPolicy Bypass -File restart.ps1

.EXAMPLE
  # Restart the demo server on port 8081
  powershell -ExecutionPolicy Bypass -File restart.ps1 `
      -Port 8081 -Module ledgrab.demo -ConfigPath 'config\demo_config.yaml'

.NOTES
  Exit codes:
    0 — server is up and accepting connections on the target port
    1 — startup timed out; process may or may not be running
    2 — could not locate a Python interpreter
#>
[CmdletBinding()]
param(
    [int]$Port = 8080,
    [string]$Module = 'ledgrab',
    [string]$ServerRoot = '',
    [string]$ConfigPath = 'config\default_config.yaml',
    [int]$StartupTimeoutSec = 30,
    [int]$ShutdownTimeoutSec = 15,
    [string]$PythonExe = '',
    [string]$PythonVersion = '3.13',
    [switch]$SkipBrowser = $true,
    [switch]$Quiet
)

$ErrorActionPreference = 'Stop'

function Write-Info {
    param([string]$Message)
    if (-not $Quiet) { Write-Host $Message }
}

# ---- Resolve paths ---------------------------------------------------------

# PS 5.1 doesn't expand $PSScriptRoot at param-binding time, so apply it here.
if (-not $ServerRoot) { $ServerRoot = $PSScriptRoot }
if (-not $ServerRoot) {
    Write-Error 'ServerRoot not provided and $PSScriptRoot is unavailable'
    exit 2
}
if (-not (Test-Path $ServerRoot)) {
    Write-Error "ServerRoot '$ServerRoot' does not exist"
    exit 2
}
$ServerRoot = (Resolve-Path $ServerRoot).Path

$resolvedConfig = $null
if ($ConfigPath) {
    $candidate = if ([IO.Path]::IsPathRooted($ConfigPath)) {
        $ConfigPath
    } else {
        Join-Path $ServerRoot $ConfigPath
    }
    if (Test-Path $candidate) { $resolvedConfig = $candidate }
}

# ---- Locate the running server ---------------------------------------------

function Get-ServerProcesses {
    param([string]$ModuleName, [string]$Root)
    # Match python.exe processes whose command line references this module AND
    # whose cwd (via command line fragment) looks like it's running from this
    # server root. Excludes unrelated python.exe (VS Code extensions, isort,
    # pip tooling, etc.) by requiring a module reference.
    $rootPattern = [regex]::Escape($Root)
    Get-CimInstance Win32_Process -Filter "Name='python.exe'" -ErrorAction SilentlyContinue |
        Where-Object {
            $cl = $_.CommandLine
            if (-not $cl) { return $false }
            # Must launch the target module via `-m <Module>` or an exact token
            $launchesModule = $cl -match ('-m\s+' + [regex]::Escape($ModuleName) + '(\s|$|\.)')
            if (-not $launchesModule) { return $false }
            # Exclude obvious tooling false-positives
            if ($cl -match '(vscode|isort|pip[-\s]|flake8|ruff|mypy|pylint|black)') {
                return $false
            }
            return $true
        }
}

function Test-PortOpen {
    param([int]$Port)
    try {
        $listener = Get-NetTCPConnection -LocalPort $Port -State Listen -ErrorAction Stop
        return [bool]$listener
    } catch {
        return $false
    }
}

$existing = Get-ServerProcesses -ModuleName $Module -Root $ServerRoot

# ---- Graceful shutdown (if the target is currently up) ---------------------

if ($existing) {
    $apiKey = $null
    if ($resolvedConfig) {
        # Pull the first api_keys entry — good enough for the local shutdown
        # endpoint; production deploys don't use this script.
        $inKeys = $false
        foreach ($line in Get-Content $resolvedConfig) {
            if ($line -match '^\s*api_keys:') { $inKeys = $true; continue }
            if ($inKeys -and $line -match '^\s+\w+:\s*"(.+)"') {
                $apiKey = $Matches[1]; break
            }
            if ($inKeys -and $line -match '^\S') { break }
        }
    }

    $shutdownRequested = $false
    if ($apiKey) {
        Write-Info 'Requesting graceful shutdown...'
        try {
            $headers = @{ Authorization = "Bearer $apiKey" }
            Invoke-RestMethod -Uri "http://localhost:$Port/api/v1/system/shutdown" `
                -Method Post -Headers $headers -TimeoutSec 5 -ErrorAction Stop | Out-Null
            $shutdownRequested = $true
        } catch {
            Write-Info "  API shutdown failed ($($_.Exception.Message)); will force-kill"
        }
    }

    if ($shutdownRequested) {
        Write-Info 'Waiting for graceful shutdown...'
        $waited = 0
        while ($waited -lt $ShutdownTimeoutSec) {
            Start-Sleep -Seconds 1
            $waited++
            if (-not (Get-ServerProcesses -ModuleName $Module -Root $ServerRoot)) {
                Write-Info "  Exited cleanly after ${waited}s"
                break
            }
        }
    }

    $still = Get-ServerProcesses -ModuleName $Module -Root $ServerRoot
    if ($still) {
        Write-Info '  Force-killing remaining processes...'
        foreach ($p in $still) {
            Write-Info "    Stop PID $($p.ProcessId)"
            Stop-Process -Id $p.ProcessId -Force -ErrorAction SilentlyContinue
        }
    }

    # Wait for Windows to release the TCP socket before we rebind. A fixed
    # 1–2 s sleep isn't enough on machines where the kernel lingers in
    # CLOSE_WAIT; poll the port state instead.
    $portDeadline = (Get-Date).AddSeconds(10)
    while ((Get-Date) -lt $portDeadline -and (Test-PortOpen -Port $Port)) {
        Start-Sleep -Milliseconds 250
    }
}

# ---- Merge per-user PATH (captures tools installed after the shell started) ----

$regUser = [Environment]::GetEnvironmentVariable('PATH', 'User')
if ($regUser) {
    $currentDirs = $env:PATH -split ';' | ForEach-Object { $_.TrimEnd('\') }
    foreach ($dir in ($regUser -split ';')) {
        if ($dir -and ($currentDirs -notcontains $dir.TrimEnd('\'))) {
            $env:PATH = "$env:PATH;$dir"
        }
    }
}

# ---- Locate a Python interpreter -------------------------------------------

# We need the Python that actually has the target module installed. Naively
# resolving `python` on PATH can pick up 3.11 or another version that doesn't
# have `ledgrab` in its site-packages, so prefer an explicit interpreter in
# this priority order:
#   1. -PythonExe (caller override)
#   2. `py -<Version>` via the Windows Python launcher
#   3. A Python<Version> install under %LOCALAPPDATA%\Programs\Python
#   4. `python` on PATH (last-resort fallback)

function Test-HasModule {
    param([string]$Exe, [string]$ModuleName)
    if (-not $Exe -or -not (Test-Path $Exe)) { return $false }
    & $Exe -c "import importlib.util, sys; sys.exit(0 if importlib.util.find_spec('$ModuleName') else 1)" 2>$null
    return ($LASTEXITCODE -eq 0)
}

$resolvedPython = $null
$launchArgs = @()

if ($PythonExe) {
    if (-not (Test-Path $PythonExe)) {
        Write-Error "PythonExe '$PythonExe' does not exist"
        exit 2
    }
    $resolvedPython = (Resolve-Path $PythonExe).Path
} else {
    # Try `py -<version>`
    $pyLauncher = (Get-Command py -ErrorAction SilentlyContinue).Source
    if ($pyLauncher) {
        $probe = & $pyLauncher "-$PythonVersion" -c "import sys; print(sys.executable)" 2>$null
        if ($LASTEXITCODE -eq 0 -and $probe) {
            $resolvedPython = $pyLauncher
            $launchArgs = @("-$PythonVersion")
        }
    }
    # Fall back to a known install path for that version
    if (-not $resolvedPython) {
        $verTag = $PythonVersion -replace '\.', ''
        $candidate = Join-Path $env:LOCALAPPDATA "Programs\Python\Python$verTag\python.exe"
        if (Test-Path $candidate) { $resolvedPython = $candidate }
    }
    # Last resort: plain `python` on PATH
    if (-not $resolvedPython) {
        $onPath = (Get-Command python -ErrorAction SilentlyContinue).Source
        if ($onPath) { $resolvedPython = $onPath }
    }
}

if (-not $resolvedPython) {
    Write-Error "No Python $PythonVersion interpreter found (tried: -PythonExe, py -$PythonVersion, %LOCALAPPDATA%\Programs\Python\Python*, PATH)"
    exit 2
}

# Verify the module is actually importable with the chosen interpreter so we
# don't launch a process that would immediately die with "No module named X".
# When using the `py` launcher, delegate to the versioned interpreter.
$effectiveExe = if ($launchArgs.Count -gt 0) {
    & $resolvedPython @launchArgs -c "import sys; print(sys.executable)" 2>$null
} else {
    $resolvedPython
}

if (-not (Test-HasModule -Exe $effectiveExe -ModuleName $Module)) {
    Write-Error "Module '$Module' is not importable with $effectiveExe. Install it (e.g. pip install -e .) or pass -PythonExe pointing to the right interpreter."
    exit 2
}

$pythonExe = $resolvedPython

# ---- Launch detached replacement -------------------------------------------

Write-Info "Starting $Module on port $Port..."
if ($SkipBrowser) { $env:LEDGRAB_RESTART = '1' }

# Launch python.exe directly with no parent-handle inheritance. We used to
# wrap it in `cmd /c python ... 1>log 2>err` so the parent powershell could
# tail crash logs, but that left an empty cmd.exe window hanging around for
# the full server lifetime (cmd had to live to hold the redirect handles).
# Instead, let python claim its own console window — the user sees the live
# server log there, and there's no spurious cmd window.
#
# Why WMI Win32_Process.Create rather than Start-Process or
# [Diagnostics.Process]::Start? Both of those go through CreateProcess with
# bInheritHandles=true, which leaks the parent shell's pipe handles into
# the new Python process. When the caller is Git-Bash (`restart.ps1 |
# tail -10`), the bash pipe then stays open for the full server lifetime,
# hanging the bash invocation even after powershell exits. WMI's
# Win32_Process.Create uses CreateProcess with bInheritHandles=FALSE.

$argList = @()
$argList += $launchArgs
$argList += @('-m', $Module)

# Quote each arg defensively in case a future caller adds whitespace.
function Quote-CmdArg {
    param([string]$Arg)
    if ($Arg -match '[\s"]') {
        return '"' + ($Arg -replace '"', '\"') + '"'
    }
    return $Arg
}
$quotedArgs = ($argList | ForEach-Object { Quote-CmdArg $_ }) -join ' '
$pyQ = Quote-CmdArg $pythonExe

$cmdLine = $pyQ + ' ' + $quotedArgs

# Win32_Process.Create starts detached with no parent-handle inheritance.
# Returns @{ ProcessId; ReturnValue (0 = success) }.
# Title sets the visible console-window title so the user can tell at a
# glance which server the window belongs to (useful when running real +
# demo side by side on different ports).
$startupInfo = New-CimInstance -ClassName Win32_ProcessStartup `
    -ClientOnly `
    -Property @{ Title = "LedGrab - $Module (port $Port)" }
$wmiResult = Invoke-CimMethod -ClassName Win32_Process -MethodName Create -Arguments @{
    CommandLine               = $cmdLine
    CurrentDirectory          = $ServerRoot
    ProcessStartupInformation = $startupInfo
} -ErrorAction SilentlyContinue

if (-not $wmiResult -or $wmiResult.ReturnValue -ne 0) {
    Write-Warning "WMI Win32_Process.Create failed (ReturnValue=$($wmiResult.ReturnValue)); falling back to Start-Process"
    # Fallback path — Start-Process inherits parent handles, so a piped
    # caller may hang. Acceptable here because this branch only runs when
    # WMI itself is broken (very rare).
    $startedProc = Start-Process -FilePath $pythonExe `
        -ArgumentList $argList `
        -WorkingDirectory $ServerRoot -PassThru
    $startedPid = if ($startedProc) { $startedProc.Id } else { 0 }
} else {
    $startedPid = [int]$wmiResult.ProcessId
}

# Confirm the process is actually our server (defensive — WMI sometimes
# returns a PID for a transient ancestor on heavily loaded boxes).
Start-Sleep -Milliseconds 250
if (-not (Get-Process -Id $startedPid -ErrorAction SilentlyContinue)) {
    $rescanned = Get-ServerProcesses -ModuleName $Module -Root $ServerRoot | Select-Object -First 1
    if ($rescanned) { $startedPid = $rescanned.ProcessId } else { $startedPid = 0 }
}

# ---- Poll readiness --------------------------------------------------------

# Port readiness is the authoritative signal — the process can be alive for
# many seconds before uvicorn finishes binding on cold starts (store init,
# etc.). Polling avoids spurious "not running" warnings that the old fixed
# 3-second sleep produced.
$deadline = (Get-Date).AddSeconds($StartupTimeoutSec)
$ready = $false
while ((Get-Date) -lt $deadline) {
    # Bail early if the process has already exited — something went wrong.
    if ($startedPid -gt 0) {
        $proc = Get-Process -Id $startedPid -ErrorAction SilentlyContinue
        if (-not $proc) {
            $rescanned = Get-ServerProcesses -ModuleName $Module -Root $ServerRoot | Select-Object -First 1
            if ($rescanned) { $startedPid = $rescanned.ProcessId } else { break }
        }
    } else {
        $rescanned = Get-ServerProcesses -ModuleName $Module -Root $ServerRoot | Select-Object -First 1
        if ($rescanned) { $startedPid = $rescanned.ProcessId }
    }
    if (Test-PortOpen -Port $Port) { $ready = $true; break }
    Start-Sleep -Milliseconds 500
}

if ($ready) {
    if ($startedPid -gt 0) {
        Write-Info "Server ready on port $Port (PID $startedPid)"
    } else {
        Write-Info "Server ready on port $Port"
    }
    exit 0
}

if ($startedPid -gt 0) {
    $proc = Get-Process -Id $startedPid -ErrorAction SilentlyContinue
    if (-not $proc) {
        Write-Warning "Server process $startedPid exited before binding port $Port (check the server console window for the error)"
    } else {
        Write-Warning "Server PID $startedPid is running but did not bind port $Port within ${StartupTimeoutSec}s"
    }
} else {
    Write-Warning "Could not locate server process; port $Port did not bind within ${StartupTimeoutSec}s"
}
exit 1
