#!/usr/bin/env bash
#
# verify-image: smoke test a built MesaStack image.
#
# Checks, in order:
#   1. static  - gpustack version, API_KEY_PREFIX, UI assets carry our brand
#                and no leftover upstream brand in index.html
#   2. runtime - boots the server, confirms / serves the MesaStack title,
#                logs in and creates an API key to prove the mesastack_ prefix
#                is live end to end, and that old gpustack_ keys are rejected
#
# Usage:
#   scripts/verify-image                     # verify ${NAMESPACE}/${REPOSITORY}:<workspace version>
#   IMAGE=mesastack/gpustack:v2.2.1 scripts/verify-image
#   scripts/verify-image --static-only       # skip booting a container
#
set -o errexit
set -o nounset
set -o pipefail

source "$(dirname "${BASH_SOURCE[0]}")/lib/common.sh"

STATIC_ONLY=0
while [[ $# -gt 0 ]]; do
  case "${1}" in
    --static-only) STATIC_ONLY=1 ;;
    -h|--help)     sed -n '2,18p' "${0}"; exit 0 ;;
    *)             fail "unknown argument: ${1}" ;;
  esac
  shift
done

require_cmd docker curl python3

if [[ -z "${IMAGE:-}" ]]; then
  require_repos
  VERSION="$(resolve_workspace_version)"
  IMAGE="${NAMESPACE}/${REPOSITORY}:${TAG:-${VERSION}}"
fi

SMOKE_PORT="${SMOKE_PORT:-18080}"
CONTAINER="${CONTAINER:-mesastack-verify}"
SITE_PACKAGES="/usr/local/lib/python3.11/dist-packages"
BRAND="${BRAND:-MesaStack}"
EXPECTED_PREFIX="${EXPECTED_PREFIX:-mesastack}"

docker image inspect "${IMAGE}" >/dev/null 2>&1 || fail "image not found locally: ${IMAGE}"
log "verifying ${IMAGE}"

failed=0
check() {
  local label="${1}"; shift
  if "$@" >/dev/null 2>&1; then
    echo -e "  \033[1;32mPASS\033[0m ${label}"
  else
    echo -e "  \033[1;31mFAIL\033[0m ${label}"
    failed=1
  fi
}

# ---------- 1. Static checks ----------
step "Static checks"
docker run --rm --entrypoint bash "${IMAGE}" -c "
set -e
echo \"version: \$(gpustack version)\"
echo \"prefix : \$(python3 -c 'from gpustack.security import API_KEY_PREFIX; print(API_KEY_PREFIX)')\"
" || fail "image failed to report version/prefix"

check "API_KEY_PREFIX == ${EXPECTED_PREFIX}" \
  docker run --rm --entrypoint python3 "${IMAGE}" -c \
    "from gpustack.security import API_KEY_PREFIX as p; assert p == '${EXPECTED_PREFIX}', p"

check "old gpustack_ keys rejected" \
  docker run --rm --entrypoint python3 "${IMAGE}" -c \
    "from gpustack.security import is_valid_format as v; assert not v('gpustack_aaaa_bbbb')[0]"

check "masked value uses ${EXPECTED_PREFIX}_" \
  docker run --rm --entrypoint python3 "${IMAGE}" -c \
    "from gpustack.utils.api_keys import get_masked_api_key_value as m; assert m('abcd1234') == '${EXPECTED_PREFIX}_abcd***', m('abcd1234')"

check "UI index.html present" \
  docker run --rm --entrypoint test "${IMAGE}" -f "${SITE_PACKAGES}/gpustack/ui/index.html"

check "UI carries ${BRAND} brand" \
  docker run --rm --entrypoint grep "${IMAGE}" -q "${BRAND}" "${SITE_PACKAGES}/gpustack/ui/index.html"

check "no upstream GPUStack brand in UI index.html" \
  docker run --rm --entrypoint bash "${IMAGE}" -c \
    "! grep -q 'GPUStack' ${SITE_PACKAGES}/gpustack/ui/index.html"

if [[ "${STATIC_ONLY}" == "1" ]]; then
  step "Result"
  [[ "${failed}" == "0" ]] && { log "static checks passed"; exit 0; } || fail "static checks failed"
fi

# ---------- 2. Runtime checks ----------
step "Runtime checks (booting ${CONTAINER} on :${SMOKE_PORT})"

cleanup() {
  docker rm -f "${CONTAINER}" >/dev/null 2>&1 || true
  rm -f "${COOKIES:-/dev/null}" 2>/dev/null || true
}
trap cleanup EXIT

docker rm -f "${CONTAINER}" >/dev/null 2>&1 || true
docker run -d --name "${CONTAINER}" -p "${SMOKE_PORT}:80" "${IMAGE}" --disable-worker >/dev/null \
  || fail "failed to start container"

BASE="http://127.0.0.1:${SMOKE_PORT}"
log "waiting for server to become ready"
ready=0
for _ in $(seq 1 60); do
  if curl -sS -m 5 -o /dev/null "${BASE}/" 2>/dev/null; then ready=1; break; fi
  if ! docker ps --filter "name=${CONTAINER}" --format '{{.ID}}' | grep -q .; then
    docker logs "${CONTAINER}" 2>&1 | tail -20
    fail "container exited during startup"
  fi
  sleep 5
done
[[ "${ready}" == "1" ]] || { docker logs "${CONTAINER}" 2>&1 | tail -20; fail "server did not become ready"; }

check "/ serves ${BRAND} title" \
  bash -c "curl -sS -m 15 '${BASE}/' | grep -q '<title>${BRAND}</title>'"

COOKIES="$(mktemp)"
PASSWORD="$(docker exec "${CONTAINER}" cat /var/lib/gpustack/initial_admin_password 2>/dev/null || echo "")"
if [[ -z "${PASSWORD}" ]]; then
  warn "bootstrap password not found; skipping API key check"
else
  login_status="$(curl -sS -c "${COOKIES}" -m 20 -o /dev/null -w '%{http_code}' \
    -X POST "${BASE}/auth/login" \
    -H 'Content-Type: application/x-www-form-urlencoded' \
    --data-urlencode "username=admin" \
    --data-urlencode "password=${PASSWORD}" || echo "000")"
  check "admin login succeeds" test "${login_status}" = "200"

  if [[ "${login_status}" == "200" ]]; then
    key_value="$(curl -sS -b "${COOKIES}" -m 20 -X POST "${BASE}/v2/api-keys" \
      -H 'Content-Type: application/json' \
      -d '{"name":"verify-image","description":"prefix check"}' \
      | python3 -c 'import json,sys; print(json.load(sys.stdin).get("value",""))' 2>/dev/null || echo "")"
    log "generated key: ${key_value%%_*}_<redacted>"
    check "generated key uses ${EXPECTED_PREFIX}_ prefix" \
      bash -c "[[ '${key_value}' == ${EXPECTED_PREFIX}_* ]]"
  fi
fi

step "Result"
if [[ "${failed}" == "0" ]]; then
  log "all checks passed for ${IMAGE}"
else
  fail "one or more checks failed for ${IMAGE}"
fi