<#
.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' }

# Redirect the child's stdout/stderr to a log file. Without this, inheriting
# the parent shell's handles via Start-Process -WindowStyle Hidden can cause
# the child to exit immediately when those handles aren't real console fds
# (e.g. when restart.ps1 is driven from WSL/Git-Bash).
$logPath = Join-Path $env:TEMP ("ledgrab-{0}-{1}.log" -f $Module, $Port)
$errPath = "$logPath.err"
$argList = @()
$argList += $launchArgs
$argList += @('-m', $Module)
$startedProc = Start-Process -FilePath $pythonExe `
    -ArgumentList $argList `
    -WorkingDirectory $ServerRoot `
    -WindowStyle Hidden `
    -RedirectStandardOutput $logPath `
    -RedirectStandardError $errPath `
    -PassThru
$startedPid = $startedProc.Id

# ---- 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.
    $proc = Get-Process -Id $startedPid -ErrorAction SilentlyContinue
    if (-not $proc) { break }
    if (Test-PortOpen -Port $Port) { $ready = $true; break }
    Start-Sleep -Milliseconds 500
}

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

$proc = Get-Process -Id $startedPid -ErrorAction SilentlyContinue
if (-not $proc) {
    Write-Warning "Server process $startedPid exited before binding port $Port"
} else {
    Write-Warning "Server PID $startedPid is running but did not bind port $Port within ${StartupTimeoutSec}s"
}
if (Test-Path $errPath) {
    $tail = Get-Content $errPath -Tail 20 -ErrorAction SilentlyContinue
    if ($tail) {
        Write-Warning "Last stderr lines from $errPath :"
        $tail | ForEach-Object { Write-Warning "  $_" }
    }
}
exit 1