#!/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 v3.0:
# - Start Ray (head + stateless workers) + SFTPGo (compose)
# - Start API server with v3.0 config
# - Create user (API triggers SFTPGo user create) and return one-time SFTP password
# - (Optional) verify SFTP login
# - Submit PPO/GRPO/SFT referencing user dataset paths 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}"
RESET_SFTPGO="${RESET_SFTPGO:-0}"
EXPECTED_RAY_NODES="${EXPECTED_RAY_NODES:-3}" # head + 2 workers
CLUSTER_NAME="${CLUSTER_NAME:-argus-ray}"

CONFIG_IN_CONTAINER="${CONFIG_IN_CONTAINER:-/workspace/mvp/configs/dev.yaml}"
SFTPGO_ADMIN_PASSWORD="${SFTPGO_ADMIN_PASSWORD:-my-dev-sftpgo-admin}"
export SFTPGO_ADMIN_PASSWORD

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
}

sftpgo_wait_ready() {
  local tries="${1:-60}"
  local url="${2:-http://127.0.0.1:8081/api/v2/token}"
  for i in $(seq 1 "${tries}"); do
    # SFTPGo admin endpoints require auth; readiness means "HTTP reachable and can issue token".
    if curl -sS -m 2 -u "admin:${SFTPGO_ADMIN_PASSWORD}" "${url}" >/dev/null 2>&1; then
      echo "[host] sftpgo_ready: ${url} (token ok)"
      return 0
    fi
    echo "[host] waiting sftpgo... (${i}/${tries})"
    sleep 2
  done
  echo "ERROR: sftpgo not ready: ${url}" >&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_inline() {
  local token="$1"
  local yaml_body="$2"
  local resp
  resp="$(curl -sS -H "Authorization: Bearer ${token}" -H "Content-Type: application/yaml" --data-binary "${yaml_body}" "${API_ADDR}/api/v2/tasks")"
  echo "[host] submit_resp: ${resp}" >&2
  printf '%s' "${resp}" | python3 -c 'import sys,json; o=json.load(sys.stdin); tid=o.get("task_id",""); print(tid) if tid else sys.exit(42)'
  local rc=$?
  if [[ "${rc}" != "0" ]]; then
    echo "ERROR: submit failed: ${resp}" >&2
    return 1
  fi
}

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_v30_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

if [[ "${RESET_SFTPGO}" == "1" ]]; then
  echo "[host] reset sftpgo metadata dir (best-effort, via helper container)"
  SFTPGO_META_DIR="${ROOT_DIR}/../../shared/common/sftpgo"
  mkdir -p "${SFTPGO_META_DIR}"
  docker run --rm --entrypoint sh -u 0:0 -v "${SFTPGO_META_DIR}:/mnt" drakkan/sftpgo:latest -lc "rm -rf /mnt/* || true"
fi

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

echo "[host] wait ray head ready"
ray_wait_ready 60

echo "[host] wait sftpgo ready"
sftpgo_wait_ready 60 "http://127.0.0.1:8081/api/v2/token"

echo "[host] render config with SFTPGo container IP (work around docker DNS issues)"
SFTPGO_IP="$(docker inspect -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' argus-sftpgo)"
RENDERED_CFG_HOST_PATH="/tmp/dev_rendered.yaml"
sed -E \
  -e "s#^(\\s*admin_api_base:) .*#\\1 \"http://${SFTPGO_IP}:8080/api/v2\"#g" \
  -e "/^  sftpgo:/,/^  retention:/ s#^(\\s*host:) .*#\\1 \"127.0.0.1\"#g" \
  "${ROOT_DIR}/configs/dev.yaml" >"${RENDERED_CFG_HOST_PATH}"
docker cp "${RENDERED_CFG_HOST_PATH}" "${HEAD_CONTAINER}:/tmp/dev_rendered.yaml"
CONFIG_IN_CONTAINER="/tmp/dev_rendered.yaml"

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"
ray_wait_nodes "${EXPECTED_RAY_NODES}" 120

echo "[host] prepare data/model (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 + sftpgo admin password via env)"
MVP_INTERNAL_TOKEN="${ADMIN_TOKEN}" CONFIG_IN_CONTAINER="${CONFIG_IN_CONTAINER}" SFTPGO_ADMIN_PASSWORD="${SFTPGO_ADMIN_PASSWORD}" "${SCRIPT_DIR}/60_start_api.sh"
api_wait_ready 60

echo "[host] create user (expect SFTP one-time password in response)"
create_resp="$(api_curl_admin -H "Content-Type: application/json" -d "{\"user_id\":\"${USER_ID}\"}" "${API_ADDR}/api/v2/users")"
echo "[host] create_user_resp: ${create_resp}"
USER_SFTP_PASSWORD="$(printf '%s' "${create_resp}" | python3 -c 'import sys,json; o=json.load(sys.stdin); print((o.get("sftp") or {}).get("password") or "")')"
if [[ -z "${USER_SFTP_PASSWORD}" ]]; then
  echo "ERROR: expected sftp.password in create user response (is data.sftpgo.enabled=true?)" >&2
  exit 1
fi

echo "[host] issue user token"
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}"

echo "[host] (optional) verify SFTP login (best-effort)"
if command -v sshpass >/dev/null 2>&1 && command -v sftp >/dev/null 2>&1; then
  tmp_batch="$(mktemp)"
  cat >"${tmp_batch}" <<EOF
pwd
ls
EOF
  # NOTE: this just proves auth works; dataset upload can be done via SFTP manually later.
  sshpass -p "${USER_SFTP_PASSWORD}" sftp -o StrictHostKeyChecking=no -P 2022 "${USER_ID}@127.0.0.1" -b "${tmp_batch}" >/dev/null 2>&1 || true
  rm -f "${tmp_batch}" || true
else
  echo "[host] skip: sshpass/sftp not found; you can test manually with: sftp -P 2022 ${USER_ID}@<host>"
fi

echo "[host] ensure user dataset paths exist (copy from common if needed; best-effort)"
echo "[host] copy dataset into user datasets path inside head container (avoid host permission issues)"
dexec "${HEAD_CONTAINER}" bash -lc "set -euo pipefail; \
  mkdir -p '/private/users/${USER_ID}/datasets/gsm8k' '/private/users/${USER_ID}/datasets/gsm8k_sft'; \
  (cp -f /private/common/datasets/gsm8k/train.parquet '/private/users/${USER_ID}/datasets/gsm8k/train.parquet' 2>/dev/null || cp -f /private/datasets/gsm8k/train.parquet '/private/users/${USER_ID}/datasets/gsm8k/train.parquet' 2>/dev/null || true); \
  (cp -f /private/common/datasets/gsm8k/test.parquet  '/private/users/${USER_ID}/datasets/gsm8k/test.parquet'  2>/dev/null || cp -f /private/datasets/gsm8k/test.parquet  '/private/users/${USER_ID}/datasets/gsm8k/test.parquet'  2>/dev/null || true); \
  (cp -f /private/common/datasets/gsm8k_sft/train.parquet '/private/users/${USER_ID}/datasets/gsm8k_sft/train.parquet' 2>/dev/null || cp -f /private/datasets/gsm8k_sft/train.parquet '/private/users/${USER_ID}/datasets/gsm8k_sft/train.parquet' 2>/dev/null || true); \
  (cp -f /private/common/datasets/gsm8k_sft/test.parquet  '/private/users/${USER_ID}/datasets/gsm8k_sft/test.parquet'  2>/dev/null || cp -f /private/datasets/gsm8k_sft/test.parquet  '/private/users/${USER_ID}/datasets/gsm8k_sft/test.parquet'  2>/dev/null || true)"

LOCAL_MODEL_PATH="Qwen/Qwen2.5-0.5B-Instruct"
echo "[host] submit PPO/GRPO/SFT via API using user dataset paths"
PPO_TASK_ID="$(submit_taskspec_inline "${USER_TOKEN}" $'workload: ppo\nnnodes: 2\nn_gpus_per_node: 4\ncode_path: /private/common/code/verl/verl_repo\ntrain_file: /private/users/'"${USER_ID}"$'/datasets/gsm8k/train.parquet\nval_file: /private/users/'"${USER_ID}"$'/datasets/gsm8k/test.parquet\nmodel_id: '"${LOCAL_MODEL_PATH}"$'\ntotal_epochs: 1\ntotal_training_steps: 10\nsave_freq: 10\n')"
GRPO_TASK_ID="$(submit_taskspec_inline "${USER_TOKEN}" $'workload: grpo\nnnodes: 2\nn_gpus_per_node: 4\ncode_path: /private/common/code/verl/verl_repo\ntrain_file: /private/users/'"${USER_ID}"$'/datasets/gsm8k/train.parquet\nval_file: /private/users/'"${USER_ID}"$'/datasets/gsm8k/test.parquet\nmodel_id: '"${LOCAL_MODEL_PATH}"$'\ntotal_epochs: 1\ntotal_training_steps: 10\nsave_freq: 10\n')"
SFT_TASK_ID="$(submit_taskspec_inline "${USER_TOKEN}" $'workload: sft\nnnodes: 1\nn_gpus_per_node: 1\ncode_path: /private/common/code/verl/verl_repo\ntrain_file: /private/users/'"${USER_ID}"$'/datasets/gsm8k_sft/train.parquet\nval_file: /private/users/'"${USER_ID}"$'/datasets/gsm8k_sft/test.parquet\nmodel_id: '"${LOCAL_MODEL_PATH}"$'\ntotal_epochs: 1\ntotal_training_steps: 10\nsave_freq: 10\n')"

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_v30_api.sh done ====="
