param(
    [string]$NodeIp = (Get-Content "$env:RIEMANN_ROOT\cloud_state\node_ip.txt").Trim(),
    [string]$LocalRoot = $env:RIEMANN_ROOT,
    [string]$ModelRepo = ""
)
$ErrorActionPreference = "Continue"
$ok = $true
function Check($label, $cond) { $script:ok = $script:ok -and $cond; "{0} {1}" -f ($(if ($cond) { "PASS" } else { "FAIL" }), $label) }

Check "bootstrap reports ALL DONE on node" ((ssh "root@$NodeIp" "grep -c 'GATE: ALL DONE' /workspace/RiemannModel/logs/bootstrap.log") -match '^[1-9]')

foreach ($s in "stage2_538m", "stage3_1.19b", "stage4_2.65b", "stage5_6.66b", "stage6_12.85b") {
    Check "local checkpoint mirrored and verified: $s" (Test-Path "$LocalRoot\checkpoints\$s\.mirrored")
    Check "local controller present: $s" ((Test-Path "$LocalRoot\ntk_controllers\$s.pt") -and (Test-Path "$LocalRoot\ntk_controllers\$s.json"))
}
Check "local training logs present" ((Get-ChildItem "$LocalRoot\cloud_logs" -Filter "pretrain_stage6*.log" -ErrorAction SilentlyContinue).Count -ge 1)

if ($ModelRepo) {
    $remoteList = ssh "root@$NodeIp" "cd /workspace/RiemannModel/checkpoints/stage6_12.85b && ls -l model-*.safetensors | awk '{print `$5, `$9}'"
    $hub = python -c "
from huggingface_hub import HfApi
for f in HfApi().list_repo_tree('$ModelRepo', recursive=True):
    s=getattr(f,'size',None)
    if s and f.path.endswith('.safetensors'): print(s, f.path.split('/')[-1])
"
    $remoteMap = @{}; foreach ($l in $remoteList) { if ($l -match '^(\d+)\s+(\S+)$') { $remoteMap[$matches[2]] = [long]$matches[1] } }
    $hubMap = @{}; foreach ($l in $hub) { if ($l -match '^(\d+)\s+(\S+)$') { $hubMap[$matches[2]] = [long]$matches[1] } }
    $shardsMatch = ($remoteMap.Count -gt 0) -and ($remoteMap.Count -eq $hubMap.Count)
    foreach ($k in $remoteMap.Keys) { if ($hubMap[$k] -ne $remoteMap[$k]) { $shardsMatch = $false } }
    Check "Hub has all $($remoteMap.Count) stage6 shards with matching sizes" $shardsMatch
    Check "Hub has stage6 controller" ((python -c "from huggingface_hub import HfApi; print(any(f.path=='ntk_controller/stage6_12.85b.pt' for f in HfApi().list_repo_tree('$ModelRepo', recursive=True)))") -eq "True")
}

""
if ($ok) { "SAFE TO DESTROY: all archive checks passed." } else { "DO NOT DESTROY: one or more checks failed." }