# FaceLoRA v3 training monitor (scheduled task, fully autonomous)
# ============================================================
# 用途：云端训练全自动监控。Windows 计划任务每 5 分钟调用一次。
# 功能：
#   1. 增量下载 checkpoint（HTTP HEAD 探测 + aria2c）——任何时刻断线已下载的都是安全的
#   2. 完成判定：e120 + final 都下载 -> 删 pod 停计费
#   3. 崩溃检测：最后 checkpoint 45min 无更新 -> 抢救 + 删 pod
#   4. 超时止损：超过预算小时数 -> 删 pod
#   5. 状态持久化到 monitor_state.txt（纯文本，PS5.1 下 .json 写入不可靠）
# 输出（全在 ASCII 目录 C:\Users\hmo\AppData\Local\Temp\opencode\face_lora\）：
#   monitor_v3.log    - 详细运行日志
#   train_events.log  - 关键事件时间线（下载/完成/止损/删pod）
#   monitor_state.txt - 持久化状态（pod_start/downloaded/done/started）
#   checkpoints_v3/   - 已下载的 checkpoint
#
# !! PS 5.1 铁律（血泪教训，2026-08-09）：
#   - 变量名不区分大小写！$STATE(路径) 会被 $state(hashtable) 覆盖 → 路径变量必须用独特名（$STATEFILE）
#   - 中文路径下 Add-Content 写新 .json 文件不可靠 → 输出目录用纯 ASCII，状态用纯文本
#   - train_env.json 读取失败必须退出，绝不用空值执行删 pod（曾因此差点误删）
# ============================================================
$ErrorActionPreference = "Continue"

# ---- 项目路径（注意：脚本必须存 GBK 编码，PowerShell 5.1 用 ANSI 读）----
$PROJ = "D:\F\NewI\opencode\daily-workspace\projects\脸部LoRA训练-Qwen-Image"
# 输出目录用纯 ASCII（避开中文路径 + PS5.1 编码 bug）
$OUTDIR = "C:\Users\hmo\AppData\Local\Temp\opencode\face_lora"

# ---- 读取训练环境配置（train_env.json：pod_id/dl_url/ssh/max_hours/total_steps）----
$CFG = Join-Path $PROJ "temp\train_env.json"
if (-not (Test-Path $CFG)) {
    Write-Host "missing train_env.json"
    exit 1
}
$env_cfg = Get-Content $CFG -Raw | ConvertFrom-Json
$PODID = $env_cfg.pod_id
$DLURL = $env_cfg.dl_url
$SSHHOST = $env_cfg.ssh_host
$SSHPORT = $env_cfg.ssh_port
$MAX_HOURS = [double]$env_cfg.max_hours
$TOTAL_STEPS = [int]$env_cfg.total_steps
# 安全阀：任何关键字段为空 → 立即退出，绝不用空值执行删 pod 等破坏性操作
if (-not $PODID -or -not $DLURL -or $MAX_HOURS -le 0) {
    Write-Host ("train_env invalid: pod=" + $PODID + " max_hours=" + $MAX_HOURS)
    exit 1
}

$MLOG = Join-Path $OUTDIR "monitor_v3.log"
$EVENT = Join-Path $OUTDIR "train_events.log"
$CDIR = Join-Path $OUTDIR "checkpoints_v3"
$STATEFILE = Join-Path $OUTDIR "monitor_state.txt"
$RPKEY = (Get-Content (Join-Path $PROJ ".runpod_api_key") -Raw).Trim()
New-Item -ItemType Directory -Force $CDIR, (Split-Path $MLOG) | Out-Null

# ---- 单例守卫（防计划任务重复触发）----
$LOCK = Join-Path $OUTDIR "monitor_v3.lock"
if (Test-Path $LOCK) {
    $oldPid = [int](Get-Content $LOCK)
    if (Get-Process -Id $oldPid -ErrorAction SilentlyContinue) {
        exit 0
    }
}
Set-Content $LOCK $PID

function Log($msg) {
    $ts = Get-Date -Format "MM-dd HH:mm:ss"
    Add-Content -Path $MLOG -Value ("[" + $ts + "] " + $msg) -Encoding UTF8
}
function Event($msg) {
    $ts = Get-Date -Format "yyyy-MM-dd HH:mm:ss"
    Add-Content -Path $EVENT -Value ("[" + $ts + "] " + $msg) -Encoding UTF8
}

try {
    # ---- 读持久化状态（纯文本）----
    $state = @{ downloaded = @(); pod_start = ""; done = $false; started = $false }
    if (Test-Path $STATEFILE) {
        try {
            foreach ($ln in Get-Content $STATEFILE) {
                if ($ln -like "downloaded=*") { $state.downloaded = @($ln.Substring(11).Split(";") | Where-Object { $_ }) }
                elseif ($ln -like "pod_start=*") { $state.pod_start = $ln.Substring(10) }
                elseif ($ln -eq "done=true") { $state.done = $true }
                elseif ($ln -eq "started=true") { $state.started = $true }
            }
        } catch {}
    }
    if (-not $state.pod_start) {
        $state.pod_start = Get-Date -Format "yyyy-MM-dd HH:mm:ss"
        Event ("monitor start: pod=" + $PODID + " dl_url=" + $DLURL + " max_hours=" + $MAX_HOURS)
    }

    if ($state.done) {
        Log "done, skip"
        exit 0
    }

    # ---- 0. 探测训练是否已开始（train.log 有 steps）----
    if (-not $state.started) {
        try {
            $tail = & ssh -i "$env:USERPROFILE\.ssh\id_rsa" -p $SSHPORT -o BatchMode=yes -o ConnectTimeout=10 $SSHHOST "tail -c 2000 /workspace/train.log 2>/dev/null | tr '\r' '\n' | tail -2" 2>$null
            if ($tail -match ("(\d+)/" + $TOTAL_STEPS)) {
                $state.started = $true
                Event ("training started: step " + $Matches[1] + "/" + $TOTAL_STEPS)
            }
        } catch {}
    }

    # ---- 1. 探测云端 checkpoint（HTTP HEAD）----
    $epochs = @(10,20,30,40,50,60,70,80,90,100,110,120)
    $newFiles = @()
    foreach ($e in $epochs) {
        $fn = "myface_lora-{0:D6}.safetensors" -f $e
        if ($fn -in $state.downloaded) { continue }
        try {
            $resp = Invoke-WebRequest -Uri ($DLURL + "/" + $fn) -Method Head -TimeoutSec 15 -UseBasicParsing
            if ($resp.StatusCode -eq 200) { $newFiles += $fn }
        } catch {}
    }
    try {
        $resp = Invoke-WebRequest -Uri ($DLURL + "/myface_lora.safetensors") -Method Head -TimeoutSec 15 -UseBasicParsing
        if ($resp.StatusCode -eq 200) { $newFiles += "myface_lora.safetensors" }
    } catch {}

    # ---- 2. 增量下载 ----
    foreach ($f in $newFiles) {
        $fp = Join-Path $CDIR $f
        if (Test-Path $fp) {
            $sz = [math]::Round((Get-Item $fp).Length/1MB,0)
            if ($sz -gt 100) {
                Log ("already have: " + $f + " (" + $sz + " MB), skip")
                $state.downloaded += $f
                continue
            }
            Remove-Item $fp -Force -ErrorAction SilentlyContinue
        }
        Log ("new checkpoint: " + $f + " downloading")
        & aria2c -x16 -s16 -k1M --continue -d "$CDIR" -o $f ($DLURL + "/" + $f) 2>&1 | Out-Null
        if (Test-Path $fp) {
            $sz = [math]::Round((Get-Item $fp).Length/1MB,0)
            if ($sz -gt 100) {
                Log ("download OK: " + $f + " (" + $sz + " MB)")
                Event ("CHECKPOINT DOWNLOADED: " + $f + " (" + $sz + " MB)")
                $state.downloaded += $f
            } else {
                Remove-Item $fp -Force -ErrorAction SilentlyContinue
                Log ("incomplete (" + $sz + " MB), retry next round: " + $f)
            }
        } else {
            Log ("download FAIL: " + $f)
        }
        break
    }

    # ---- 3. 完成判定 ----
    $hasFinal = "myface_lora.safetensors" -in $state.downloaded
    if ($hasFinal) {
        Log "=== TRAINING COMPLETE: final checkpoint downloaded ==="
        Event ("=== TRAINING COMPLETE: all checkpoints in checkpoints_v3\ ===")
        try {
            Invoke-RestMethod -Uri ("https://rest.runpod.io/v1/pods/" + $PODID) -Method Delete -Headers @{Authorization = "Bearer $RPKEY"} -TimeoutSec 30 | Out-Null
            Log ("pod " + $PODID + " deleted (complete)")
            Event ("POD DELETED: " + $PODID + " (complete, billing stopped)")
        } catch { Log ("pod delete fail: " + $_.Exception.Message) }
        $state.done = $true
    }

    # ---- 4. 崩溃检测 ----
    if ($state.downloaded.Count -gt 0 -and -not $state.done) {
        $lastFile = $state.downloaded[-1]
        $lastPath = Join-Path $CDIR $lastFile
        if (Test-Path $lastPath) {
            $ageMin = [int]((Get-Date) - (Get-Item $lastPath).LastWriteTime).TotalMinutes
            if ($ageMin -gt 45) {
                Log ("CRASH: last ckpt (" + $lastFile + ") " + $ageMin + "min stale -> rescue + delete pod")
                Event ("CRASH: last ckpt " + $lastFile + " " + $ageMin + "min stale -> rescue + delete pod")
                foreach ($e in $epochs) {
                    $fn = "myface_lora-{0:D6}.safetensors" -f $e
                    if ($fn -in $state.downloaded) { continue }
                    try {
                        $r = Invoke-WebRequest -Uri ($DLURL + "/" + $fn) -Method Head -TimeoutSec 15 -UseBasicParsing
                        if ($r.StatusCode -eq 200) {
                            & aria2c -x16 -s16 -k1M -d "$CDIR" -o $fn ($DLURL + "/" + $fn) 2>&1 | Out-Null
                            if (Test-Path (Join-Path $CDIR $fn)) { $state.downloaded += $fn; Event ("RESCUED: " + $fn) }
                        }
                    } catch {}
                }
                try {
                    Invoke-RestMethod -Uri ("https://rest.runpod.io/v1/pods/" + $PODID) -Method Delete -Headers @{Authorization = "Bearer $RPKEY"} -TimeoutSec 30 | Out-Null
                    Log ("pod " + $PODID + " deleted (crash)")
                    Event ("POD DELETED: " + $PODID + " (crash stop)")
                } catch { Log ("pod delete fail: " + $_.Exception.Message) }
                $state.done = $true
            }
        }
    }

    # ---- 5. 超时止损 ----
    if (-not $state.done) {
        try {
            $start = [datetime]::ParseExact($state.pod_start, "yyyy-MM-dd HH:mm:ss", $null)
            $runHours = ((Get-Date) - $start).TotalHours
            if ($runHours -gt $MAX_HOURS) {
                Log ("TIMEOUT: " + [math]::Round($runHours,1) + "h > budget " + $MAX_HOURS + "h")
                Event ("TIMEOUT: " + [math]::Round($runHours,1) + "h > budget " + $MAX_HOURS + "h -> delete pod")
                try {
                    Invoke-RestMethod -Uri ("https://rest.runpod.io/v1/pods/" + $PODID) -Method Delete -Headers @{Authorization = "Bearer $RPKEY"} -TimeoutSec 30 | Out-Null
                    Log ("pod " + $PODID + " deleted (timeout, downloaded " + $state.downloaded.Count + ")")
                    Event ("POD DELETED: " + $PODID + " (timeout, downloaded " + $state.downloaded.Count + ")")
                } catch { Log ("pod delete fail: " + $_.Exception.Message) }
                $state.done = $true
            } else {
                Log ("running: " + [math]::Round($runHours,1) + "/" + $MAX_HOURS + "h, downloaded " + $state.downloaded.Count)
            }
        } catch {
            Log ("time parse fail: " + $_.Exception.Message)
        }
    }

    # ---- 持久化状态（纯文本行，PS5.1 可靠）----
    try {
        $lines = @()
        $lines += "pod_start=" + $state.pod_start
        $lines += "downloaded=" + ($state.downloaded -join ";")
        $lines += ("done=" + $state.done)
        $lines += ("started=" + $state.started)
        if (Test-Path $STATEFILE) { Remove-Item $STATEFILE -Force }
        Add-Content -Path $STATEFILE -Value $lines -Encoding UTF8
    } catch {
        Log ("state write fail: " + $_.Exception.Message)
    }
    $state.downloaded | ForEach-Object { Event ("state: downloaded " + $_) }
    if ($state.done) { Event ("=== MONITOR DONE: " + $state.downloaded.Count + " checkpoints ===") }
    else { Event ("state: running, downloaded=" + $state.downloaded.Count) }
    if ($state.done) { Log "=== monitor done ===" }
} finally {
    Remove-Item $LOCK -Force -ErrorAction SilentlyContinue
}