# Windows Runner Startup Script
#
# This script runs when a new ephemeral Windows GPU VM boots. It:
# 1. Reads the JIT (just-in-time) runner config from GCP instance metadata
# 2. Configures the GitHub Actions runner with the JIT config
# 3. Runs the runner (executes one job, then exits)
# 4. Shuts down the VM (so the scaler can delete it)
#
# The JIT config is a base64-encoded blob generated by the Scale Set Client
# that contains everything the runner needs to register with GitHub.

$ErrorActionPreference = "Stop"

$runnerDir = "C:\actions-runner"
$logFile = "C:\actions-runner\startup.log"

function Write-Log {
    param([string]$Message)
    $timestamp = Get-Date -Format "yyyy-MM-dd HH:mm:ss"
    $line = "$timestamp - $Message"
    Write-Host $line
    Add-Content -Path $logFile -Value $line
}

Write-Log "=== Windows GPU Runner Startup ==="

# Step 1: Read JIT config from GCP instance metadata
Write-Log "Reading JIT config from instance metadata..."
$metadataUrl = "http://metadata.google.internal/computeMetadata/v1/instance/attributes/jit-config"
$maxRetries = 10
$jitConfig = $null

for ($i = 1; $i -le $maxRetries; $i++) {
    try {
        $jitConfig = Invoke-RestMethod -Uri $metadataUrl -Headers @{ "Metadata-Flavor" = "Google" } -TimeoutSec 10
        break
    }
    catch {
        Write-Log "  Attempt ${i}/${maxRetries}: Metadata not available yet, waiting..."
        Start-Sleep -Seconds 5
    }
}

if (-not $jitConfig) {
    Write-Log "ERROR: Failed to read JIT config from metadata after $maxRetries attempts"
    Stop-Computer -Force
    exit 1
}

Write-Log "JIT config retrieved (${($jitConfig.Length)} chars)"

# Step 2: Log GPU and system info
Write-Log "=== System Information ==="
try {
    nvidia-smi 2>&1 | ForEach-Object { Write-Log "  $_" }
}
catch {
    Write-Log "WARNING: nvidia-smi not available"
}

# Step 3: Start sccache server (for build caching)
Write-Log "Starting sccache server..."
try {
    sccache --start-server 2>&1 | ForEach-Object { Write-Log "  $_" }
}
catch {
    Write-Log "WARNING: sccache start failed (builds will be uncached)"
}

# Step 4: Configure and run the GitHub Actions runner with JIT config
Write-Log "Starting runner with JIT config..."
Set-Location $runnerDir

try {
    # The --jitconfig flag configures and runs the runner in one step.
    # In ephemeral mode, it runs exactly one job and then exits.
    & .\run.cmd --jitconfig $jitConfig
    $exitCode = $LASTEXITCODE
    Write-Log "Runner exited with code $exitCode"
}
catch {
    Write-Log "ERROR: Runner failed: $_"
    $exitCode = 1
}

# Step 5: Stop sccache and show stats
Write-Log "=== sccache stats ==="
try {
    sccache --show-stats 2>&1 | ForEach-Object { Write-Log "  $_" }
    sccache --stop-server 2>&1 | Out-Null
}
catch {
    Write-Log "WARNING: sccache stats unavailable"
}

# Step 6: Shut down the VM
# The scaler's cleanup loop will detect the TERMINATED state and delete the VM.
Write-Log "=== Runner complete, shutting down VM ==="
Stop-Computer -Force
