#!/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:-}"
USER_ID="${USER_ID:-alice}"
EXPECTED_RAY_NODES="${EXPECTED_RAY_NODES:-3}" # head + 2 workers

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
}

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

openai_wait_ready() {
  local base_url="$1"
  local tries="${2:-120}"
  for i in $(seq 1 "${tries}"); do
    if curl -sS -m 2 "${base_url}/models" >/dev/null 2>&1; then
      echo "[host] openai_ready: ${base_url}"
      return 0
    fi
    echo "[host] waiting openai... (${i}/${tries})"
    sleep 2
  done
  echo "ERROR: openai not ready: ${base_url}" >&2
  return 1
}

wait_model_state() {
  local token="$1"
  local model_key="$2"
  local want="$3"
  local tries="${4:-120}"
  for i in $(seq 1 "${tries}"); do
    local body state
    body="$(curl -sS -H "Authorization: Bearer ${token}" "${API_ADDR}/api/v2/serve/models/${model_key}")"
    state="$(printf '%s' "${body}" | python3 -c 'import sys,json; print(json.load(sys.stdin)["model"]["state"])' 2>/dev/null || true)"
    echo "[host] model ${model_key}: ${state}"
    if [[ "${state}" == "${want}" ]]; then
      return 0
    fi
    if [[ "${state}" == "FAILED" ]]; then
      echo "[host] model failed; detail:" >&2
      printf '%s\n' "${body}" | python3 -m json.tool >&2 || true
      return 1
    fi
    sleep 2
  done
  echo "ERROR: model not in state ${want} after timeout" >&2
  return 1
}

model_openai_base_url() {
  local token="$1"
  local model_key="$2"
  curl -sS -H "Authorization: Bearer ${token}" "${API_ADDR}/api/v2/serve/models/${model_key}" \
    | python3 -c 'import sys,json; print(json.load(sys.stdin)["model"]["endpoint"]["openai_base_url"])'
}

serve_replica_fingerprint_json() {
  local model_key="$1"
  # Best-effort: list ServeReplica actors for this model app and return a stable JSON signature.
  dexec "${HEAD_CONTAINER}" python3 - "${model_key}" <<'PY' | tail -n 1
import json
import sys

import ray

model_key = sys.argv[1]
ray.init(address="auto", ignore_reinit_error=True, log_to_driver=False)
try:
    from ray.util.state import list_actors
except Exception:
    from ray.experimental.state.api import list_actors  # type: ignore

actors = list_actors(limit=10000)
want = f"ServeReplica:{model_key}:"
out = []
for a in actors or []:
    if isinstance(a, dict):
        cls = a.get("class_name")
        st = a.get("state")
        actor_id = a.get("actor_id")
        node_id = a.get("node_id")
        pid = a.get("pid")
    else:
        cls = getattr(a, "class_name", None)
        st = getattr(a, "state", None)
        actor_id = getattr(a, "actor_id", None)
        node_id = getattr(a, "node_id", None)
        pid = getattr(a, "pid", None)

    cls_s = str(cls or "")
    if not cls_s.startswith(want):
        continue
    if str(st or "") and str(st) != "ALIVE":
        continue
    out.append(
        {
            "actor_id": actor_id,
            "node_id": node_id,
            "class_name": cls_s,
            "pid": pid,
        }
    )
out = sorted(out, key=lambda x: (x.get("class_name") or "", x.get("actor_id") or ""))
print(json.dumps(out, sort_keys=True))
PY
}

echo "[host] ===== run_all_v38_serving.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 (Ray + SFTPGo + W&B)"
# Default: don't rebuild if image already exists (01_up.sh will auto-build if missing).
BUILD="${BUILD:-0}" "${SCRIPT_DIR}/01_up.sh"

echo "[host] wait ray ready"
ray_wait_ready 60
ray_wait_nodes "${EXPECTED_RAY_NODES}" 120

echo "[host] prepare data/model (best-effort; uses shared caches)"
"${SCRIPT_DIR}/30_prepare_data_and_model.sh" || true

echo "[host] start api"
CONFIG_IN_CONTAINER="${CONFIG_IN_CONTAINER}" MVP_INTERNAL_TOKEN="${ADMIN_TOKEN}" "${SCRIPT_DIR}/60_start_api.sh"
api_wait_ready 60

echo "[host] create user (idempotent)"
api_curl_admin -X POST "${API_ADDR}/api/v2/users" -H "Content-Type: application/json" --data-binary "{\"user_id\":\"${USER_ID}\"}" >/dev/null || true

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] resolve local model snapshot path (offline)"
LOCAL_MODEL_PATH="$(dexec "${HEAD_CONTAINER}" bash -lc "python3 -c \"import os; from huggingface_hub import snapshot_download; os.environ.setdefault('HF_HOME','/private/hf'); print(snapshot_download(repo_id='Qwen/Qwen2.5-0.5B-Instruct', local_files_only=True))\" " | tail -n 1)"
if [[ -z "${LOCAL_MODEL_PATH}" || "${LOCAL_MODEL_PATH}" != /* ]]; then
  echo "ERROR: failed to resolve LOCAL_MODEL_PATH: ${LOCAL_MODEL_PATH}" >&2
  exit 1
fi
echo "[host] local_model_path: ${LOCAL_MODEL_PATH}"

echo "[host] submit serving model A via API"
SERVE_SPEC_A=$'model_id: qwen-0.5b-a\nmodel_source: '"${LOCAL_MODEL_PATH}"$'\nnum_replicas: 1\ngpus_per_replica: 1\n'
CREATE_RESP_A="$(curl -sS -H "Authorization: Bearer ${USER_TOKEN}" -H "Content-Type: application/yaml" --data-binary "${SERVE_SPEC_A}" "${API_ADDR}/api/v2/serve/models")"
echo "[host] create_model_a_resp: ${CREATE_RESP_A}"
MODEL_KEY_A="$(printf '%s' "${CREATE_RESP_A}" | python3 -c 'import sys,json; print(json.load(sys.stdin)["model_key"])')"

echo "[host] wait model A RUNNING"
wait_model_state "${USER_TOKEN}" "${MODEL_KEY_A}" "RUNNING" 300
OPENAI_BASE_A="$(model_openai_base_url "${USER_TOKEN}" "${MODEL_KEY_A}")"
echo "[host] openai_base_a: ${OPENAI_BASE_A}"
openai_wait_ready "${OPENAI_BASE_A}" 120

echo "[host] verify A /models contains model"
MODEL_ID_A="$(
  curl -sS "${OPENAI_BASE_A}/models" \
    | python3 -c 'import sys,json; obj=json.load(sys.stdin); print("\n".join([m.get("id","") for m in obj.get("data",[]) if isinstance(m,dict)]))' \
    | grep -E "^${USER_ID}-[0-9]{12}-qwen-0\\.5b-a$" \
    | head -n1 \
    || true
)"
if [[ -z "${MODEL_ID_A}" ]]; then
  echo "ERROR: model id A not found in /models" >&2
  curl -sS "${OPENAI_BASE_A}/models" | python3 -m json.tool >&2 || true
  exit 1
fi
echo "[host] model_id_a: ${MODEL_ID_A}"

echo "[host] chat completion A (best-effort)"
CHAT_RESP_A="$(curl -sS -H "Content-Type: application/json" -X POST "${OPENAI_BASE_A}/chat/completions" --data-binary "{\"model\":\"${MODEL_ID_A}\",\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}],\"max_tokens\":16,\"stream\":false}")"
printf '%s\n' "${CHAT_RESP_A}" | python3 -m json.tool >/dev/null 2>&1 || {
  echo "ERROR: invalid chat response A" >&2
  printf '%s\n' "${CHAT_RESP_A}" >&2
  exit 1
}
echo "[host] chat_a_ok"

echo "[host] fingerprint A (serve replica actors)"
FP_A_1="$(serve_replica_fingerprint_json "${MODEL_KEY_A}")"
echo "[host] fp_a_1: ${FP_A_1}"
if [[ -z "${FP_A_1}" || "${FP_A_1}" == "[]" ]]; then
  echo "ERROR: failed to fingerprint A replicas (empty)" >&2
  exit 1
fi

echo "[host] submit serving model B via API"
SERVE_SPEC_B=$'model_id: qwen-0.5b-b\nmodel_source: '"${LOCAL_MODEL_PATH}"$'\nnum_replicas: 1\ngpus_per_replica: 1\n'
CREATE_RESP_B="$(curl -sS -H "Authorization: Bearer ${USER_TOKEN}" -H "Content-Type: application/yaml" --data-binary "${SERVE_SPEC_B}" "${API_ADDR}/api/v2/serve/models")"
echo "[host] create_model_b_resp: ${CREATE_RESP_B}"
MODEL_KEY_B="$(printf '%s' "${CREATE_RESP_B}" | python3 -c 'import sys,json; print(json.load(sys.stdin)["model_key"])')"

echo "[host] wait model B RUNNING"
wait_model_state "${USER_TOKEN}" "${MODEL_KEY_B}" "RUNNING" 300
OPENAI_BASE_B="$(model_openai_base_url "${USER_TOKEN}" "${MODEL_KEY_B}")"
echo "[host] openai_base_b: ${OPENAI_BASE_B}"
openai_wait_ready "${OPENAI_BASE_B}" 120

echo "[host] verify isolation: A replica fingerprint unchanged after creating B"
FP_A_2="$(serve_replica_fingerprint_json "${MODEL_KEY_A}")"
echo "[host] fp_a_2: ${FP_A_2}"
if [[ "${FP_A_2}" != "${FP_A_1}" ]]; then
  echo "ERROR: A replicas changed after creating B (unexpected redeploy)" >&2
  exit 1
fi

echo "[host] delete model B"
curl -sS -H "Authorization: Bearer ${USER_TOKEN}" -X DELETE "${API_ADDR}/api/v2/serve/models/${MODEL_KEY_B}" >/dev/null
wait_model_state "${USER_TOKEN}" "${MODEL_KEY_B}" "DELETED" 300

echo "[host] verify isolation: A replica fingerprint unchanged after deleting B"
FP_A_3="$(serve_replica_fingerprint_json "${MODEL_KEY_A}")"
echo "[host] fp_a_3: ${FP_A_3}"
if [[ "${FP_A_3}" != "${FP_A_1}" ]]; then
  echo "ERROR: A replicas changed after deleting B (unexpected redeploy)" >&2
  exit 1
fi

echo "[host] delete model A"
curl -sS -H "Authorization: Bearer ${USER_TOKEN}" -X DELETE "${API_ADDR}/api/v2/serve/models/${MODEL_KEY_A}" >/dev/null
wait_model_state "${USER_TOKEN}" "${MODEL_KEY_A}" "DELETED" 300

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