#!/usr/bin/env pwsh
# update.ps1 - One-command update & restart for ocr-sprint-service (local dev)

param(
    [ValidateSet("cpu", "gpu")]
    [string] $OcrMode
)

$ErrorActionPreference = "Stop"

$Port = 8000
$ProjectRoot = $PSScriptRoot
$VenvDir = Join-Path $ProjectRoot ".venv"
$Python = Join-Path $VenvDir "Scripts\python.exe"

function Invoke-Step {
    param(
        [Parameter(Mandatory = $true)]
        [scriptblock] $Command,
        [Parameter(Mandatory = $true)]
        [string] $FailureMessage
    )

    & $Command
    if ($LASTEXITCODE -ne 0) {
        Write-Host "  $FailureMessage" -ForegroundColor Red
        exit $LASTEXITCODE
    }
}

function Get-DotEnvValue {
    param(
        [Parameter(Mandatory = $true)]
        [string] $Name
    )

    $envFile = Join-Path $ProjectRoot ".env"
    if (Test-Path $envFile) {
        $line = Get-Content $envFile | Where-Object { $_ -match "^\s*$Name\s*=" } | Select-Object -Last 1
        if ($line) {
            return (($line -split "=", 2)[1] -split "\s+#", 2)[0].Trim()
        }
    }
    return [Environment]::GetEnvironmentVariable($Name)
}

function Set-DotEnvValue {
    param(
        [Parameter(Mandatory = $true)]
        [string] $Name,
        [Parameter(Mandatory = $true)]
        [string] $Value
    )

    $envFile = Join-Path $ProjectRoot ".env"
    if (-not (Test-Path $envFile)) {
        New-Item -Path $envFile -ItemType File | Out-Null
    }

    $lines = @(Get-Content $envFile)
    $updated = $false
    for ($i = 0; $i -lt $lines.Count; $i++) {
        if ($lines[$i] -match "^\s*$Name\s*=") {
            $comment = ""
            if ($lines[$i] -match "(\s+#.*)$") {
                $comment = $Matches[1]
            }
            $lines[$i] = "$Name=$Value$comment"
            $updated = $true
        }
    }
    if (-not $updated) {
        $lines += "$Name=$Value"
    }
    Set-Content -Path $envFile -Value $lines
}

function Test-PythonPackage {
    param(
        [Parameter(Mandatory = $true)]
        [string] $Name
    )

    & $Python -m pip show $Name *> $null
    return $LASTEXITCODE -eq 0
}

function Add-NvidiaDllPaths {
    $dllDirs = @(
        (Join-Path $VenvDir "Lib\site-packages\nvidia\cudnn\bin"),
        (Join-Path $VenvDir "Lib\site-packages\nvidia\cublas\bin"),
        (Join-Path $VenvDir "Lib\site-packages\nvidia\cuda_nvrtc\bin")
    )
    foreach ($dir in $dllDirs) {
        if ((Test-Path $dir) -and (($env:PATH -split ";") -notcontains $dir)) {
            $env:PATH = "$dir;$env:PATH"
        }
    }
}

Set-Location $ProjectRoot

if (-not (Test-Path $Python)) {
    Write-Host "Virtualenv not found at $VenvDir. Creating one..." -ForegroundColor Yellow
    $venvCreated = $false
    $pythonLauncher = Get-Command py -ErrorAction SilentlyContinue
    if ($pythonLauncher) {
        foreach ($version in @("3.12", "3.11", "3.10")) {
            & py "-$version" -m venv $VenvDir 2>$null
            if ($LASTEXITCODE -eq 0) {
                $venvCreated = $true
                break
            }
        }
    }
    if (-not $venvCreated) {
        $systemPython = Get-Command python -ErrorAction SilentlyContinue
        if (-not $systemPython) {
            Write-Host "  Python was not found. Install Python 3.10-3.12, then rerun this script." -ForegroundColor Red
            exit 1
        }
        & python -m venv $VenvDir
        $venvCreated = ($LASTEXITCODE -eq 0)
    }
    if (-not $venvCreated) {
        Write-Host "  Failed to create virtualenv." -ForegroundColor Red
        exit $LASTEXITCODE
    }
}

$env:VIRTUAL_ENV = $VenvDir
$env:PATH = "$(Join-Path $VenvDir 'Scripts');$env:PATH"

if ($PSBoundParameters.ContainsKey("OcrMode")) {
    $ocrUseGpuValue = if ($OcrMode -eq "gpu") { "true" } else { "false" }
    Set-DotEnvValue "OCR_USE_GPU" $ocrUseGpuValue
    $env:OCR_USE_GPU = $ocrUseGpuValue
    Write-Host "OCR mode set to $($OcrMode.ToUpperInvariant()) and saved to .env." -ForegroundColor Green
}

# ── [1/5] Git pull ──────────────────────────────────────────────────────────
Write-Host "`n[1/5] Pulling latest code..." -ForegroundColor Cyan
Invoke-Step { git pull } "Git pull failed."

# ── [2/5] Install/update dependencies ───────────────────────────────────────
Write-Host "`n[2/5] Installing/updating dependencies..." -ForegroundColor Cyan
Invoke-Step { & $Python -m pip install -e ".[dev]" -q } "Dependency install failed."

$ocrUseGpu = (Get-DotEnvValue "OCR_USE_GPU")
if ($ocrUseGpu -and $ocrUseGpu.ToLowerInvariant() -in @("1", "true", "yes", "on")) {
    Write-Host "  GPU mode enabled; checking Paddle CUDA runtime..." -ForegroundColor Cyan
    if (-not (Test-PythonPackage "paddlepaddle-gpu")) {
        Invoke-Step {
            & $Python -m pip install paddlepaddle-gpu==2.6.2 -i https://www.paddlepaddle.org.cn/packages/stable/cu118/ -q
        } "Paddle GPU install failed."
    }
    if (-not (Test-PythonPackage "nvidia-cudnn-cu11")) {
        Invoke-Step { & $Python -m pip install nvidia-cudnn-cu11==8.9.5.29 -q } "NVIDIA cuDNN install failed."
    }
    Add-NvidiaDllPaths
} else {
    Write-Host "  CPU mode enabled; checking Paddle CPU runtime..." -ForegroundColor Cyan
    if (-not ((Test-PythonPackage "paddlepaddle") -or (Test-PythonPackage "paddlepaddle-gpu"))) {
        Invoke-Step { & $Python -m pip install paddlepaddle==2.6.2 -q } "Paddle CPU install failed."
    }
}

# ── [3/5] Database migration ─────────────────────────────────────────────────
Write-Host "`n[3/5] Running database migrations..." -ForegroundColor Cyan
& $Python -m alembic upgrade head
if ($LASTEXITCODE -ne 0) {
    Write-Host "  Migration conflict detected, stamping current state as head..." -ForegroundColor Yellow
    Invoke-Step { & $Python -m alembic stamp head } "Alembic stamp failed."
    Write-Host "  Retrying upgrade for any remaining new migrations..." -ForegroundColor Yellow
    & $Python -m alembic upgrade head
    if ($LASTEXITCODE -ne 0) {
        Write-Host "  Migration still failed. Please check alembic manually." -ForegroundColor Red
        exit 1
    }
}
Write-Host "  Migrations OK." -ForegroundColor Green

# ── [4/5] Free up port ───────────────────────────────────────────────────────
Write-Host "`n[4/5] Checking port $Port..." -ForegroundColor Cyan

# Use Get-NetTCPConnection for reliable port detection on Windows
$connections = Get-NetTCPConnection -LocalPort $Port -State Listen -ErrorAction SilentlyContinue
if ($connections) {
    foreach ($conn in $connections) {
        $procId = $conn.OwningProcess
        $procName = (Get-Process -Id $procId -ErrorAction SilentlyContinue).Name
        Write-Host "  Port $Port used by '$procName' (PID $procId), killing..." -ForegroundColor Yellow
        Stop-Process -Id $procId -Force -ErrorAction SilentlyContinue
    }
    # Wait until port is actually released (max 5 seconds)
    $waited = 0
    do {
        Start-Sleep -Milliseconds 500
        $waited += 500
        $still = Get-NetTCPConnection -LocalPort $Port -State Listen -ErrorAction SilentlyContinue
    } while ($still -and $waited -lt 5000)

    if ($still) {
        Write-Host "  Port $Port still in use after waiting. Try a different port or restart manually." -ForegroundColor Red
        exit 1
    }
    Write-Host "  Port $Port freed." -ForegroundColor Green
} else {
    Write-Host "  Port $Port is free." -ForegroundColor Green
}

# ── [5/5] Start dev server ───────────────────────────────────────────────────
Write-Host "`n[5/5] Starting dev server on port $Port (Ctrl+C to stop)..." -ForegroundColor Cyan
& $Python -m uvicorn ocr_sprint.main:app --reload --host 0.0.0.0 --port $Port