param(
    [string]$NodeIp = (Get-Content "$env:RIEMANN_ROOT\cloud_state\node_ip.txt").Trim(),
    [string]$Stage = "stage6_12.85b",
    [string]$LocalRoot = $env:RIEMANN_ROOT,
    [string]$RemoteRoot = "/workspace/RiemannModel"
)
$ErrorActionPreference = "Continue"
$node = "root@$NodeIp"
$dest = "$LocalRoot\checkpoints\$Stage"
New-Item -ItemType Directory -Force -Path $dest, "$LocalRoot\ntk_controllers", "$LocalRoot\cloud_logs" | Out-Null
function Log($m) { "$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss') [pull] $m" }

$manifest = ssh $node "cd $RemoteRoot/checkpoints/$Stage && for f in *; do [ -f `"`$f`" ] && echo `"`$(stat -c %s `"`$f`") `$(sha256sum `"`$f`" | cut -d' ' -f1) `$f`"; done"
$files = @()
foreach ($line in $manifest) {
    if ($line -match '^(\d+)\s+([0-9a-f]{64})\s+(.+)$') { $files += [pscustomobject]@{ Size=[long]$matches[1]; Sha=$matches[2]; Name=$matches[3].Trim() } }
}
Log "$($files.Count) files, $([math]::Round(($files | Measure-Object Size -Sum).Sum/1GB,1)) GB total"

foreach ($f in ($files | Sort-Object Name)) {
    $local = Join-Path $dest $f.Name
    $attempt = 0
    while ($true) {
        $attempt++
        if ((Test-Path $local) -and ((Get-Item $local).Length -eq $f.Size)) {
            $h = (Get-FileHash -Algorithm SHA256 $local).Hash.ToLower()
            if ($h -eq $f.Sha) { Log "OK   $($f.Name) ($([math]::Round($f.Size/1MB)) MB)"; break }
            Log "sha mismatch on $($f.Name), re-pulling"
        }
        if ($attempt -gt 4) { Log "FAILED $($f.Name) after $attempt attempts"; break }
        $t = [Diagnostics.Stopwatch]::StartNew()
        scp -q -o ServerAliveInterval=15 -o ServerAliveCountMax=4 "${node}:$RemoteRoot/checkpoints/$Stage/$($f.Name)" "$local" 2>$null
        $t.Stop()
        if ($LASTEXITCODE -ne 0) { Log "scp exit $LASTEXITCODE on $($f.Name) (attempt $attempt), retrying"; Start-Sleep 10 }
        else { Log "pulled $($f.Name) in $([math]::Round($t.Elapsed.TotalMinutes,1)) min" }
    }
}

$bad = @()
foreach ($f in $files) {
    $local = Join-Path $dest $f.Name
    if (-not (Test-Path $local) -or (Get-Item $local).Length -ne $f.Size -or (Get-FileHash -Algorithm SHA256 $local).Hash.ToLower() -ne $f.Sha) { $bad += $f.Name }
}
if ($bad.Count -eq 0) {
    Set-Content "$dest\.mirrored" (Get-Date -Format o)
    Log "$Stage fully mirrored and verified"
} else { Log "VERIFY FAILED: $($bad -join ', ')" }

scp -q "${node}:$RemoteRoot/ntk_controllers/*" "$LocalRoot\ntk_controllers\" 2>$null
scp -q "${node}:$RemoteRoot/logs/*.log" "$LocalRoot\cloud_logs\" 2>$null
Log "controllers and logs refreshed"