#requires -Version 5.1
<#
  Install-MajicMediaNode.ps1 — MAJIC media-node installer (Phase 2b).

  On a node it:
    1. Inspects the host GPU (nvidia-smi / WMI) -> model + VRAM.
    2. Picks the BEST image/video model that FITS the VRAM (documented, generalizable
       tier logic in MajicMediaNode.psm1; first target GTX 1060 6GB -> SD 1.5, image-only).
    3. Installs + configures that model behind a small local OpenAI-ish HTTP endpoint
       (media_node_server.py). Model PERSISTS in VRAM but YIELDS when the user needs the
       card and reloads when idle (idle-only contribution).
    4. Provisions THIS node's Cloudflare tunnel + DNS for the media endpoint (reuses the
       Enable-MajicCfAccess pattern) so the orchestrator can reach it regardless of WAN IP.
    5. CAPABILITY PROBE: image runtime, video feasibility + max_video_seconds, runtimes.
    6. SELF-REGISTERS {id, cf endpoint, gpu, model, capabilities, idle_policy} to the
       STANDALONE enroll endpoint, and installs a heartbeat task for idle-online transitions.

  GUARDRAILS: additive; backs up an existing config; never touches the chat bridge or its
  services; idle-yield leaves headroom so it does not fight the user or HV1's existing role.

  USAGE (on the node, elevated):
    powershell -ExecutionPolicy Bypass -File Install-MajicMediaNode.ps1 `
        -RegistryUrl https://registry.majicholdings.com `
        -ShareUnc \\MEDIA-HOST\media `
        -CfAccount 5d5385be79ab1dfa09def1fe76e04d73 `
        -CfGlobalKey <key> -CfEmail nmorris@greencommllc.com
#>
[CmdletBinding()]
param(
    [string]$RegistryUrl = 'https://registry.majicholdings.com',
    [string]$ShareUnc    = '\\MEDIA-HOST\media',
    [string]$Zone        = 'majicholdings.com',
    [int]$MediaPort      = 9120,
    [string]$CfAccount   = '5d5385be79ab1dfa09def1fe76e04d73',
    [string]$CfGlobalKey = '',
    [string]$CfEmail     = 'nmorris@greencommllc.com',
    [string]$Root        = 'C:\Majic\media',
    [switch]$SkipDeps,                 # skip the (large) torch/diffusers pip install
    [switch]$NoTunnel                  # skip CF tunnel (LAN-only test)
)
$ErrorActionPreference = 'Stop'
[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12

$here = Split-Path -Parent $MyInvocation.MyCommand.Path
Import-Module (Join-Path $here 'MajicMediaNode.psm1') -Force

Write-Step '==================== MAJIC media-node installer ====================' Cyan
New-Item -ItemType Directory -Force -Path $Root | Out-Null

# ---- 1. GPU ----------------------------------------------------------------
Write-Step '[1/7] Inspecting GPU ...' Cyan
$gpu = Get-MajicGpuInfo
if (-not $gpu.name) { Write-Step '  ERROR: no GPU detected (nvidia-smi/WMI). A media node needs an NVIDIA GPU.' Red; exit 2 }
$gpuHash = @{ name=$gpu.name; vramMB=$gpu.vramMB; driver=$gpu.driver }
Write-Step ("  GPU: {0} | VRAM {1} MB | driver {2} | via {3}" -f $gpu.name,$gpu.vramMB,$gpu.driver,$gpu.source) Green

# ---- 2. Model selection ----------------------------------------------------
Write-Step '[2/7] Selecting best-fit media model ...' Cyan
$model = Select-MajicMediaModel -Gpu $gpuHash
Write-Step ("  -> {0} (kind={1}, tier={2}, arch={3}) image={4} video={5}" -f `
    $model.model_id,$model.model_kind,$model.tier,$model.arch,$model.image,$model.video) Green

# ---- 3. Python deps + server + config + service ----------------------------
Write-Step '[3/7] Installing media server + dependencies ...' Cyan
$py = (Get-Command python -ErrorAction SilentlyContinue).Source
if (-not $py) { $py = (Get-Command python3 -ErrorAction SilentlyContinue).Source }
if (-not $py) { Write-Step '  ERROR: python not found on PATH. Install Python 3.10+ first.' Red; exit 3 }
Write-Step "  python: $py"

# Copy the server next to the install root.
Copy-Item (Join-Path $here 'media_node_server.py') (Join-Path $Root 'media_node_server.py') -Force

# Dedicated venv so we never touch any system/bridge python env.
$venv = Join-Path $Root 'venv'
if (-not (Test-Path (Join-Path $venv 'Scripts\python.exe'))) {
    & $py -m venv $venv
}
$venvPy = Join-Path $venv 'Scripts\python.exe'

if (-not $SkipDeps) {
    Write-Step '  installing CUDA torch + diffusers + transformers + accelerate (large, be patient) ...' Cyan
    & $venvPy -m pip install --upgrade pip 2>&1 | Out-Null

    # CRITICAL: install a CUDA torch FIRST and verify it, THEN diffusers/etc. If we let
    # diffusers pull torch as a transitive dep it resolves the CPU wheel from PyPI, which
    # leaves the GPU unused. The right CUDA channel depends on the venv's Python tag:
    # cu121 has no cp313 wheels, so on Python 3.13 we must use cu124+. A GTX 1060 is
    # Pascal (sm_61) — supported through CUDA 12.4 (torch 2.6 cu124); newer cu126/cu128
    # builds still carry sm_61 too. We try channels in order and keep the first whose
    # torch.cuda.is_available() comes back true; CPU is the explicit last resort.
    $cudaOk = $false
    foreach ($ch in @('cu121','cu124','cu126')) {
        Write-Step "    trying torch from $ch ..." Gray
        & $venvPy -m pip install --index-url "https://download.pytorch.org/whl/$ch" torch 2>&1 | Select-Object -Last 2 | ForEach-Object { Write-Step "      $_" }
        $probe = (& $venvPy -c "import torch,sys; sys.stdout.write('1' if torch.cuda.is_available() else '0')" 2>$null)
        if ($probe -eq '1') { $cudaOk = $true; Write-Step "    CUDA torch OK via $ch." Green; break }
        Write-Step "    $ch did not yield a CUDA-capable torch; trying next." Yellow
    }
    if (-not $cudaOk) {
        Write-Step '  ! no CUDA torch wheel matched this Python/GPU; falling back to CPU torch (slow).' Yellow
        & $venvPy -m pip install torch 2>&1 | Out-Null
    }
    # Pin the resolved torch so the next install can't downgrade it to a CPU build.
    $torchVer = (& $venvPy -c "import torch,sys; sys.stdout.write(torch.__version__)" 2>$null)
    Write-Step "  torch resolved: $torchVer" Green
    # Pin transformers/diffusers to a known-compatible pair: transformers 5.x dropped the
    # top-level CLIPTextModel import path that diffusers' SD pipelines rely on, so the
    # latest-of-both combo fails at pipeline import ("Could not import module CLIPTextModel").
    # transformers 4.49 + diffusers 0.32 is a verified-good SD1.5/SDXL pair on torch 2.6.
    & $venvPy -m pip install "diffusers==0.32.2" "transformers==4.49.0" "huggingface-hub<1.0" "tokenizers<0.22" accelerate safetensors pillow "torch==$torchVer" 2>&1 | Select-Object -Last 3 | ForEach-Object { Write-Step "    $_" }
} else {
    Write-Step '  -SkipDeps set: not installing torch/diffusers (server will report model_loaded=false until installed).' Yellow
}

# ---- write node config -----------------------------------------------------
$cfgPath = Join-Path $Root 'node-config.json'
if (Test-Path $cfgPath) { Copy-Item $cfgPath "$cfgPath.bak-$(Get-Date -Format yyyyMMddHHmmss)" -Force; Write-Step "  backed up existing config." }
$config = [ordered]@{
    node_id            = $env:COMPUTERNAME.ToLower()
    listen_host        = '127.0.0.1'
    listen_port        = $MediaPort
    model_id           = $model.model_id
    model_kind         = $model.model_kind
    video_enabled      = [bool]$model.video
    max_video_seconds  = 0
    share_unc          = $ShareUnc
    vram_total_mb      = [int]$gpu.vramMB
    yield_free_vram_floor_mb = [int]$model.yield_free_vram_floor_mb
    idle_reload_seconds= 20
    image_steps        = [int]$model.image_steps
    image_size         = [int]$model.image_size
}
$config | ConvertTo-Json -Depth 6 | Set-Content -LiteralPath $cfgPath -Encoding UTF8
Write-Step "  config -> $cfgPath" Green

# ---- register the server as a SYSTEM service (schtask, auto-restart) --------
# Use a SYSTEM scheduled task at boot (no NSSM dependency). The yield-watcher inside
# the server handles GPU yield/reload; the task restarts the process if it exits.
$svcScript = Join-Path $Root 'run-media-node.cmd'
@"
@echo off
set MAJIC_MEDIA_NODE_CONFIG=$cfgPath
:loop
"$venvPy" "$Root\media_node_server.py"
timeout /t 5 /nobreak >nul
goto loop
"@ | Set-Content -LiteralPath $svcScript -Encoding ASCII
$taskName = 'MajicMediaNode'
try { schtasks /Query /TN $taskName >$null 2>&1; if ($LASTEXITCODE -eq 0) { schtasks /Delete /TN $taskName /F | Out-Null } } catch {}
schtasks /Create /TN $taskName /SC ONSTART /RU SYSTEM /RL HIGHEST /TR "cmd /c `"$svcScript`"" /F | Out-Null
schtasks /Run /TN $taskName | Out-Null
Write-Step "  media server registered as SYSTEM task '$taskName' (starts on boot, self-restarts)." Green

# ---- 4. CF tunnel ----------------------------------------------------------
Write-Step '[4/7] Provisioning Cloudflare tunnel + DNS for the media endpoint ...' Cyan
$cfEndpoint = $null
if (-not $NoTunnel) {
    $cfEndpoint = Set-MajicMediaTunnel -Port $MediaPort -Zone $Zone -CfAccount $CfAccount -CfGlobalKey $CfGlobalKey -CfEmail $CfEmail
}
if (-not $cfEndpoint) {
    $cfEndpoint = "http://$($env:COMPUTERNAME.ToLower()):$MediaPort"   # LAN fallback so the row is still routable on-LAN
    Write-Step "  ! using LAN fallback endpoint $cfEndpoint (no CF tunnel)." Yellow
}
Write-Step "  media endpoint: $cfEndpoint" Green

# ---- 5. Capability probe ---------------------------------------------------
Write-Step '[5/7] Probing capabilities (runtimes + video second-limit) ...' Cyan
$caps = Test-MajicNodeCapabilities -Model $model -Gpu $gpuHash
# reflect the probed max_video_seconds back into the running config
$config.max_video_seconds = [int]$caps.max_video_seconds
$config | ConvertTo-Json -Depth 6 | Set-Content -LiteralPath $cfgPath -Encoding UTF8
Write-Step ("  image={0} video={1} max_video_seconds={2} img_runtime~{3}s video_runtime~{4}s" -f `
    $caps.image,$caps.video,$caps.max_video_seconds,$caps.expected_image_runtime,$caps.expected_video_runtime) Green

# ---- 6. wait for the local endpoint, then self-register --------------------
Write-Step '[6/7] Waiting for the local media endpoint + self-registering ...' Cyan
$healthUrl = "http://127.0.0.1:$MediaPort/health"
$up = $false
for ($i=1; $i -le 20 -and -not $up; $i++) {
    try { $h = Invoke-RestMethod -Uri $healthUrl -TimeoutSec 5; if ($h.ok) { $up = $true; break } } catch { Start-Sleep -Seconds 2 }
}
if ($up) { Write-Step "  local endpoint UP: model_loaded=$($h.model_loaded) $(if(-not $h.model_loaded){"(loading: $($h.load_error))"})" Green }
else { Write-Step "  ! local endpoint did not answer on $healthUrl yet (deps may still be installing); registering anyway." Yellow }

Register-MajicMediaNode -RegistryUrl $RegistryUrl -CfEndpoint $cfEndpoint -Gpu $gpuHash -Model $model -Capabilities $caps -IdlePolicy 'idle-only' | Out-Null

# ---- 7. heartbeat task (idle-online transitions) ---------------------------
Write-Step '[7/7] Installing heartbeat task (re-register on idle-online) ...' Cyan
$hbScript = Join-Path $Root 'heartbeat.ps1'
@"
Import-Module '$here\MajicMediaNode.psm1' -Force -ErrorAction SilentlyContinue
if (-not (Get-Command Register-MajicMediaNode -EA SilentlyContinue)) { Import-Module '$Root\MajicMediaNode.psm1' -Force -EA SilentlyContinue }
`$cfg = Get-Content -Raw '$cfgPath' | ConvertFrom-Json
`$gpu = @{ name='$($gpu.name)'; vramMB=$($gpu.vramMB) }
`$model = @{ model_id=`$cfg.model_id }
`$caps = @{ image=`$cfg.video_enabled -or `$true; video=`$cfg.video_enabled; max_video_seconds=`$cfg.max_video_seconds; expected_image_runtime=$($caps.expected_image_runtime); expected_video_runtime=$($caps.expected_video_runtime) }
Register-MajicMediaNode -RegistryUrl '$RegistryUrl' -CfEndpoint '$cfEndpoint' -Gpu `$gpu -Model `$model -Capabilities `$caps -Heartbeat | Out-Null
"@ | Set-Content -LiteralPath $hbScript -Encoding UTF8
Copy-Item (Join-Path $here 'MajicMediaNode.psm1') (Join-Path $Root 'MajicMediaNode.psm1') -Force
$hbTask = 'MajicMediaNodeHeartbeat'
try { schtasks /Query /TN $hbTask >$null 2>&1; if ($LASTEXITCODE -eq 0) { schtasks /Delete /TN $hbTask /F | Out-Null } } catch {}
schtasks /Create /TN $hbTask /SC MINUTE /MO 2 /RU SYSTEM /RL HIGHEST /TR "powershell -NoProfile -ExecutionPolicy Bypass -File `"$hbScript`"" /F | Out-Null
Write-Step "  heartbeat task '$hbTask' every 2 min." Green

Write-Step ''
Write-Step '==================== media node install complete ====================' Cyan
Write-Step ("  node      : {0}" -f $env:COMPUTERNAME.ToLower())
Write-Step ("  endpoint  : {0}" -f $cfEndpoint)
Write-Step ("  model     : {0} (image={1}, video={2}, max_video_seconds={3})" -f $model.model_id,$caps.image,$caps.video,$caps.max_video_seconds)
Write-Step ("  registry  : {0}" -f $RegistryUrl)
Write-Step ("  share     : {0}" -f $ShareUnc)
exit 0