# lab_train_fakes.sh — lab_train.sh --dry-run が読む差し替え (単体では動かさない. 依頼 49)
#
# 本物の Pod を借りずに lab_train.sh の流れを最後まで流す:
#   - 書く場所 (DATA_ROOT・LAB_ROOT・checkpoints)・鍵を一時フォルダ LAB_DRY_ROOT (なければ作る) へ
#   - runpod_api / pod_ssh / pod_scp / Mac の hf download / 通知 / caffeinate を偽物にする
#     (呼ばれた中身を $LAB_DRY_ROOT/calls.log・ssh.log・scp.log・notify.log に残す)
#   - 待ちを短くする (進みを見るのは 1 秒おき)
# どうなるかは LAB_FAKE_SCENARIO で決める (test_lab_train.sh が使う):
#   success (既定)  学習が最後まで終わる
#   crash           学習が終了コード 1 で落ちる
#   killed          学習のプロセスが終わりの印を書かずに消える
#   timeout         学習が終わらない (--limit-hours を小さくして使う)
#   stall           学習のログが長く書かれない
#   hf_fail         Hugging Face に上げられない (Pod から scp で KIOXIA へ取る)
#   dl_mismatch     Hugging Face から取ったものが Pod 側の一覧と合わない
#   gpu_busy        A40 が借りられない (RTX A6000 で借りる)
#   all_busy        どの GPU も借りられない
#   delete_fail     Pod を消せない
#   delete_flaky    1 回目の DELETE だけ失敗する
#   late_finish     上限を超えた直後の確かめで, 学習が終わっている (消さずに持ち帰る)
#   wd_missing      Pod 用の鍵が見つからず, Pod 側の見張りが働かない
#   list_empty      Pod 側のチェックポイントの一覧が取れない
#   list_partial    Pod 側の一覧の途中で find が失敗する (一部の行と終了コード 1)
#   late_start      学習を始めるのに手間取り, 見守りに入ったときにはもう上限を過ぎている (学習は終わっている)

LAB_DRY_ROOT=${LAB_DRY_ROOT:-$(mktemp -d /tmp/lab_train_dry.XXXXXX)}
LAB_FAKE_SCENARIO=${LAB_FAKE_SCENARIO:-success}
FAKE_POD_ID=dryrunpod0001
echo "[DRY   ] ドライラン: Pod は借りません. 書くのは $LAB_DRY_ROOT の中だけ (場合: $LAB_FAKE_SCENARIO)"
if [ -d /Volumes/KIOXIA/r2r2r-data ]; then
    for f in "/Volumes/KIOXIA/r2r2r-data/datasets/mecharm_redblock_box_teach_v1_trim2_corr_v3/mecharm_redblock_box_teach_v1_trim2_corr_v3_raw_episodes.tar" \
             /Volumes/KIOXIA/r2r2r-data/lab/Log11/pod/add_trim_config.py "$RUNPOD_KEY_FILE" "$HF_TOKEN_FILE"; do
        echo "[DRY   ] 本物の $(basename "$f"): $([ -s "$f" ] && echo ある || echo ない)"
    done
fi

DATA_ROOT="$LAB_DRY_ROOT/data"
LAB_ROOT="$DATA_ROOT/lab"; LAB_CURRENT="$LAB_ROOT/.lab_current"
RUNPOD_KEY_FILE="$LAB_DRY_ROOT/keys/runpod_api_key.txt"; HF_TOKEN_FILE="$LAB_DRY_ROOT/keys/hf_token.txt"
POD_SSH_KEY="$LAB_DRY_ROOT/keys/id_ed25519"
mkdir -p "$LAB_DRY_ROOT/keys" "$LAB_DRY_ROOT/fake" "$DATA_ROOT/lab/Log11/pod" \
         "$DATA_ROOT/datasets/mecharm_redblock_box_teach_v1_trim2_corr_v3"
[ -s "$RUNPOD_KEY_FILE" ] || echo "DRY-RUNPOD-SECRET-KEY" > "$RUNPOD_KEY_FILE"
[ -s "$HF_TOKEN_FILE" ] || echo "hf_DRYSECRETTOKEN" > "$HF_TOKEN_FILE"
[ -s "$POD_SSH_KEY" ] || echo "dry-ssh-key" > "$POD_SSH_KEY"
[ -s "$POD_SSH_KEY.pub" ] || echo "ssh-ed25519 AAAA dry" > "$POD_SSH_KEY.pub"
[ -f "$DATA_ROOT/lab/Log11/pod/add_trim_config.py" ] || echo "# dry" > "$DATA_ROOT/lab/Log11/pod/add_trim_config.py"
TARF="$DATA_ROOT/datasets/mecharm_redblock_box_teach_v1_trim2_corr_v3/mecharm_redblock_box_teach_v1_trim2_corr_v3_raw_episodes.tar"
[ -f "$TARF" ] || echo "dry tar" > "$TARF"
touch "$LAB_DRY_ROOT/calls.log" "$LAB_DRY_ROOT/ssh.log" "$LAB_DRY_ROOT/scp.log"

CREATE_WAIT=0; SSH_WAIT=0; LOGIN_WAIT=0; GONE_WAIT=0; FETCH_WAIT=0
POLL_SEC=${LAB_DRY_POLL_SEC:-1}
NOTIFY_TITLE="$NOTIFY_TITLE (ドライラン)"

# 偽物のチェックポイント (Pod 側の一覧と同じ中身を Hugging Face / scp で取ったことにする)
FAKE_CKPT_FILES=$'_CHECKPOINT_METADATA\t426\nassets/local/mecharm_redblock_box_teach_v1_trim2_corr_v3/norm_stats.json\t2120\nparams/_METADATA\t33052\nparams/d/0123abcd\t1048576\nparams/manifest.ocdbt\t118'
fake_make_ckpt() {          # fake_make_ckpt フォルダ [ずらす] → FAKE_CKPT_FILES のとおりのファイルを作る
    local d="$1" p n
    while IFS=$'\t' read -r p n; do
        [ "${2:-}" = shift ] && [ "$p" = params/manifest.ocdbt ] && n=$((n + 1))
        mkdir -p "$d/$(dirname "$p")"; head -c "$n" /dev/zero > "$d/$p"
    done <<<"$FAKE_CKPT_FILES"
}
fake_n() { echo x >> "$LAB_DRY_ROOT/fake/$1.n"; wc -l < "$LAB_DRY_ROOT/fake/$1.n" | tr -d ' '; }

runpod_api() {
    local m="$1" p="$2" data="${3:-}" gpu
    echo "$m $p" >> "$LAB_DRY_ROOT/calls.log"
    case "$m $p" in
        "POST /pods")
            gpu=$(jq -r '.gpuTypeIds[0]' <<<"$data"); echo "POST gpu=$gpu name=$(jq -r .name <<<"$data")" >> "$LAB_DRY_ROOT/calls.log"
            if [ "$LAB_FAKE_SCENARIO" = all_busy ] || { [ "$LAB_FAKE_SCENARIO" = gpu_busy ] && [ "$gpu" = "NVIDIA A40" ]; }; then
                echo '{"error":"There are no longer any instances available with the requested specifications."}'; RUNPOD_HTTP=500
            else
                echo "{\"id\":\"$FAKE_POD_ID\",\"costPerHr\":0.44}"; RUNPOD_HTTP=200
                echo "$gpu" > "$LAB_DRY_ROOT/fake/created"
            fi ;;
        "GET /pods")
            if [ -f "$LAB_DRY_ROOT/fake/created" ] && [ ! -f "$LAB_DRY_ROOT/fake/deleted" ]; then
                echo "[{\"id\":\"$FAKE_POD_ID\",\"name\":\"p2m-train-x\",\"desiredStatus\":\"RUNNING\",\"costPerHr\":0.44}]"
            else
                cat "$LAB_DRY_ROOT/fake/others.json" 2>/dev/null || echo '[]'
            fi
            RUNPOD_HTTP=200 ;;
        "GET /pods/"*)
            if [ "${p#/pods/}" = "$FAKE_POD_ID" ] && [ -f "$LAB_DRY_ROOT/fake/created" ] && [ ! -f "$LAB_DRY_ROOT/fake/deleted" ]; then
                echo "{\"id\":\"$FAKE_POD_ID\",\"desiredStatus\":\"RUNNING\",\"publicIp\":\"10.9.8.7\",\"portMappings\":{\"22\":40123}}"
                RUNPOD_HTTP=200
            else
                echo '{"error":"pod not found"}'; RUNPOD_HTTP=404
            fi ;;
        "DELETE /pods/"*)
            if [ "$LAB_FAKE_SCENARIO" = delete_fail ] || { [ "$LAB_FAKE_SCENARIO" = delete_flaky ] && [ "$(fake_n delete)" = 1 ]; }; then
                echo '{"error":"internal"}'; RUNPOD_HTTP=500
            else echo '{}'; RUNPOD_HTTP=200; touch "$LAB_DRY_ROOT/fake/deleted"; fi ;;
        *) echo '{}'; RUNPOD_HTTP=404 ;;
    esac
    case "$RUNPOD_HTTP" in 2??) return 0 ;; *) return 1 ;; esac
}

pod_ssh() {
    printf '%s\n----\n' "$*" >> "$LAB_DRY_ROOT/ssh.log"
    cat > /dev/null
    [ -f "$LAB_DRY_ROOT/fake/deleted" ] && return 255
    local n
    case "$*" in
        *": lt=watchdog"*) [ "$LAB_FAKE_SCENARIO" = wd_missing ] && echo "WATCHDOG run=1 pod=1 key=0 ctl=0" || echo "WATCHDOG run=1 pod=1 key=1 ctl=1" ;;
        *": lt=start_train"*) [ "$LAB_FAKE_SCENARIO" = late_start ] && sleep 9; echo STARTED ;;
        *": lt=status"*)
            n=$(fake_n status)
            if [ "$LAB_FAKE_SCENARIO" = late_start ]; then echo "STATUS run=0 exit=0 age=4 bar=1500 logstep=1400 loss=0.0215 ckpt=1499 cm=1"; return 0; fi
            if [ "$LAB_FAKE_SCENARIO" = late_finish ]; then     # 確かめている間に上限を過ぎ, 学習も終わる
                sleep 9; echo "STATUS run=0 exit=0 age=4 bar=1500 logstep=1400 loss=0.0215 ckpt=1499 cm=1"; return 0
            fi
            case "$LAB_FAKE_SCENARIO:$n" in
                crash:1|killed:1|stall:1|timeout:*|*:1)
                    if [ "$LAB_FAKE_SCENARIO" = stall ]; then echo "STATUS run=1 exit=- age=99999 bar=40 logstep=- loss=- ckpt=- cm=0"
                    else echo "STATUS run=1 exit=- age=3 bar=$((n * 100)) logstep=$((n * 100)) loss=0.0412 ckpt=- cm=0"; fi ;;
                crash:*) echo "STATUS run=0 exit=1 age=5 bar=650 logstep=600 loss=0.0399 ckpt=- cm=0" ;;
                killed:*) echo "STATUS run=0 exit=- age=5 bar=650 logstep=600 loss=0.0399 ckpt=- cm=0" ;;
                *) echo "STATUS run=0 exit=0 age=4 bar=1500 logstep=1400 loss=0.0215 ckpt=1499 cm=1" ;;
            esac ;;
        *": lt=list_ckpt"*|*": lt=ls_all"*)
            if [ "$LAB_FAKE_SCENARIO" = list_partial ] && [[ "$*" == *"lt=list_ckpt"* ]]; then
                sed 's/^/F\t/' <<<"$FAKE_CKPT_FILES" | head -2; echo "find: './params/d': Input/output error"; return 1
            fi
            if [ -f "$LAB_DRY_ROOT/fake/status.n" ] && [ "$LAB_FAKE_SCENARIO" != crash ] && [ "$LAB_FAKE_SCENARIO" != list_empty ] && [ "$LAB_FAKE_SCENARIO" != killed ] \
               && [ "$LAB_FAKE_SCENARIO" != timeout ] && [ "$LAB_FAKE_SCENARIO" != stall ]; then
                sed 's/^/F\t/' <<<"$FAKE_CKPT_FILES"
            fi ;;
        *"hf_upload_ckpt.py"*)
            if [ "$LAB_FAKE_SCENARIO" = hf_fail ]; then echo "HFUP error HfHubHTTPError: 401 Client Error"; return 6; fi
            echo "HFUP ok PixelstoMotors/pi05-mecharm-trim-corr-lora-v1 に 5 ファイルを上げ, 名前・バイト数が全部合うのを確かめた" ;;
        *"pod_prep_trim_corr.sh"*) echo "== 5. 準備できました (ドライラン)" ;;
        *"hf auth whoami"*) echo "user: PixelstoMotors (ドライラン)" ;;
    esac
    return 0
}

pod_scp() {
    echo "$*" >> "$LAB_DRY_ROOT/scp.log"
    [ -f "$LAB_DRY_ROOT/fake/deleted" ] && return 1
    local last="${!#}"
    case "$*" in
        *":/workspace/openpi/checkpoints/"*) fake_make_ckpt "$last" ;;
        *":/workspace/logs/watchdog.log"*) return 1 ;;      # 見張りが働いていなければ, ないのがふつう
        *" root@"*":/workspace/logs/"*|"root@"*":/workspace/logs/"*) echo "dry log line (${1##*/})" > "$last" ;;
    esac
    return 0
}

lab_have_hf_cli() { return 0; }
lab_hf_download_mac() {
    echo "hf download $1 → $2" >> "$LAB_DRY_ROOT/calls.log"
    fake_make_ckpt "$2" "$([ "$LAB_FAKE_SCENARIO" = dl_mismatch ] && echo shift)"
    mkdir -p "$2/.cache/huggingface/download"; echo x > "$2/.cache/huggingface/download/x.metadata"; echo "*" > "$2/.gitattributes"
    echo x > "$2/._params"; echo x > "$2/.DS_Store"         # KIOXIA (exFAT) が付けるもの
}
lab_notify() { echo "[NOTIFY] $1: $2" | tee -a "$LAB_DRY_ROOT/notify.log"; }
lab_keep_awake() { echo "[MAC   ] (ドライラン) caffeinate は使わない"; }