param(
    [string]$SshTarget = "",
    [string]$ModelsRoot = "Z:\LLM\models",
    [int]$ConnectTimeoutSeconds = 5
)

$ErrorActionPreference = "Stop"

function Invoke-LocalCheck {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Name,
        [Parameter(Mandatory = $true)]
        [scriptblock]$Block
    )

    Write-Host ""
    Write-Host "== $Name"
    & $Block
}

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

    Write-Host ""
    Write-Host "== $Name"
    ssh -o BatchMode=yes -o ConnectTimeout=$ConnectTimeoutSeconds $SshTarget "powershell -NoProfile -Command $Command"
    if ($LASTEXITCODE -ne 0) {
        throw "Remote check failed on ${SshTarget}: ${Name}"
    }
}

if ($SshTarget) {
    Invoke-RemoteCheck "Host identity" "whoami; hostname; `$PSVersionTable.PSVersion.ToString()"
    Invoke-RemoteCheck "GPU driver" "nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader"
    Invoke-RemoteCheck "Docker version" "docker version"
    Invoke-RemoteCheck "Docker compose" "docker compose version"
    Invoke-RemoteCheck "Docker info" "docker info --format '{{json .}}'"
    Invoke-RemoteCheck "Models drive" "if (-not (Test-Path '$ModelsRoot')) { New-Item -ItemType Directory -Force '$ModelsRoot' | Out-Null }; Get-Item '$ModelsRoot' | Select-Object FullName,Exists"
    Invoke-RemoteCheck "CUDA container" "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 "Windows GPU host preflight completed through SSH."
    exit 0
}

Invoke-LocalCheck "Host identity" {
    whoami
    hostname
    $PSVersionTable.PSVersion.ToString()
}

Invoke-LocalCheck "GPU driver" {
    nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
}

Invoke-LocalCheck "Docker version" {
    docker version
}

Invoke-LocalCheck "Docker compose" {
    docker compose version
}

Invoke-LocalCheck "Docker info" {
    docker info --format '{{json .}}'
}

Invoke-LocalCheck "Models drive" {
    if (-not (Test-Path $ModelsRoot)) {
        New-Item -ItemType Directory -Force $ModelsRoot | Out-Null
    }
    Get-Item $ModelsRoot | Select-Object FullName,Exists
}

Invoke-LocalCheck "CUDA container" {
    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 "Windows GPU host preflight completed."