#!/bin/bash
# lab_stop.sh — lab_start.sh で作った Pod を片付ける (依頼 37)
#
#   ./lab_stop.sh --log Log17      serve.log を持ち帰り, Pod を Terminate まで消し, 費用の見込みを出す
#   ./lab_stop.sh                  --log なしは, 最後に lab_start.sh を動かした Log を使う
#   ./lab_stop.sh --check          Pod が残っていないかを確かめるだけ (消さない. 残っていれば終了コード 1)
#   ./lab_stop.sh --log Log17 --skip-log   serve.log の持ち帰りは1回だけ試し, だめでもやり直さずに消す
#
# 1. Pod の /workspace/logs/serve.log を /Volumes/KIOXIA/r2r2r-data/lab/<Log>/ へ持ち帰る
#    (serve.log がすでにあれば serve_<時刻>.log にする). 持ち帰れなければ 10 秒待ってやり直す (合わせて 3 回).
#    3 回ともだめなら「serve.log を持ち帰れなかった」と画面と lab_stop のログに残し, そのまま 2 へ進んで消す
# 2. Pod を Terminate (DELETE) する. 一覧からも消えたことを API で確かめる
# 3. 借りていた時間 × 1時間の料金 で費用の見込みを出す (ドルと円. 1ドル=158円)
# 4. 裏の 8000 番のトンネルが残っていれば閉じる

set -u
set -o pipefail
HERE="$(cd "$(dirname "$0")" && pwd)"
# shellcheck source=lab_lib.sh
. "$HERE/lab_lib.sh"

LOG=""; CHECK=0; SKIP_LOG=0
while [ $# -gt 0 ]; do
    case "$1" in
        --log) LOG="$2"; shift 2 ;;
        --check) CHECK=1; shift ;;
        --skip-log) SKIP_LOG=1; shift ;;
        -h|--help) sed -n '2,16p' "$0"; exit 0 ;;
        *) echo "[ERROR] 知らない引数: $1  (./lab_stop.sh --help)" >&2; exit 2 ;;
    esac
done
lab_need_key "$RUNPOD_KEY_FILE" "Runpod の API キー"

close_tunnel() {
    local pids
    pids=$(lab_tunnel_pids)
    if [ -n "$pids" ]; then
        kill $pids 2>/dev/null
        echo "[TUNNEL] 残っていた 8000 番のトンネルを閉じました (pid $(echo $pids))"
    else
        echo "[TUNNEL] 8000 番のトンネルは残っていません"
    fi
}

list_pods() {               # 残っている Pod を出す. 戻り値は数
    local j n
    j=$(runpod_api GET /pods) || { echo "[ERROR] Pod の一覧を API から取れない (HTTP $RUNPOD_HTTP)" >&2; return 255; }
    n=$(jq 'length' <<<"$j")
    if [ "$n" -eq 0 ]; then
        echo "[CHECK ] 残っている Pod はありません"
    else
        echo "[CHECK ] 残っている Pod が $n 個あります:"
        jq -r '.[] | "         \(.id)  \(.name)  \(.desiredStatus)  \(.costPerHr) ドル/時"' <<<"$j"
    fi
    return "$n"
}

if [ $CHECK = 1 ]; then
    list_pods; n=$?
    if [ -n "$(lab_tunnel_pids)" ]; then
        echo "[CHECK ] 8000 番のトンネルが残っています (pid $(echo $(lab_tunnel_pids)))"
    fi
    [ $n = 0 ] && exit 0 || exit 1
fi

if [ -z "$LOG" ] && [ -s "$LAB_CURRENT" ]; then LOG=$(cat "$LAB_CURRENT"); fi
[ -n "$LOG" ] || { echo "[ERROR] --log Log17 のように Log 番号を付けてください" >&2; exit 2; }
LOGDIR="$LAB_ROOT/$LOG"
STATE="$LOGDIR/pod/lab_pod.json"
if [ -d "$LOGDIR" ]; then      # 画面に出しつつ <Log>/ に残す
    exec > >(tee -a "$LOGDIR/lab_stop_$(date +%Y%m%d_%H%M%S).txt") 2>&1
fi
[ -f "$STATE" ] || { echo "[ERROR] $STATE がありません (lab_start.sh で作った Pod の記録). ./lab_stop.sh --check で確かめる" >&2; exit 1; }
POD_ID=$(jq -r '.id' "$STATE")
COST_PER_HR=$(jq -r '.costPerHr' "$STATE")
GPU_ID=$(jq -r '.gpu' "$STATE")
POD_T0=$(jq -r '.created_epoch' "$STATE")
echo "[STOP  ] $LOG の Pod: id $POD_ID  GPU=$GPU_ID  作った時刻 $(jq -r '.created_at' "$STATE")"
lab_ssh_opts

# ---- 1. serve.log を持ち帰る ----
LOG_TRIES=${LAB_LOG_TRIES:-3}       # 合わせて何回試すか (--skip-log は 1 回)
[ $SKIP_LOG = 1 ] && LOG_TRIES=1
LOG_WAIT=${LAB_LOG_WAIT:-10}        # やり直す前に待つ秒数

fetch_serve_log() {         # serve.log を持ち帰る. LOG_TRIES 回ともだめなら戻り値 1
    local i dest
    dest="$LOGDIR/serve.log"
    [ -e "$dest" ] && dest="$LOGDIR/serve_$(date +%Y%m%d_%H%M%S).log"
    for i in $(seq 1 "$LOG_TRIES"); do
        if ! lab_pod_addr; then
            echo "[LOG   ] ($i/$LOG_TRIES) Pod の SSH の IP とポートが API から出ない (状態 ${POD_STATUS:-?})"
        elif pod_scp "root@$POD_IP:/workspace/logs/serve.log" "$dest"; then
            echo "[LOG   ] serve.log を持ち帰った → $dest ($(wc -l < "$dest" | tr -d ' ') 行)"
            return 0
        else
            echo "[LOG   ] ($i/$LOG_TRIES) serve.log の scp に失敗"
            rm -f "$dest"       # 途中まで書いたものは残さない
        fi
        if [ "$i" -lt "$LOG_TRIES" ]; then
            echo "[LOG   ] ${LOG_WAIT} 秒待ってやり直します"
            sleep "$LOG_WAIT"
        fi
    done
    return 1
}

if runpod_api GET "/pods/$POD_ID" >/dev/null; then
    if ! fetch_serve_log; then
        echo "[WARN  ] serve.log を持ち帰れなかった ($LOG_TRIES 回試した). Pod はこのまま Terminate まで消します"
    fi
else
    echo "[STOP  ] Pod (id $POD_ID) は API から見つかりません (HTTP $RUNPOD_HTTP). すでに消えている見込み"
fi

# ---- 2. Terminate ----
if runpod_api GET "/pods/$POD_ID" >/dev/null; then
    runpod_api DELETE "/pods/$POD_ID" >/dev/null || { echo "[ERROR] Pod を消せない (HTTP $RUNPOD_HTTP)"; exit 1; }
    T_END=$(date +%s)
    echo "[STOP  ] Terminate を送った ($(date -r "$T_END" '+%H:%M:%S'))"
else
    T_END=$(jq -r '.terminated_epoch // empty' "$STATE")    # 前に消したときの時刻 (なければ今)
    [ -n "$T_END" ] || T_END=$(date +%s)
fi
gone=0
for _ in $(seq 1 30); do
    if pods=$(runpod_api GET /pods) \
       && ! jq -e --arg id "$POD_ID" 'any(.[]; .id == $id and .desiredStatus != "TERMINATED")' <<<"$pods" >/dev/null; then
        gone=1; break
    fi
    sleep 5
done
if [ $gone = 1 ]; then
    echo "[STOP  ] Pod (id $POD_ID) が消えたことを API で確かめた"
else
    echo "[ERROR] 2 分半待っても Pod (id $POD_ID) が API に残っています. Runpod の画面で確かめてください"
fi

# ---- 3. 費用の見込み ----
secs=$((T_END - POD_T0))
awk -v s="$secs" -v c="$COST_PER_HR" -v y="$YEN_PER_USD" 'BEGIN {
    h = s / 3600; usd = h * c
    printf "[COST  ] 借りていた時間 %d 分 %d 秒 (%.2f 時間) × 1時間 %.2f ドル = 約 %.2f ドル (約 %d 円. 1ドル=%d円)\n",
           int(s / 60), s % 60, h, c, usd, int(usd * y + 0.5), y }'
jq --argjson t "$T_END" --argjson s "$secs" \
   '. + {terminated_epoch:$t, terminated_at:($t|strflocaltime("%Y-%m-%d %H:%M:%S")), rented_sec:$s}' \
   "$STATE" > "$STATE.tmp" && mv "$STATE.tmp" "$STATE"

# ---- 4. トンネル ----
close_tunnel
list_pods >/dev/null; n=$?
[ $gone = 1 ] || exit 1
[ $n = 0 ] || { echo "[WARN  ] ほかの Pod が $n 個残っています (./lab_stop.sh --check)"; }