#!/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 回まで作り直す
# 2. Pod が立ち上がり, SSH の IP とポートが API から出るまで待つ (画面・ログには出さない)
# 3. Pod を準備する (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
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,27p' "$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=215abfb
TEMPLATE_ID=runpod-torch-v280
IMAGE=runpod/pytorch:1.0.2-cu1281-torch280-ubuntu2404

# ---- 段階ごとの秒数 ----
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 "=== かかった時間"
    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=""
cleanup() {
    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

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 を作る ----
stage "1. Pod を作る"
if [ -f "$STATE" ]; then
    OLD_ID=$(jq -r '.id' "$STATE")
    if runpod_api GET "/pods/$OLD_ID" >/dev/null 2>&1; then
        echo "[WARN  ] $LOG の Pod (id $OLD_ID) がまだ残っています. 新しく作らずに止まります."
        echo "         消すなら ./lab_stop.sh --log $LOG"
        exit 1
    fi
fi
others=$(runpod_api GET /pods) || fail "Pod の一覧を API から取れない (HTTP $RUNPOD_HTTP)"
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)"
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 から受け取った (<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 を準備する ----
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

stage "3b. huggingface_hub を入れ, 鍵を渡す"
pod_ssh "pip install -q -U --break-system-packages huggingface_hub" || fail "huggingface_hub を入れられない"
# 鍵は ssh の標準入力で渡す (コマンドの引数・履歴に残らない). hf はこのファイルを読む
tr -d '\r\n' < "$HF_TOKEN_FILE" | pod_ssh 'umask 077; mkdir -p ~/.cache/huggingface; cat > ~/.cache/huggingface/token' \
    || fail "Hugging Face の鍵を渡せない"
pod_ssh 'hf auth whoami' || fail "Hugging Face に入れない (鍵が違う: $HF_TOKEN_FILE を確かめる)"
stage_done

stage "3c. モデルを取る (hf download)"
pod_ssh "hf download $HF_REPO --local-dir /workspace/ckpt/3999 && du -sh /workspace/ckpt/3999" \
    || fail "hf download に失敗"
stage_done

stage "3d. openpi を取る ($OPENPI_COMMIT)"
pod_ssh "cd /workspace && { [ -d openpi/.git ] || git clone -q --recurse-submodules https://github.com/Physical-Intelligence/openpi.git; } \
         && cd openpi && git checkout -q $OPENPI_COMMIT && git log --oneline -1" \
    || fail "openpi を取れない"
stage_done

stage "3e. uv sync"
pod_ssh "curl -LsSf https://astral.sh/uv/install.sh | sh >/dev/null && source \$HOME/.local/bin/env \
         && cd /workspace/openpi && GIT_LFS_SKIP_SMUDGE=1 uv sync" || fail "uv sync に失敗"
stage_done

stage "3f. addon・設定・assets"
pod_ssh "cd /workspace/openpi && python3 /workspace/r2r2r/pi05/openpi_addon/install_addon.py /workspace/openpi \
         && python3 /workspace/r2r2r/add_trim_config.py /workspace/openpi \
         && mkdir -p assets/pi05_mecharm_trim_lora/local \
         && cp -r /workspace/ckpt/3999/assets/local/mecharm_redblock_box_teach_v1_trim2 assets/pi05_mecharm_trim_lora/local/" \
    || 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 を入れられない"
pod_ssh "tmux kill-session -t serve 2>/dev/null; : > /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

# ---- 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"