#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)"
ACTIVE_STATE_FILE="${TRAINING_STATE_FILE:-${REPO_ROOT}/.runtime/active_training.txt}"
LAST_STATE_FILE="${LAST_TRAINING_STATE_FILE:-${REPO_ROOT}/.runtime/last_training.txt}"
RESUME_DRY_RUN="${RESUME_DRY_RUN:-0}"

read_state_value() {
  local state_file="$1"
  local key="$2"
  awk -v key="$key" '
    index($0, key "=") == 1 {
      sub(/^[^=]*=/, "")
      print
      exit
    }
  ' "$state_file"
}

if [[ -f "$ACTIVE_STATE_FILE" ]]; then
  STATE_FILE="$ACTIVE_STATE_FILE"
elif [[ -f "$LAST_STATE_FILE" ]]; then
  STATE_FILE="$LAST_STATE_FILE"
else
  echo "ERROR: no active or previous training state found." >&2
  echo "checked: $ACTIVE_STATE_FILE" >&2
  echo "checked: $LAST_STATE_FILE" >&2
  exit 2
fi

CONTAINER_NAME="$(read_state_value "$STATE_FILE" container_name)"
CHECKPOINT_DIR="$(read_state_value "$STATE_FILE" checkpoint_dir)"
RUN_DIR="$(read_state_value "$STATE_FILE" run_dir)"
LOG_FILE="$(read_state_value "$STATE_FILE" log_file)"
LAUNCHER_SCRIPT="$(read_state_value "$STATE_FILE" launcher_script)"

CONTAINER_NAME="${CONTAINER_NAME:-laoyao-2b-pretraining}"
if [[ "$(docker inspect -f '{{.State.Running}}' "$CONTAINER_NAME" 2>/dev/null || true)" == "true" ]]; then
  echo "ERROR: training container is already running: $CONTAINER_NAME" >&2
  exit 3
fi
if [[ -z "$LAUNCHER_SCRIPT" || ! -f "$LAUNCHER_SCRIPT" ]]; then
  echo "ERROR: recorded launcher script is missing: ${LAUNCHER_SCRIPT:-<empty>}" >&2
  exit 4
fi
if [[ -z "$CHECKPOINT_DIR" || ! -f "${CHECKPOINT_DIR}/latest_checkpointed_iteration.txt" ]]; then
  echo "ERROR: resumable checkpoint tracker is missing: ${CHECKPOINT_DIR:-<empty>}" >&2
  exit 5
fi

latest_iteration="$(tr -d '[:space:]' < "${CHECKPOINT_DIR}/latest_checkpointed_iteration.txt")"
cat <<EOF
state_file=$STATE_FILE
container_name=$CONTAINER_NAME
launcher_script=$LAUNCHER_SCRIPT
run_dir=$RUN_DIR
checkpoint_dir=$CHECKPOINT_DIR
latest_checkpointed_iteration=$latest_iteration
log_file=$LOG_FILE
EOF

if [[ "$RESUME_DRY_RUN" == "1" ]]; then
  echo "Dry run passed; training was not launched."
  exit 0
fi

CONTAINER_NAME="$CONTAINER_NAME" bash "$LAUNCHER_SCRIPT"