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

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

# E2E v2.5:
# - Clean legacy env
# - Start containers
# - Ray head/worker auto-start (supervisor in-container)
# - Verify discovery + worker join
# - Prepare data/model/code (reuses existing downloads)
# - Start API
# - Create user + issue token
# - Submit PPO/GRPO/SFT via API and wait

API_ADDR="${API_ADDR:-http://127.0.0.1:8080}"
ADMIN_TOKEN="${MVP_INTERNAL_TOKEN:-}"
USER_ID="${USER_ID:-alice}"
RESET_DB="${RESET_DB:-1}"
EXPECTED_RAY_NODES="${EXPECTED_RAY_NODES:-3}" # head + 2 workers
CLUSTER_NAME="${CLUSTER_NAME:-argus-ray}"

if [[ -z "${ADMIN_TOKEN}" ]]; then
  echo "ERROR: MVP_INTERNAL_TOKEN must be set in host env (admin token)" >&2
  exit 1
fi

api_curl_admin() {
  curl -sS -H "Authorization: Bearer ${ADMIN_TOKEN}" "$@"
}

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

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_taskspec() {
  local token="$1"
  local taskspec_path="$2"
  echo "[host] submit via API (user=${USER_ID}): ${taskspec_path}" >&2
  local resp
  resp="$(curl -sS -H "Authorization: Bearer ${token}" -H "Content-Type: application/yaml" --data-binary @"${taskspec_path}" "${API_ADDR}/api/v2/tasks")"
  echo "[host] submit_resp: ${resp}" >&2
  printf '%s' "${resp}" | python3 -c 'import sys,json; print(json.load(sys.stdin)["task_id"])'
}

wait_task() {
  local token="$1"
  local task_id="$2"
  while true; do
    local body state
    body="$(curl -sS -H "Authorization: Bearer ${token}" "${API_ADDR}/api/v2/tasks/${task_id}")"
    state="$(printf '%s' "${body}" | python3 -c 'import sys,json; print(json.load(sys.stdin)["state"])')"
    echo "[host] task ${task_id}: ${state}"

    if [[ "${state}" == "SUCCEEDED" ]]; then
      return 0
    fi
    if [[ "${state}" == "FAILED" || "${state}" == "CANCELED" ]]; then
      echo "[host] terminal=${state}; tail logs (best-effort):" >&2
      curl -sS -H "Authorization: Bearer ${token}" "${API_ADDR}/api/v2/tasks/${task_id}/logs?tail=200" >&2 || true
      return 1
    fi
    sleep 10
  done
}

echo "[host] ===== run_all_v25_api.sh begin ====="

"${SCRIPT_DIR}/00_prereq_check.sh"
"${SCRIPT_DIR}/03_cleanup_v1_legacy.sh"
"${SCRIPT_DIR}/04_cleanup_v2_legacy.sh"

echo "[host] bring down existing containers (best-effort)"
"${SCRIPT_DIR}/02_down.sh" || true

echo "[host] (re)create containers"
"${SCRIPT_DIR}/01_up.sh"

echo "[host] wait ray head ready (ray is supervised in-container; head has no compute)"
ray_wait_ready 60

echo "[host] verify head discovery record (supervised in-container)"
HEAD_IP_FILE="${SHARED_ROOT}/ray/discovery/${CLUSTER_NAME}/head.json"
dexec "${HEAD_CONTAINER}" bash -lc "test -f '${HEAD_IP_FILE}' && python3 -c 'import json,sys; print(json.load(open(sys.argv[1]))[\"job_server_url\"])' '${HEAD_IP_FILE}' || true"

echo "[host] wait workers join (watchdog auto-connect; supervised in-container)"
ray_wait_nodes "${EXPECTED_RAY_NODES}" 120

echo "[host] prepare data/model/code snapshot (idempotent; reuse cache)"
"${SCRIPT_DIR}/30_prepare_data_and_model.sh"

echo "[host] install api deps in head container"
"${SCRIPT_DIR}/12_install_api_deps.sh"

echo "[host] stop api (best-effort)"
"${SCRIPT_DIR}/61_stop_api.sh" || true

if [[ "${RESET_DB}" == "1" ]]; then
  echo "[host] reset api sqlite db in container (best-effort)"
  docker exec -i "${HEAD_CONTAINER}" bash -lc "rm -f /private/common/db/mvp.sqlite3 /private/common/db/mvp.sqlite3-wal /private/common/db/mvp.sqlite3-shm || true"
else
  echo "[host] keep existing api sqlite db (RESET_DB=${RESET_DB})"
fi

echo "[host] start api (admin token via env)"
MVP_INTERNAL_TOKEN="${ADMIN_TOKEN}" "${SCRIPT_DIR}/60_start_api.sh"
api_wait_ready 60

echo "[host] ensure user exists + issue token"
api_curl_admin -H "Content-Type: application/json" -d "{\"user_id\":\"${USER_ID}\"}" "${API_ADDR}/api/v2/users" >/dev/null 2>&1 || true
USER_TOKEN="$(api_curl_admin -X POST "${API_ADDR}/api/v2/users/${USER_ID}/tokens" | python3 -c 'import sys,json; print(json.load(sys.stdin)["token"])')"
echo "[host] user_token_issued: user=${USER_ID}"

PPO_TASK_ID="$(submit_taskspec "${USER_TOKEN}" "${ROOT_DIR}/taskspecs/ppo.yaml")"
GRPO_TASK_ID="$(submit_taskspec "${USER_TOKEN}" "${ROOT_DIR}/taskspecs/grpo.yaml")"
SFT_TASK_ID="$(submit_taskspec "${USER_TOKEN}" "${ROOT_DIR}/taskspecs/sft.yaml")"

echo "[host] submitted task ids:"
echo "  ppo=${PPO_TASK_ID}"
echo "  grpo=${GRPO_TASK_ID}"
echo "  sft=${SFT_TASK_ID}"

echo "[host] wait for tasks (in submission order)"
wait_task "${USER_TOKEN}" "${PPO_TASK_ID}"
wait_task "${USER_TOKEN}" "${GRPO_TASK_ID}"
wait_task "${USER_TOKEN}" "${SFT_TASK_ID}"

echo "[host] ===== run_all_v25_api.sh done ====="
