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

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

API_ADDR="${API_ADDR:-http://127.0.0.1:8080}"
ADMIN_TOKEN="${MVP_INTERNAL_TOKEN:-}"

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

require_cmd jq
require_cmd python3

api_admin() { curl -sS -H "Authorization: Bearer ${ADMIN_TOKEN}" "$@"; }
api_user() { local tok="$1"; shift; curl -sS -H "Authorization: Bearer ${tok}" "$@"; }

wait_state() {
  local tok="$1"
  local task_id="$2"
  local want="$3"
  local tries="${4:-120}"
  for i in $(seq 1 "${tries}"); do
    local st
    st="$(api_user "${tok}" "${API_ADDR}/api/v2/tasks/${task_id}" | jq -r .state)"
    echo "[case] ${task_id} state=${st} (want=${want})"
    if [[ "${st}" == "${want}" ]]; then
      return 0
    fi
    if [[ "${st}" == "FAILED" || "${st}" == "CANCELED" || "${st}" == "SUCCEEDED" ]]; then
      return 1
    fi
    sleep 5
  done
  return 2
}

issue_token() {
  local user_id="$1"
  api_admin -H "Content-Type: application/json" -d "{\"user_id\":\"${user_id}\"}" "${API_ADDR}/api/v2/users" >/dev/null 2>&1 || true
  api_admin -X POST "${API_ADDR}/api/v2/users/${user_id}/tokens" | jq -r .token
}

submit_yaml() {
  local tok="$1"
  local yaml_path="$2"
  api_user "${tok}" -H "Content-Type: application/yaml" --data-binary @"${yaml_path}" "${API_ADDR}/api/v2/tasks" | jq -r .task_id
}

echo "[case] ===== v2.5 e2e cases ====="

echo "[case] HP-1: run_all_v25_api.sh (happy path)"
RESET_DB="${RESET_DB:-1}"
MVP_INTERNAL_TOKEN="${ADMIN_TOKEN}" RESET_DB="${RESET_DB}" "${SCRIPT_DIR}/run_all_v25_api.sh"

echo "[case] E-Auth-1: missing bearer token"
code="$(curl -sS -o /dev/null -w '%{http_code}' "${API_ADDR}/api/v2/queue" || true)"
[[ "${code}" == "401" ]] || { echo "expected 401 got ${code}" >&2; exit 1; }

echo "[case] E-Auth-2: invalid token"
code="$(curl -sS -o /dev/null -w '%{http_code}' -H 'Authorization: Bearer nope' "${API_ADDR}/api/v2/queue" || true)"
[[ "${code}" == "401" ]] || { echo "expected 401 got ${code}" >&2; exit 1; }

echo "[case] setup users alice/bob"
ALICE_TOKEN="$(issue_token alice)"
BOB_TOKEN="$(issue_token bob)"

echo "[case] E-Isolation-1: cross-user task visibility (404)"
TID="$(submit_yaml "${ALICE_TOKEN}" "${ROOT_DIR}/taskspecs/ppo.yaml")"
code="$(api_user "${BOB_TOKEN}" -o /dev/null -w '%{http_code}' "${API_ADDR}/api/v2/tasks/${TID}" || true)"
[[ "${code}" == "404" ]] || { echo "expected 404 got ${code}" >&2; exit 1; }

# Keep the queue clean for subsequent cases (avoid consuming scheduler capacity).
api_user "${ALICE_TOKEN}" -X POST "${API_ADDR}/api/v2/tasks/${TID}:cancel" >/dev/null 2>&1 || true
wait_state "${ALICE_TOKEN}" "${TID}" "CANCELED" 30 || true

echo "[case] E-Input-1: reject non-common inputs"
bad_yaml="$(mktemp)"
cat >"${bad_yaml}" <<'YAML'
workload: "ppo"
submission_id: ""
code_path: "/private/common/code/verl/verl_repo"
model_id: "Qwen/Qwen2.5-0.5B-Instruct"
train_file: "/private/datasets/gsm8k/train.parquet"
val_file: "/private/common/datasets/gsm8k/test.parquet"
nnodes: 2
n_gpus_per_node: 4
total_epochs: 1
total_training_steps: 10
save_freq: 10
test_freq: -1
YAML
code="$(api_user "${ALICE_TOKEN}" -H 'Content-Type: application/yaml' --data-binary @"${bad_yaml}" -o /dev/null -w '%{http_code}' "${API_ADDR}/api/v2/tasks" || true)"
rm -f "${bad_yaml}"
[[ "${code}" == "400" ]] || { echo "expected 400 got ${code}" >&2; exit 1; }

echo "[case] B-Queue-1: pending_resources when demand > cluster capacity"
big_yaml="$(mktemp)"
cat >"${big_yaml}" <<'YAML'
workload: "ppo"
submission_id: ""
code_path: "/private/common/code/verl/verl_repo"
model_id: "Qwen/Qwen2.5-0.5B-Instruct"
train_file: "/private/common/datasets/gsm8k/train.parquet"
val_file: "/private/common/datasets/gsm8k/test.parquet"
nnodes: 3
n_gpus_per_node: 4
total_epochs: 1
total_training_steps: 10
save_freq: 10
test_freq: -1
YAML
BIG_TID="$(api_user "${ALICE_TOKEN}" -H 'Content-Type: application/yaml' --data-binary @"${big_yaml}" "${API_ADDR}/api/v2/tasks" | jq -r .task_id)"
rm -f "${big_yaml}"
wait_state "${ALICE_TOKEN}" "${BIG_TID}" "PENDING_RESOURCES" 60

echo "[case] B-Cancel-1: cancel a running task (best-effort)"
long_yaml="$(mktemp)"
cat >"${long_yaml}" <<'YAML'
workload: "ppo"
submission_id: ""
code_path: "/private/common/code/verl/verl_repo"
model_id: "Qwen/Qwen2.5-0.5B-Instruct"
train_file: "/private/common/datasets/gsm8k/train.parquet"
val_file: "/private/common/datasets/gsm8k/test.parquet"
nnodes: 2
n_gpus_per_node: 4
total_epochs: 1
total_training_steps: 200
save_freq: 10
test_freq: -1
YAML
LONG_TID="$(api_user "${ALICE_TOKEN}" -H 'Content-Type: application/yaml' --data-binary @"${long_yaml}" "${API_ADDR}/api/v2/tasks" | jq -r .task_id)"
rm -f "${long_yaml}"

for i in $(seq 1 60); do
  st="$(api_user "${ALICE_TOKEN}" "${API_ADDR}/api/v2/tasks/${LONG_TID}" | jq -r .state)"
  echo "[case] wait running: ${LONG_TID} state=${st}"
  if [[ "${st}" == "RUNNING" ]]; then
    break
  fi
  if [[ "${st}" == "FAILED" || "${st}" == "SUCCEEDED" ]]; then
    break
  fi
  sleep 5
done

api_user "${ALICE_TOKEN}" -X POST "${API_ADDR}/api/v2/tasks/${LONG_TID}:cancel" | jq .
st="$(api_admin "${API_ADDR}/api/v2/tasks/${LONG_TID}" | jq -r .state)"
[[ "${st}" == "CANCELED" ]] || { echo "expected CANCELED got ${st}" >&2; exit 1; }

echo "[case] ===== all v2.5 e2e cases passed ====="
