#!/bin/bash
# lab_start.sh — Runpod の Pod の準備から mechArm 270 の実機推論までを1行で行う (依頼 37 / 手順は LOG407「起動の手順」)
#
#   ./lab_start.sh --gpu A40 --episode run_0001_20261019_194818 --log Log17
#   ./lab_start.sh --log Log17 --dry-run      アームなしで試す (デーモン・HOME・爪・位置合わせを飛ばし,
#                                             run_policy_client.py を --execute なし・--dry-keys で呼ぶ)
#
#   --gpu      GPU の種類 (既定 A40). 「A40」は「NVIDIA A40」として API に渡す. 名前に空白を含む種類は
#              「--gpu "NVIDIA RTX A6000"」のように Runpod の名前のまま渡す
#   --episode  位置合わせ (camera_align.py) で重ねるお手本 (既定 run_0001_20261019_194818)
#   --log      ターミナルの出力を残す /Volumes/KIOXIA/r2r2r-data/lab/<Log>/ (なければ聞く)
#
# 1. Pod を作る (A40×1・Runpod Pytorch 2.8.0・コンテナのディスク 200GB・22/tcp を外へ出す)
#    「resources」で借りられなければ 30 秒待って, 合わせて 5 回まで作り直す
#    <Log>/pod/lab_pod.json の Pod が API で RUNNING なら, 作らずにつなぎ直す (依頼 44)
#      - 3. で Pod の中を見る. サーバーが 8000 番で待っていれば 3b〜4 を飛ばす. 落ちていれば前の serve.log を
#        持ち帰り, モデル・openpi・.venv のないものだけやり直し, 3f と 4 を流す. 3a (スクリプトを送る) は毎回
#        プロセスはあるのに 8000 番で待っていなければ 30 秒待ち, それでも待っていなければ起動し直す (依頼 47)
#      - 止まっている (EXITED など) なら lab_stop.sh --skip-log で消してから作る. API にない・TERMINATED なら作る
#        どちらも古い lab_pod.json は lab_pod_<時刻>.json に名前を変えて残す
#      - API で様子が分からないとき・記録にない p2m-<Log> の Pod が残っているときは, 二重に借りないよう止まる
# 2. Pod が立ち上がり, SSH の IP とポートが API から出るまで待つ (画面・ログには出さない)
# 3. Pod を準備する (3b〜3f のコマンドは lab_lib.sh. lab_train.sh と同じもの) (pi05/ と add_trim_config.py → huggingface_hub → hf download → openpi 215abfb → uv sync
#    → addon・設定・assets → serve_policy.py を tmux の中で起動し /workspace/logs/serve.log に残す)
# 4. serve.log に「server listening on 0.0.0.0:8000」が出るまで待つ
# 5. Mac の 8000 番へのトンネルを裏で張る (このスクリプトが終わるときに閉じる)
# 6. tools/daemon_ctl.sh status で 5555 の LISTEN を見る. なければ start
# 7. HOME → 爪を開く → camera_align.py (q で閉じる)
# 8. run_policy_client.py (LOG407 と同じ引数). 終わったら, もう1回やるかを聞く (y なら 7 へ)
# 9. Pod は消さない. 片付けは ./lab_stop.sh
#
# Pod の情報 (id・GPU・1時間の料金・作った時刻) は <Log>/pod/lab_pod.json に残す (IP とポートは残さない).

set -u
set -o pipefail

# 中身は全部 main の中に置き, 最後の1行で呼ぶ (依頼 47). bash はスクリプトを少しずつ読みながら流すので,
# 動いている途中でこのファイルが書き換わると, 続きを新しい中身の途中から読んで壊れる. 関数は呼ぶ前に最後まで読む
main() {
HERE="$(cd "$(dirname "$0")" && pwd)"
# shellcheck source=lab_lib.sh
. "$HERE/lab_lib.sh"

ORIG_ARGS=("$@")
GPU=A40
EPISODE=run_0001_20261019_194818
LOG=""
DRY_RUN=0
while [ $# -gt 0 ]; do
    case "$1" in
        --gpu) GPU="$2"; shift 2 ;;
        --episode) EPISODE="$2"; shift 2 ;;
        --log) LOG="$2"; shift 2 ;;
        --dry-run) DRY_RUN=1; shift ;;
        -h|--help) sed -n '2,32p' "$0"; exit 0 ;;
        *) echo "[ERROR] 知らない引数: $1  (./lab_start.sh --help)" >&2; exit 2 ;;
    esac
done
case "$GPU" in *" "*) GPU_ID="$GPU" ;; *) GPU_ID="NVIDIA $GPU" ;; esac

if [ ! -d /Volumes/KIOXIA/r2r2r-data ]; then
    echo "[ERROR] SSD 未マウント: /Volumes/KIOXIA/r2r2r-data" >&2; exit 1
fi
if [ -z "$LOG" ]; then
    if [ -t 0 ]; then read -r -p "Log 番号 (例 Log17) > " LOG; fi
    [ -n "$LOG" ] || { echo "[ERROR] --log Log17 のように Log 番号を付けてください" >&2; exit 2; }
fi
LOGDIR="$LAB_ROOT/$LOG"
STATE="$LOGDIR/pod/lab_pod.json"

# ---- ターミナルの出力を画面に出しつつ <Log>/ に残す ----
if [ -z "${LAB_LOGGED:-}" ]; then
    mkdir -p "$LOGDIR/pod"
    LOGFILE="$LOGDIR/lab_start_$(date +%Y%m%d_%H%M%S).txt"
    echo "[LOG   ] ターミナルの出力を残します → $LOGFILE"
    if [ -t 0 ] && [ -t 1 ]; then
        # script は端末のまま記録する (camera_align.py・run_policy_client.py の入力待ちもそのまま使える)
        exec env LAB_LOGGED=1 script -q -a "$LOGFILE" "$0" "${ORIG_ARGS[@]}" --log "$LOG"
    fi
    export LAB_LOGGED=1
    exec > >(tee -a "$LOGFILE") 2>&1
fi
export PYTHONUNBUFFERED=1

lab_need_key "$RUNPOD_KEY_FILE" "Runpod の API キー"
lab_need_key "$HF_TOKEN_FILE" "Hugging Face の鍵"
[ -s "$POD_SSH_KEY" ] || { echo "[ERROR] Pod に入る SSH の鍵がありません: $POD_SSH_KEY" >&2; exit 1; }
[ -s "$POD_SSH_KEY.pub" ] || { echo "[ERROR] SSH の公開鍵がありません: $POD_SSH_KEY.pub" >&2; exit 1; }
ADD_TRIM=/Volumes/KIOXIA/r2r2r-data/lab/Log11/pod/add_trim_config.py
[ -f "$ADD_TRIM" ] || { echo "[ERROR] add_trim_config.py がありません: $ADD_TRIM" >&2; exit 1; }
RAW=/Volumes/KIOXIA/r2r2r-data/datasets/mecharm_redblock_box_teach_v1_trim2/raw_episodes
[ -d "$RAW/$EPISODE" ] || { echo "[ERROR] お手本がありません: $RAW/$EPISODE" >&2; exit 1; }
lab_ssh_opts

HF_REPO=PixelstoMotors/pi05-mecharm-trim-lora-v1
OPENPI_COMMIT=$LAB_OPENPI_COMMIT                                          # lab_lib.sh (lab_train.sh と同じ)
CKPT=/workspace/ckpt/3999                                                  # 3c で hf download が置く場所
NORM_STATS=assets/local/mecharm_redblock_box_teach_v1_trim2/norm_stats.json  # 3f が写す assets (CKPT からの相対)
LISTEN_WAIT=30; LISTEN_POLL=5    # プロセスはあるのに 8000 番で待っていないとき, 何秒待つか・何秒おきに見るか
TEMPLATE_ID=$LAB_TEMPLATE_ID
IMAGE=$LAB_IMAGE

# ---- 段階ごとの秒数 ----
T_NAMES=(); T_SECS=()
CUR_STAGE="(始める前)"
stage() { CUR_STAGE="$1"; STAGE_T0=$SECONDS; echo; echo "=== $1"; }
stage_done() { T_NAMES+=("$CUR_STAGE"); T_SECS+=($((SECONDS - STAGE_T0))); echo "    ($((SECONDS - STAGE_T0)) 秒)"; }
summary() {
    echo
    echo "=== かかった時間$([ "${REATTACH:-0}" = 1 ] && echo " ($LOG の Pod (id $POD_ID) につなぎ直し)")"
    local i
    for i in "${!T_NAMES[@]}"; do printf '  %-32s %6d 秒\n' "${T_NAMES[$i]}" "${T_SECS[$i]}"; done
}
POD_ID=""
fail() {
    echo
    echo "[STOP  ] 「${CUR_STAGE}」で止まりました: $*"
    if [ -n "$POD_ID" ]; then
        echo "[STOP  ] Pod は残っています (id $POD_ID). 消すなら ./lab_stop.sh --log $LOG"
    fi
    summary
    exit 1
}

TUNNEL_PID=""
LOCK="$LOGDIR/pod/.lab_start.lock"
LOCKED=0
cleanup() {
    if [ $LOCKED = 1 ] && [ "$(cat "$LOCK/pid" 2>/dev/null)" = "$$" ]; then rm -rf "$LOCK"; fi
    if [ -n "$TUNNEL_PID" ] && kill -0 "$TUNNEL_PID" 2>/dev/null; then
        kill "$TUNNEL_PID" 2>/dev/null
        echo "[TUNNEL] 8000 番のトンネルを閉じました"
    fi
}
trap cleanup EXIT
trap 'fail "Ctrl-C で止めました"' INT

# ---- 同じ Log の lab_start.sh を2つ同時に動かさない (mkdir は1つだけが成り立つ. 終わるときに外す) ----
# 印が残っていれば, 動いていてもいなくても止まる (自動で取り直すと, 2つが同時に取り直したときに両方進めるため)
if ! mkdir "$LOCK" 2>/dev/null; then
    other=$(cat "$LOCK/pid" 2>/dev/null)
    if [ -n "$other" ] && kill -0 "$other" 2>/dev/null; then
        echo "[ERROR] $LOG の lab_start.sh がもう動いています (pid $other). 二重に借りないよう止まります"
    else
        echo "[ERROR] $LOG の lab_start.sh の印が残っています (pid ${other:-?}・動いていないように見える). 二重に借りないよう止まります"
        echo "        ほかに $LOG の lab_start.sh が動いていないことを確かめてから, 印を消して打ち直してください:"
        echo "        rm -r $LOCK"
    fi
    exit 1
fi
LOCKED=1
echo $$ > "$LOCK/pid" || { echo "[ERROR] $LOCK/pid を書けない"; exit 1; }

echo "$LOG" > "$LAB_CURRENT"
echo "[LAB   ] $(date '+%Y-%m-%d %H:%M:%S')  GPU=$GPU_ID  episode=$EPISODE  log=$LOGDIR  $([ $DRY_RUN = 1 ] && echo 'ドライラン (アームなし)')"

# ---- 1. Pod を作る (同じ Log の Pod が動いていれば, 作らずにつなぎ直す) ----
REATTACH=0
stage "1. Pod を作る"
if [ -f "$STATE" ]; then
    OLD_ID=$(jq -r '.id // ""' "$STATE" 2>/dev/null)
    if [ -z "$OLD_ID" ]; then
        lab_archive_state "$STATE" || fail "読めない $STATE を残せない"
        echo "[POD   ] $STATE に Pod の id がない. $LAB_ARCHIVED に名前を変えて残し, 新しく作ります"
    elif ! lab_pod_kind "$OLD_ID"; then
        fail "記録の Pod (id $OLD_ID) の様子を API で確かめられない (HTTP ${RUNPOD_HTTP:-?}). 二重に借りないよう止まります"
    else
        case "$LAB_POD_KIND" in
            running)
                REATTACH=1 ;;
            stopped)
                echo "[POD   ] $LOG の Pod (id $OLD_ID) は止まっています (状態 $LAB_POD_STATUS). 止まっていてもディスクの料金がかかるので,"
                echo "         ./lab_stop.sh と同じ消し方で消してから新しく作ります"
                "$HERE/lab_stop.sh" --log "$LOG" --skip-log || fail "止まっている Pod (id $OLD_ID) を消せない (上の lab_stop.sh の出力を見る)"
                lab_archive_state "$STATE" || fail "古い $STATE を残せない"
                echo "[POD   ] 古い記録は $LAB_ARCHIVED に残しました" ;;
            gone)
                lab_archive_state "$STATE" || fail "古い $STATE を残せない"
                echo "[POD   ] 記録の Pod (id $OLD_ID) は API で見ると残っていない${LAB_POD_STATUS:+ (状態 $LAB_POD_STATUS)}. 新しく作ります"
                echo "         古い記録は $LAB_ARCHIVED に残しました" ;;
        esac
    fi
fi
if [ $REATTACH = 1 ]; then
    POD_ID=$OLD_ID
    COST_PER_HR=$(jq -r '.costPerHr // 0' "$STATE")
    POD_T0=$(jq -r '.created_epoch // 0' "$STATE")
    OLD_GPU=$(jq -r '.gpu // ""' "$STATE")
    echo "[POD   ] $LOG の Pod (id $POD_ID) につなぎ直します (状態 RUNNING  GPU=$OLD_GPU  1時間 $COST_PER_HR ドル  作った時刻 $(jq -r '.created_at // "?"' "$STATE"))"
    [ "$OLD_GPU" = "$GPU_ID" ] || echo "[POD   ] --gpu は $GPU_ID ですが, 残っている Pod の $OLD_GPU をそのまま使います"
    CUR_STAGE="1. Pod につなぎ直す"
else
    others=$(runpod_api GET /pods) || fail "Pod の一覧を API から取れない (HTTP $RUNPOD_HTTP)"
    lab_valid_pods "$others" || fail "Pod の一覧の答えが読めない. 二重に借りないよう止まります"
    same=$(jq -r --arg name "p2m-$LOG" '.[] | select(.name == $name and .desiredStatus != "TERMINATED") | .id' <<<"$others")
    if [ -n "$same" ]; then
        fail "記録にない $LOG の Pod (p2m-$LOG, id $(echo $same)) が残っている. 二重に借りないよう止まります (Runpod の画面か ./lab_stop.sh --check で確かめる)"
    fi
    n_others=$(jq 'length' <<<"$others")
    if [ "$n_others" -gt 0 ]; then
        echo "[WARN  ] ほかに Pod が $n_others 個あります (消し忘れなら ./lab_stop.sh --check で確かめる):"
        jq -r '.[] | "         \(.id)  \(.name)  \(.desiredStatus)  \(.costPerHr) ドル/時"' <<<"$others"
    fi
    body=$(jq -n --arg name "p2m-$LOG" --arg gpu "$GPU_ID" --arg tpl "$TEMPLATE_ID" --arg img "$IMAGE" \
              --arg pk "$(cat "$POD_SSH_KEY.pub")" \
              '{name:$name, templateId:$tpl, imageName:$img, gpuTypeIds:[$gpu], gpuCount:1,
                cloudType:"SECURE", containerDiskInGb:200, volumeInGb:0, ports:["22/tcp"],
                env:{PUBLIC_KEY:$pk}}')
    for attempt in 1 2 3 4 5; do
        resp=$(runpod_api POST /pods "$body") && break
        msg=$(jq -r '.error // .message // .' <<<"$resp" 2>/dev/null || echo "$resp")
        echo "[POD   ] 作れなかった ($attempt/5, HTTP $RUNPOD_HTTP): $msg"
        case "$msg" in
            *resources*|*"no longer any instances"*|*"not available"*|*"no instances"*)
                [ "$attempt" -lt 5 ] && { echo "[POD   ] 30 秒待って作り直します"; sleep 30; continue; } ;;
        esac
        fail "Pod を作れない: $msg"
    done
    POD_ID=$(jq -r '.id // ""' <<<"$resp")
    [ -n "$POD_ID" ] || fail "API の答えに Pod の id がない"
    COST_PER_HR=$(jq -r '.costPerHr // 0' <<<"$resp")
    POD_T0=$(date +%s)
    jq -n --arg id "$POD_ID" --arg gpu "$GPU_ID" --arg log "$LOG" --argjson cost "$COST_PER_HR" --argjson t0 "$POD_T0" \
          '{id:$id, gpu:$gpu, log:$log, costPerHr:$cost, created_epoch:$t0,
            created_at:($t0|strflocaltime("%Y-%m-%d %H:%M:%S"))}' > "$STATE"
    echo "[POD   ] 作った: id $POD_ID  GPU=$GPU_ID  1時間 $COST_PER_HR ドル  (記録 $STATE)"
fi
stage_done

# ---- 2. Pod が立ち上がり, SSH の IP とポートが出るまで待つ ----
stage "2. Pod が立ち上がるまで待つ"
for _ in $(seq 1 90); do
    lab_pod_addr && break
    sleep 10
done
[ -n "${POD_IP:-}" ] && [ -n "${SSH_PORT:-}" ] || fail "15 分待っても SSH の IP とポートが出ない"
echo "[POD   ] SSH の IP とポートを API から$([ $REATTACH = 1 ] && echo '取り直した' || echo '受け取った') (<POD_IP>:<SSH_PORT>)"
for _ in $(seq 1 60); do
    pod_ssh true >/dev/null && break
    sleep 5
done
pod_ssh 'echo "[POD   ] 入れた: $(hostname)"; nvidia-smi --query-gpu=name,memory.total --format=csv,noheader' \
    || fail "Pod に SSH で入れない"
stage_done

# ---- 3. Pod を準備する (つなぎ直しは, 済んでいるものを飛ばす) ----
NEED_PREP=1; NEED_CKPT=1; NEED_OPENPI=1; NEED_VENV=1
if [ $REATTACH = 1 ]; then
    stage "3. Pod の中を確かめる (つなぎ直し)"
    lab_pod_probe "$OPENPI_COMMIT" "$CKPT" "$NORM_STATS" || fail "Pod の中を確かめられない"
    yn() { [ "$1" = 1 ] && echo ある || echo ない; }
    if [ "$PROBE_CKPT" = 0 ]; then ckpt_msg=ない
    elif [ "$PROBE_MARK" = 1 ]; then ckpt_msg=ある
    else ckpt_msg="ある (取り終えた印 .lab_download_ok はない. 依頼 44 より前の lab_start.sh で取ったものなど)"; fi
    echo "[POD   ] serve_policy.py のプロセス: $(yn "$PROBE_SERVE")  8000 番の LISTEN: $(yn "$PROBE_LISTEN")"
    echo "[POD   ] モデル ($CKPT): $ckpt_msg  openpi ($OPENPI_COMMIT): $(yn "$PROBE_OPENPI")  .venv: $(yn "$PROBE_VENV")"
    if [ "$PROBE_SERVE" = 1 ] && [ "$PROBE_LISTEN" = 0 ]; then
        # 読み込みの途中かもしれないので少し待つ (エイイチの決め: 30 秒待ち, それでも待っていなければ起動し直す)
        echo "[POD   ] serve_policy.py は動いているが 8000 番で待っていない. 読み込み中かもしれないので $LISTEN_WAIT 秒まで待ちます"
        waited=0
        while [ $waited -lt $LISTEN_WAIT ]; do
            sleep $LISTEN_POLL; waited=$((waited + LISTEN_POLL))
            lab_pod_probe "$OPENPI_COMMIT" "$CKPT" "$NORM_STATS" || fail "Pod の中を確かめられない"
            [ "$PROBE_LISTEN" = 1 ] || [ "$PROBE_SERVE" = 0 ] && break
        done
        if [ "$PROBE_LISTEN" = 1 ]; then
            echo "[POD   ] $waited 秒で 8000 番で待ち始めました"
        else
            echo "[POD   ] $waited 秒待っても 8000 番で待っていない$([ "$PROBE_SERVE" = 0 ] && echo " (その間にプロセスが終わった)"). 起動し直します"
        fi
    fi
    if [ "$PROBE_SERVE" = 1 ] && [ "$PROBE_LISTEN" = 1 ]; then
        NEED_PREP=0
        echo "[POD   ] サーバーは 8000 番で待っています. 3b〜4 を飛ばします"
        pod_ssh "grep 'server listening' /workspace/logs/serve.log | tail -1"
    else
        echo "[POD   ] サーバーが動いていません. ないものだけやり直してサーバーを起動します"
        # 起動し直すと serve.log を空にするので, 前の serve.log を先に持ち帰る (落ちたわけを残す. だめでも続ける)
        dest="$LOGDIR/serve_before_restart_$(date +%Y%m%d_%H%M%S).log"
        if pod_scp "root@$POD_IP:/workspace/logs/serve.log" "$dest"; then
            echo "[LOG   ] 前の serve.log を持ち帰った → $dest ($(wc -l < "$dest" | tr -d ' ') 行)"
        else
            rm -f "$dest"; echo "[LOG   ] 前の serve.log は持ち帰れなかった (ないのかもしれない). 続けます"
        fi
        # 中身があっても印がなければ 3b・3c を流し直す (hf download は取り終えたファイルを取り直さない)
        [ "$PROBE_CKPT" = 1 ] && [ "$PROBE_MARK" = 1 ] && NEED_CKPT=0
        [ "$PROBE_OPENPI" = 1 ] && NEED_OPENPI=0
        [ "$PROBE_VENV" = 1 ] && [ $NEED_OPENPI = 0 ] && NEED_VENV=0     # openpi を取り直したら uv sync もやり直す
    fi
    stage_done
fi
skip() { T_NAMES+=("$1 (飛ばした)"); T_SECS+=(0); echo; echo "=== $1 → 飛ばす ($2)"; }

# スクリプトはつなぎ直しでも毎回送る (Mac 側の pi05/ が変わっているかもしれないため. 動いているサーバーには効かない)
stage "3a. スクリプトを送る"
pod_ssh "mkdir -p /workspace/r2r2r /workspace/logs /workspace/ckpt" || fail "Pod にフォルダを作れない"
pod_scp -r "$HERE/pi05" "root@$POD_IP:/workspace/r2r2r/" || fail "pi05/ を送れない"
pod_scp "$ADD_TRIM" "root@$POD_IP:/workspace/r2r2r/" || fail "add_trim_config.py を送れない"
stage_done

if [ $NEED_PREP = 0 ]; then
    skip "3b〜4. 準備とサーバーの起動" "サーバーが 8000 番で待っている"
else
    if [ $NEED_CKPT = 1 ]; then
        stage "3b. huggingface_hub を入れ, 鍵を渡す"
        pod_ssh "$(lab_cmd_hf_install)" || fail "huggingface_hub を入れられない"
        lab_send_hf_token || fail "Hugging Face の鍵を渡せない"     # 鍵は ssh の標準入力で渡す (lab_lib.sh)
        pod_ssh 'hf auth whoami' || fail "Hugging Face に入れない (鍵が違う: $HF_TOKEN_FILE を確かめる)"
        stage_done

        stage "3c. モデルを取る (hf download)"
        pod_ssh "$(lab_cmd_ckpt "$HF_REPO" "$CKPT")" || fail "hf download に失敗"
        stage_done
    else
        skip "3b・3c. huggingface_hub・モデル" "モデルがある"
    fi

    if [ $NEED_OPENPI = 1 ]; then
        stage "3d. openpi を取る ($OPENPI_COMMIT)"
        pod_ssh "$(lab_cmd_openpi "$OPENPI_COMMIT")" || fail "openpi を取れない"
        stage_done
    else
        skip "3d. openpi を取る ($OPENPI_COMMIT)" "openpi がある"
    fi

    if [ $NEED_VENV = 1 ]; then
        stage "3e. uv sync"
        pod_ssh "$(lab_cmd_uv_sync)" || fail "uv sync に失敗"
        stage_done
    else
        skip "3e. uv sync" ".venv がある"
    fi

    # 3f は何度流しても同じ結果 (install_addon.py・add_trim_config.py とも二重に入れない) なので毎回流す
    stage "3f. addon・設定・assets"
    pod_ssh "$(lab_cmd_addon)" || fail "addon・設定・assets に失敗"
    stage_done

    # ---- 3g・4. サーバーを tmux の中で起動し, 受け付けを始めるまで待つ ----
    stage "4. サーバーの起動"
    pod_ssh 'command -v tmux >/dev/null || { apt-get update -qq && apt-get install -y -qq tmux >/dev/null; }' \
        || fail "tmux を入れられない"
    # 起動し直すときは, 前の serve_policy.py が終わってから起動する (8000 番・GPU を空ける. 止めるのは別の ssh で)
    pod_ssh "$(lab_serve_stop_cmd)" || fail "前の serve_policy.py を止められない"
    pod_ssh ": > /workspace/logs/serve.log; \
             tmux new -d -s serve 'cd /workspace/openpi && source \$HOME/.local/bin/env && uv run scripts/serve_policy.py policy:checkpoint --policy.config=pi05_mecharm_trim_lora --policy.dir=/workspace/ckpt/3999 2>&1 | tee /workspace/logs/serve.log'" \
        || fail "serve_policy.py を起動できない"
    listening=0
    for _ in $(seq 1 120); do
        if pod_ssh "grep -q 'server listening on 0.0.0.0:8000' /workspace/logs/serve.log"; then listening=1; break; fi
        if ! pod_ssh "tmux has-session -t serve 2>/dev/null"; then
            pod_ssh "tail -20 /workspace/logs/serve.log"
            fail "serve_policy.py が終わってしまった (上は serve.log の末尾)"
        fi
        sleep 10
    done
    [ $listening = 1 ] || { pod_ssh "tail -20 /workspace/logs/serve.log"; fail "20 分待っても server listening が出ない"; }
    pod_ssh "grep 'server listening' /workspace/logs/serve.log | tail -1"
    stage_done
fi

# ---- 5. トンネル ----
stage "5. Mac の 8000 番へのトンネル"
if lsof -nP -iTCP:$TUNNEL_PORT -sTCP:LISTEN >/dev/null 2>&1; then
    fail "Mac の $TUNNEL_PORT 番がすでに使われている (前のトンネルなら ./lab_stop.sh で閉じる)"
fi
ssh -N -L $TUNNEL_PORT:localhost:$TUNNEL_PORT "${SSH_OPTS[@]}" -o ExitOnForwardFailure=yes \
    -p "$SSH_PORT" "root@$POD_IP" 2>&1 | lab_mask &
sleep 1
TUNNEL_PID=$(lab_tunnel_pids | head -1)
for _ in $(seq 1 15); do
    lsof -nP -iTCP:$TUNNEL_PORT -sTCP:LISTEN >/dev/null 2>&1 && break
    sleep 1
done
lsof -nP -iTCP:$TUNNEL_PORT -sTCP:LISTEN >/dev/null 2>&1 || fail "トンネルを張れない"
echo "[TUNNEL] localhost:$TUNNEL_PORT → Pod の 8000 番 (pid $TUNNEL_PID. このスクリプトが終わると閉じる)"
stage_done

cd "$HERE" || fail "$HERE に移れない"

# ---- 6. デーモン ----
if [ $DRY_RUN = 0 ]; then
    stage "6. デーモンの状態"
    st=$("$HERE/tools/daemon_ctl.sh" status 2>&1); echo "$st"
    if ! grep -q ':5555 ' <<<"$st"; then
        echo "[DAEMON] 5555 番で受け付けていないので起動します"
        st=$("$HERE/tools/daemon_ctl.sh" start 2>&1); echo "$st"
        grep -q ':5555 ' <<<"$st" || fail "デーモンが 5555 番で受け付けない (tools/daemon_ctl.sh log で確かめる)"
    fi
    stage_done
fi

# ---- 7・8. 位置合わせと実機推論 (くり返せる) ----
# ここから先の Ctrl-C は, 動いている Python (位置合わせ・推論) だけを止める (このスクリプトは続ける)
trap ':' INT
round=0
while :; do
    round=$((round + 1))
    if [ $DRY_RUN = 0 ]; then
        stage "7. HOME → 爪を開く → 位置合わせ ($round 回目)"
        python3 dynamic_tracking_limits/place/place_modules/08_home_from_anywhere.py || fail "HOME に戻せない"
        python3 -c "from visual_servo_pick import get_daemon_socket, daemon_set_gripper, close_daemon_socket; get_daemon_socket(); daemon_set_gripper(0); close_daemon_socket()" \
            || fail "爪を開けない"
        python3 pi05/camera_align.py --raw-dir "$RAW" --episode "$EPISODE" --scale 5 --alpha 0.7
        rc=$?
        [ $rc = 0 ] || [ $rc = 130 ] || fail "camera_align.py が失敗 (終了コード $rc)"
        stage_done
        stage "8. 実機推論 ($round 回目)"
        python3 pi05/run_policy_client.py --raw-dir "$RAW" --steps-per-chunk 10 --via-stride 3 \
            --settle-retries 0 --state-from-target --execute --ignore-framing
        rc=$?
    else
        stage "8. 推論のクライアント (ドライラン・--dry-keys. アームには触らない)"
        python3 pi05/run_policy_client.py --raw-dir "$RAW" --steps-per-chunk 10 --via-stride 3 \
            --settle-retries 0 --state-from-target --dry-keys --ignore-framing
        rc=$?
    fi
    echo "[CLIENT] run_policy_client.py が終わった (終了コード $rc)"
    stage_done
    if [ -t 0 ]; then
        read -r -p "もう1回やりますか? y=やる / 他=終わる > " again
    else
        again=n
    fi
    [ "$again" = y ] || break
done

# ---- 9. Pod は残す ----
summary
echo
echo "[LAB   ] Pod は残しています (id $POD_ID). 片付けは ./lab_stop.sh --log $LOG"
}

main "$@"; exit $?