param(
    [string]$SshTarget = "docker-gpu.cin.su",
    [int]$ConnectTimeoutSeconds = 5
)

$ErrorActionPreference = "Stop"

function Invoke-Remote {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Command
    )

    ssh `
        -o BatchMode=yes `
        -o ConnectTimeout=$ConnectTimeoutSeconds `
        $SshTarget `
        $Command
    if ($LASTEXITCODE -ne 0) {
        throw "Remote command failed on ${SshTarget} with exit code ${LASTEXITCODE}: ${Command}"
    }
}

Write-Host "Checking SSH access to $SshTarget..."
Invoke-Remote "hostname"

Write-Host ""
Write-Host "Checking NVIDIA GPU..."
Invoke-Remote "nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader"

Write-Host ""
Write-Host "Checking Docker..."
Invoke-Remote "docker version --format '{{.Server.Version}}'"

Write-Host ""
Write-Host "Checking Docker Compose..."
Invoke-Remote "docker compose version"

Write-Host ""
Write-Host "Checking Docker GPU runtime..."
Invoke-Remote "docker run --rm --gpus all nvidia/cuda:12.4.1-base-ubuntu22.04 nvidia-smi --query-gpu=name,memory.total --format=csv,noheader"

Write-Host ""
Write-Host "GPU host preflight completed."