# atlas-train.ps1 — train a new Atlas model from R2 trajectories and upload it.
#
# Usage:
#   First-time setup:   .\atlas-train.ps1 -Setup
#   Train + upload:     .\atlas-train.ps1
#   Train, no upload:   .\atlas-train.ps1 -NoUpload
#
# Requires Python 3.11+ in PATH. Reads config from .\atlas-trainer.json.

param(
    [switch]$Setup,
    [switch]$NoUpload,
    [string]$ConfigPath = "atlas-trainer.json"
)

$ErrorActionPreference = "Stop"

# ─── Helpers ───────────────────────────────────────────────────────────────

function Write-Step($msg) {
    Write-Host ""
    Write-Host "═══ $msg ═══" -ForegroundColor Cyan
}

function Test-Command($cmd) {
    return [bool](Get-Command $cmd -ErrorAction SilentlyContinue)
}

# ─── Setup ─────────────────────────────────────────────────────────────────

if ($Setup) {
    Write-Step "First-time setup"

    if (-not (Test-Command "python")) {
        Write-Host "ERROR: Python is not in PATH. Install Python 3.11+ from python.org first." -ForegroundColor Red
        exit 1
    }
    $pyver = (python --version) -replace "Python ", ""
    Write-Host "Found Python $pyver"

    if (-not (Test-Path ".venv")) {
        Write-Host "Creating .venv..."
        python -m venv .venv
    }

    Write-Host "Installing PyTorch (CUDA 12.4)..."
    .\.venv\Scripts\python.exe -m pip install --upgrade pip
    .\.venv\Scripts\python.exe -m pip install torch --index-url https://download.pytorch.org/whl/cu124

    Write-Host "Installing other deps..."
    .\.venv\Scripts\python.exe -m pip install onnx onnxscript pyarrow numpy

    Write-Host ""
    Write-Host "✓ Setup complete." -ForegroundColor Green

    if (-not (Test-Path $ConfigPath)) {
        Write-Host ""
        Write-Host "Next: create $ConfigPath with your R2 + admin credentials. Sample:" -ForegroundColor Yellow
        Write-Host @"
{
  "apiUrl": "https://atlas-engine.landover.io",
  "adminKey": "ASK_JAY_FOR_THIS",
  "r2Endpoint": "https://...r2.cloudflarestorage.com",
  "r2AccessKeyId": "...",
  "r2SecretAccessKey": "...",
  "r2Bucket": "atlas-trajectories",
  "maxGames": 50000,
  "epochs": 20,
  "hidden": 128,
  "arch": "deep"
}
"@
    }
    Write-Host ""
    Write-Host "Then run:  .\atlas-train.ps1" -ForegroundColor Yellow
    exit 0
}

# ─── Main flow ─────────────────────────────────────────────────────────────

if (-not (Test-Path $ConfigPath)) {
    Write-Host "ERROR: $ConfigPath not found. Run '.\atlas-train.ps1 -Setup' first." -ForegroundColor Red
    exit 1
}
$cfg = Get-Content $ConfigPath | ConvertFrom-Json

if (-not (Test-Path ".\atlas-cli.exe")) {
    Write-Host "ERROR: atlas-cli.exe not found in current directory." -ForegroundColor Red
    Write-Host "Download from https://downloads.landover.io/atlas-cli-win-x64.zip" -ForegroundColor Yellow
    exit 1
}
if (-not (Test-Path ".venv")) {
    Write-Host "ERROR: .venv not found. Run '.\atlas-train.ps1 -Setup' first." -ForegroundColor Red
    exit 1
}

$timestamp = (Get-Date).ToUniversalTime().ToString("yyyyMMddTHHmmssZ")
$parquetPath = "data\$timestamp.parquet"
$onnxPath = "models\$timestamp.onnx"
$version = "v-$timestamp"

New-Item -ItemType Directory -Force -Path data, models | Out-Null

# Step 1: Pull trajectories from R2 and convert to Parquet.
Write-Step "Step 1/3 — Replay R2 trajectories → Parquet"
$env:R2_ENDPOINT          = $cfg.r2Endpoint
$env:R2_ACCESS_KEY_ID     = $cfg.r2AccessKeyId
$env:R2_SECRET_ACCESS_KEY = $cfg.r2SecretAccessKey
$env:R2_BUCKET            = $cfg.r2Bucket
& .\atlas-cli.exe replay-r2 trajectories/ $parquetPath --max-games $cfg.maxGames --game-type 2p
if ($LASTEXITCODE -ne 0) { throw "replay-r2 failed (exit $LASTEXITCODE)" }
$parquetSize = (Get-Item $parquetPath).Length / 1MB
Write-Host ("✓ Parquet: {0:N1} MB" -f $parquetSize) -ForegroundColor Green

# Step 2: Train.
Write-Step "Step 2/3 — Train neural model"
$trainArgs = @(
    "atlas-ml\train.py",
    "--data", $parquetPath,
    "--gametype", "2p",
    "--out", $onnxPath,
    "--epochs", $cfg.epochs,
    "--hidden", $cfg.hidden,
    "--arch", $cfg.arch,
    "--policy-weight", "1.0",
    "--dropout", "0.1"
)
& .\.venv\Scripts\python.exe @trainArgs
if ($LASTEXITCODE -ne 0) { throw "training failed (exit $LASTEXITCODE)" }
$modelSize = (Get-Item $onnxPath).Length / 1KB
Write-Host ("✓ Model: {0:N1} KB at {1}" -f $modelSize, $onnxPath) -ForegroundColor Green

# Step 3: Upload.
if ($NoUpload) {
    Write-Step "Step 3/3 — Skipped (--NoUpload)"
    Write-Host "Model saved locally at $onnxPath"
    exit 0
}

Write-Step "Step 3/3 — Upload model to API"
$resp = Invoke-RestMethod -Uri "$($cfg.apiUrl)/api/model/upload?version=$version" `
                           -Method POST `
                           -Headers @{ Authorization = "Bearer $($cfg.adminKey)" } `
                           -ContentType "application/octet-stream" `
                           -InFile $onnxPath
Write-Host ("✓ Upload: {0}" -f ($resp | ConvertTo-Json -Compress)) -ForegroundColor Green

Write-Step "Done!"
Write-Host "Model $version is live. Workers will pick it up on their next heartbeat (~5 min)."
Write-Host "Watch atlas.landover.io for the Model field to flip to '$version'."
