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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck source=lib.sh
source "${SCRIPT_DIR}/lib.sh"

# Run the end-to-end flow using the CLI submitter (argus.cli).
# This script restarts the Ray cluster and submits PPO/GRPO/SFT sequentially.

EXPECTED_RAY_NODES="${EXPECTED_RAY_NODES:-3}" # head + 2 workers

ray_wait_ready() {
  local tries="${1:-60}"
  for i in $(seq 1 "${tries}"); do
    if curl -sS -m 2 "${RAY_DASHBOARD_ADDR}/api/version" >/dev/null 2>&1; then
      echo "[host] ray_dashboard_ready: ${RAY_DASHBOARD_ADDR}"
      return 0
    fi
    echo "[host] waiting ray dashboard... (${i}/${tries})"
    sleep 2
  done
  echo "ERROR: ray dashboard not ready: ${RAY_DASHBOARD_ADDR}" >&2
  return 1
}

ray_wait_nodes() {
  local want="${1:-3}"
  local tries="${2:-60}"
  for i in $(seq 1 "${tries}"); do
    local out n
    out="$(docker exec -i "${HEAD_CONTAINER}" python3 -c "import ray; ray.init(address='auto', ignore_reinit_error=True, log_to_driver=False, logging_level='ERROR'); print(sum(1 for n in ray.nodes() if n.get('Alive')))" 2>/dev/null || true)"
    n="$(printf '%s\n' "${out}" | tail -n 1 | tr -cd '0-9' || true)"
    if [[ "${n}" =~ ^[0-9]+$ ]]; then
      echo "[host] ray_nodes_alive=${n} (want>=${want})"
      if [[ "${n}" -ge "${want}" ]]; then
        return 0
      fi
    else
      echo "[host] waiting ray nodes... (${i}/${tries})"
    fi
    sleep 2
  done
  echo "ERROR: ray nodes not ready (want>=${want})" >&2
  docker exec -i "${HEAD_CONTAINER}" bash -lc "ray status || true" >&2 || true
  return 1
}

submit_and_wait() {
  local taskspec_in_container="$1"
  local sid
  local out

  echo "[host] submit via SDK: ${taskspec_in_container}"
  out="$(dexec "${HEAD_CONTAINER}" bash -lc "cd /workspace/mvp/py && python3 -m argus.cli.run --config /workspace/mvp/configs/dev.yaml --taskspec '${taskspec_in_container}' --action submit --no-wait" | tr -d '\r')"
  sid="$(printf '%s\n' "${out}" | tail -n 1)"

  if [[ -z "${sid}" ]]; then
    echo "[host] failed to parse submission id from output:" >&2
    printf '%s\n' "${out}" >&2
    exit 1
  fi

  echo "[host] submitted: ${sid}"

  while true; do
    st="$(dexec "${HEAD_CONTAINER}" bash -lc "cd /workspace/mvp/py && python3 -m argus.cli.run --config /workspace/mvp/configs/dev.yaml --action status --submission-id '${sid}'" | tr -d '\r' | tail -n 1)"
    echo "[host] status: ${sid} -> ${st}"
    case "${st}" in
      *SUCCEEDED*)
        return 0
        ;;
      *FAILED*|*STOPPED*)
        echo "[host] job failed: ${sid} (${st})" >&2
        echo "[host] last logs:" >&2
        dexec "${HEAD_CONTAINER}" bash -lc "cd /workspace/mvp/py && python3 -m argus.cli.run --config /workspace/mvp/configs/dev.yaml --action logs --submission-id '${sid}' --tail 200" >&2 || true
        return 1
        ;;
      *)
        sleep 10
        ;;
    esac
  done
}

"${SCRIPT_DIR}/00_prereq_check.sh"
"${SCRIPT_DIR}/03_cleanup_v1_legacy.sh"
"${SCRIPT_DIR}/05_ensure_verl_repo.sh"
"${SCRIPT_DIR}/01_up.sh"
ray_wait_ready 60
ray_wait_nodes "${EXPECTED_RAY_NODES}" 120
"${SCRIPT_DIR}/30_prepare_data_and_model.sh"
"${SCRIPT_DIR}/12_install_py_deps.sh"
submit_and_wait /workspace/mvp/taskspecs/ppo.yaml
submit_and_wait /workspace/mvp/taskspecs/grpo.yaml
submit_and_wait /workspace/mvp/taskspecs/sft.yaml
"${SCRIPT_DIR}/50_status.sh"
