param(
    [string]$NodeIp = (Get-Content "$env:RIEMANN_ROOT\cloud_state\node_ip.txt").Trim(),
    [string]$LocalRoot = $env:RIEMANN_ROOT,
    [string]$RemoteRoot = "/workspace/RiemannModel",
    [int]$PollSeconds = 300
)
$ErrorActionPreference = "Continue"
$node = "root@$NodeIp"
function Log($m) { "$(Get-Date -Format 'yyyy-MM-dd HH:mm:ss') [mirror] $m" }

New-Item -ItemType Directory -Force -Path "$LocalRoot\checkpoints", "$LocalRoot\ntk_controllers", "$LocalRoot\cloud_logs" | Out-Null

function Pull-Small {
    cmd /c "ssh $node ""tar -C $RemoteRoot -cf - ntk_controllers logs"" | tar -xf - -C ""$LocalRoot\_incoming""" 2>$null
    if (Test-Path "$LocalRoot\_incoming\ntk_controllers") { Copy-Item "$LocalRoot\_incoming\ntk_controllers\*" "$LocalRoot\ntk_controllers\" -Force }
    if (Test-Path "$LocalRoot\_incoming\logs") { Copy-Item "$LocalRoot\_incoming\logs\*" "$LocalRoot\cloud_logs\" -Force }
}

function Pull-Stage($stage) {
    Log "pulling $stage ..."
    $dest = "$LocalRoot\checkpoints"
    $t = [Diagnostics.Stopwatch]::StartNew()
    cmd /c "ssh $node ""tar -C $RemoteRoot/checkpoints -cf - $stage"" | tar -xf - -C ""$dest"""
    $t.Stop()
    $remote = ssh $node "cd $RemoteRoot/checkpoints/$stage && sha256sum *.safetensors *.json 2>/dev/null"
    $ok = $true
    foreach ($line in $remote) {
        if ($line -match '^([0-9a-f]{64})\s+\*?(.+)$') {
            $rsha = $matches[1]; $name = $matches[2].Trim()
            $local = Join-Path "$dest\$stage" $name
            if (-not (Test-Path $local)) { Log "MISSING $name"; $ok = $false; continue }
            $lsha = (Get-FileHash -Algorithm SHA256 $local).Hash.ToLower()
            if ($lsha -ne $rsha) { Log "SHA MISMATCH $name"; $ok = $false }
        }
    }
    $gb = [math]::Round((Get-ChildItem "$dest\$stage" -File | Measure-Object Length -Sum).Sum / 1GB, 2)
    if ($ok) {
        Set-Content "$dest\$stage\.mirrored" (Get-Date -Format o)
        Log "$stage mirrored and verified: $gb GB in $([math]::Round($t.Elapsed.TotalMinutes,1)) min"
    } else {
        Log "$stage verification FAILED; will retry next cycle"
        Remove-Item "$dest\$stage\.mirrored" -ErrorAction SilentlyContinue
    }
}

New-Item -ItemType Directory -Force -Path "$LocalRoot\_incoming" | Out-Null
Log "mirroring $node -> $LocalRoot every $PollSeconds s"
while ($true) {
    try {
        Pull-Small
        $done = ssh $node "ls $RemoteRoot/checkpoints/*/.rung_done 2>/dev/null"
        foreach ($m in $done) {
            if ($m -match 'checkpoints/([^/]+)/\.rung_done') {
                $stage = $matches[1]
                if (-not (Test-Path "$LocalRoot\checkpoints\$stage\.mirrored")) { Pull-Stage $stage }
            }
        }
        $allDone = ssh $node "grep -c 'GATE: ALL DONE' $RemoteRoot/logs/bootstrap.log 2>/dev/null"
        if ("$allDone" -match '^[1-9]') {
            Pull-Small
            Log "bootstrap reports ALL DONE and every finished rung is mirrored; exiting"
            break
        }
    } catch { Log "cycle error: $($_.Exception.Message)" }
    Start-Sleep $PollSeconds
}