# EthoVision worker — one-step installer for Windows (PowerShell).
#
# Run with:
#   powershell -ExecutionPolicy Bypass -File install.ps1 -Server <url> -Token <token>
# or set the env vars and call with no args.

param(
    [string]$Server = $env:ETHOVISION_SERVER,
    [string]$Token = $env:ETHOVISION_TOKEN,
    [string]$EnvName = "ethovision"
)

$ErrorActionPreference = "Stop"
$Here = Split-Path -Parent $MyInvocation.MyCommand.Path
$PyVer = "3.11"

function Log  ($msg) { Write-Host ">> $msg" -ForegroundColor Cyan }
function Warn ($msg) { Write-Host "** $msg" -ForegroundColor Yellow }
function Ok   ($msg) { Write-Host "OK $msg" -ForegroundColor Green }
function Err  ($msg) { Write-Host "!! $msg" -ForegroundColor Red }

Log "EthoVision worker installer (Windows)"
Log "working dir: $Here, env: $EnvName"

# ---------- 1. Conda ----------
function Use-Conda {
    # Prefer a conda already on PATH
    $inPath = Get-Command conda -ErrorAction SilentlyContinue
    if ($inPath) {
        Ok "conda already installed"
        return $inPath.Source
    }
    # Existing user Miniconda
    $userConda = Join-Path $env:USERPROFILE "miniconda3\Scripts\conda.exe"
    if (Test-Path $userConda) {
        Ok "using existing $userConda"
        return $userConda
    }
    # Fresh install
    Log "installing Miniconda to $env:USERPROFILE\miniconda3 (silent, ~100 MB)"
    $installer = Join-Path $env:TEMP "miniconda-installer.exe"
    $url = "https://repo.anaconda.com/miniconda/Miniconda3-latest-Windows-x86_64.exe"
    Invoke-WebRequest -Uri $url -OutFile $installer -UseBasicParsing
    Start-Process -FilePath $installer -ArgumentList "/S","/AddToPath=0","/RegisterPython=0","/D=$env:USERPROFILE\miniconda3" -Wait
    Remove-Item $installer
    if (-not (Test-Path $userConda)) { throw "Miniconda install failed - $userConda not found" }
    Ok "Miniconda installed"
    return $userConda
}
$conda = Use-Conda

# Run a command inside the env without needing activation (conda run handles it).
function InEnv {
    param([Parameter(ValueFromRemainingArguments=$true)][string[]]$Args)
    & $conda run -n $EnvName --no-capture-output @Args
    if ($LASTEXITCODE -ne 0) { throw ("in-env command failed (exit " + $LASTEXITCODE + "): " + ($Args -join ' ')) }
}

# ---------- 2. Env ----------
function Test-EnvExists {
    param([string]$Name)
    # Primary check: env directory in the standard location
    try {
        $condaBase = Split-Path (Split-Path $conda -Parent) -Parent
        $envDir = Join-Path $condaBase "envs\$Name"
        if (Test-Path $envDir) { return $true }
    } catch { }
    # Fallback: parse `conda env list`
    try {
        $listing = (& $conda env list 2>&1) | Out-String
        return ($listing -match "(?m)^\s*$([regex]::Escape($Name))(\s|\*)")
    } catch { return $false }
}

if (Test-EnvExists $EnvName) {
    Ok "env '$EnvName' already exists, reusing"
} else {
    # Accept Anaconda's Terms of Service for default channels (added mid-2024).
    # Without this, `conda create` fails with CondaToSNonInteractiveError.
    # On older conda versions without the tos plugin this is a no-op.
    Log "accepting Anaconda Terms of Service for default channels"
    foreach ($ch in @(
        "https://repo.anaconda.com/pkgs/main",
        "https://repo.anaconda.com/pkgs/r",
        "https://repo.anaconda.com/pkgs/msys2"
    )) {
        & $conda tos accept --override-channels --channel $ch 2>&1 | Out-Null
    }

    Log "creating env '$EnvName' with Python $PyVer"
    # Don't swallow output — user needs to see if conda hits a proxy/SSL/disk error.
    & $conda create -n $EnvName "python=$PyVer" -y
    if ($LASTEXITCODE -ne 0) {
        Err "conda create failed with exit code $LASTEXITCODE"
        Write-Host ""
        Write-Host "If an earlier install attempt left a broken env, clean it with:" -ForegroundColor Yellow
        Write-Host "    & `"$conda`" env remove -n $EnvName -y" -ForegroundColor Yellow
        Write-Host "Then re-run install.ps1." -ForegroundColor Yellow
        throw "conda create failed"
    }
}

Log "upgrading pip"
InEnv python -m pip install --upgrade pip --quiet

# ---------- 3. PyTorch ----------
$torchVariant = "cpu"
$nvidiaSmi = "$env:SystemRoot\System32\nvidia-smi.exe"
if (Test-Path $nvidiaSmi) {
    try {
        $driver = & $nvidiaSmi --query-gpu=driver_version --format=csv,noheader
        $driverMajor = [int]($driver.Split(".")[0])
        Log "NVIDIA driver detected: $driverMajor"
        if     ($driverMajor -ge 550) { $torchVariant = "cu124" }
        elseif ($driverMajor -ge 525) { $torchVariant = "cu121" }
        elseif ($driverMajor -ge 470) { $torchVariant = "cu118" }
        else  {
            Warn "NVIDIA driver too old for current PyTorch CUDA builds (need 470+). Falling back to CPU."
            $torchVariant = "cpu"
        }
    } catch {
        Warn "could not read nvidia-smi output, falling back to CPU"
    }
} else {
    Warn "no NVIDIA GPU detected - installing CPU-only PyTorch (tracking will be ~20x slower)"
}

Log "installing PyTorch ($torchVariant) - this is the big download (~2 GB)"
if ($torchVariant -eq "cpu") {
    InEnv python -m pip install torch torchvision --index-url "https://download.pytorch.org/whl/cpu"
} else {
    InEnv python -m pip install torch torchvision --index-url "https://download.pytorch.org/whl/$torchVariant"
}

# ---------- 4. DLC ----------
# Notes for future maintainers:
#  - The [modelzoo] extra was removed in DLC 3; using it causes pip to fall back
#    to a DLC 2.x release that still depends on TensorFlow.
#  - DLC 3.x is currently only published as release candidates (3.0.0rc*), so
#    we pass --pre to let pip consider them.
Log "installing DeepLabCut 3.x (PyTorch-based, SuperAnimal built-in)"
InEnv python -m pip install --pre "deeplabcut>=3.0.0rc1" requests

# ---------- 5. Worker script ----------
$workerPath = Join-Path $Here "ethovision_worker.py"
if (-not (Test-Path $workerPath)) {
    if ($Server) {
        Log "downloading ethovision_worker.py from $Server"
        Invoke-WebRequest -Uri "$Server/api/worker/ethovision_worker.py" -OutFile $workerPath -UseBasicParsing
    } else {
        Warn "ethovision_worker.py not next to install.ps1, and -Server not set - skipping download"
    }
}

# ---------- 6. Launcher ----------
if ($Server -and $Token) {
    Log "writing run.bat with your server + token"
    $condaEscaped = $conda.Replace('"', '""')
    $runBat = @"
@echo off
set ETHOVISION_SERVER=$Server
set ETHOVISION_TOKEN=$Token
"$condaEscaped" run -n $EnvName --no-capture-output python "%~dp0ethovision_worker.py" %*
"@
    Set-Content -Path (Join-Path $Here "run.bat") -Value $runBat -Encoding ASCII
    Ok "launcher written: $Here\run.bat"
} else {
    Warn "-Server and -Token not both set - skipping run.bat launcher"
}

# ---------- 7. Sanity checks ----------
Log "verifying install"
InEnv python -c "import torch; print('  torch:', torch.__version__); print('  CUDA available:', torch.cuda.is_available()); print('  GPU:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else '(none)')"
InEnv python -c "import deeplabcut; print('  deeplabcut:', deeplabcut.__version__)"

Ok "Install complete."
Write-Host ""
if (Test-Path (Join-Path $Here "run.bat")) {
    Write-Host "Start the worker by double-clicking run.bat, or from this PowerShell:"
    Write-Host "    .\run.bat"
} else {
    Write-Host "Start the worker with:"
    Write-Host "    `$env:ETHOVISION_SERVER='http://<your-server>:3012'"
    Write-Host "    `$env:ETHOVISION_TOKEN='<your-token>'"
    Write-Host "    & '$conda' run -n $EnvName --no-capture-output python '$Here\ethovision_worker.py'"
}