#!/usr/bin/env bash
#
# Gemma 4 家族服务验证脚本（Gemma-4-31B / GEV-26B-Decide）
#
# 默认行为：
#   1. 如果 docker 可用，自动找正在运行的 gemma4-bi150 容器；
#   2. 从容器挂载/环境变量或 MODEL_HOST_DIR 推断模型目录、端口、served name；
#   3. 如果模型目录可读，调用 detect_model.py 自动判断变体；
#   4. 依次检查 /health、/v1/models、/v1/chat/completions；
#   5. 可选 --with-reasoning 和 --with-image。
#
# 常用：
#   ./verify.sh
#   ./verify.sh --model-dir /mnt/disk0/models/autotrust/GEV-26B-Decide --with-image
#   ./verify.sh --container gemma4-gemma4-31b --with-reasoning
#   ./verify.sh --base-url http://127.0.0.1:8095 --served-name gemma-4-31B-it
#
set -euo pipefail

cd "$(dirname "$0")"

MODEL_DIR="${MODEL_HOST_DIR:-}"
CONTAINER="${CONTAINER:-}"
PORT="${PORT:-${HOST_PORT:-}}"
SERVED_NAME="${SERVED_NAME:-}"
BASE_URL="${BASE_URL:-}"
TIMEOUT="${TIMEOUT:-300}"
WITH_IMAGE=0
WITH_REASONING=0
SKIP_CHAT=0
KIND="unknown"

usage() {
  cat >&2 <<'EOF'
用法: ./verify.sh [options]

  --model-dir PATH     模型目录（含 config.json），用于自动判断变体
  --container NAME     运行中的容器名；默认自动查找 gemma4 容器
  --port PORT          服务端口；默认从模型/容器推断
  --served-name NAME   OpenAI API 模型名；默认从模型/容器推断
  --base-url URL       直接指定服务地址，例如 http://127.0.0.1:8095
  --timeout SEC        等待和请求超时，默认 300
  --with-reasoning     额外跑一个简单推理题（期望 27）
  --with-image         额外跑图片识别（纯红图片，期望回答红色）
  --skip-chat          只检查 /health 和 /v1/models
  -h, --help           显示帮助

环境变量也可以设置：MODEL_HOST_DIR CONTAINER PORT SERVED_NAME BASE_URL TIMEOUT
EOF
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --model-dir) MODEL_DIR="${2:?--model-dir 需要路径}"; shift 2 ;;
    --container) CONTAINER="${2:?--container 需要名称}"; shift 2 ;;
    --port) PORT="${2:?--port 需要端口}"; shift 2 ;;
    --served-name) SERVED_NAME="${2:?--served-name 需要名称}"; shift 2 ;;
    --base-url) BASE_URL="${2:?--base-url 需要 URL}"; shift 2 ;;
    --timeout) TIMEOUT="${2:?--timeout 需要秒数}"; shift 2 ;;
    --with-image) WITH_IMAGE=1; shift ;;
    --with-reasoning) WITH_REASONING=1; shift ;;
    --skip-chat) SKIP_CHAT=1; shift ;;
    -h|--help) usage; exit 0 ;;
    *) echo "ERROR: 未知参数: $1" >&2; usage; exit 2 ;;
  esac
done

fail() { echo "FAIL: $*" >&2; exit 1; }
pass() { echo "PASS: $*"; }

command -v curl >/dev/null 2>&1 || fail "缺少 curl"
command -v python3 >/dev/null 2>&1 || fail "缺少 python3"

TMP="$(mktemp -d)"
trap 'rm -rf "${TMP}"' EXIT

# ---------------------------------------------------------------- container
if [[ -z "${CONTAINER}" ]] && command -v docker >/dev/null 2>&1; then
  CONTAINER="$(docker ps --format '{{.Names}} {{.Image}}' 2>/dev/null \
    | awk '/gemma4|gev-26b/ {print $1; exit}' || true)"
fi

container_env() {
  local name="$1" key="$2"
  docker inspect "${name}" --format '{{range .Config.Env}}{{println .}}{{end}}' 2>/dev/null \
    | awk -F= -v key="${key}" '$1 == key { sub(/^[^=]*=/, ""); print; exit }' || true
}

if [[ -n "${CONTAINER}" ]] && command -v docker >/dev/null 2>&1; then
  if [[ -z "${MODEL_DIR}" ]]; then
    MODEL_DIR="$(docker inspect "${CONTAINER}" --format '{{json .Mounts}}' 2>/dev/null \
      | python3 -c '
import json, sys
try:
    mounts = json.load(sys.stdin)
except Exception:
    sys.exit(0)
for m in mounts:
    if m.get("Destination") == "/models/gemma4":
        print(m.get("Source", ""))
        break
' || true)"
  fi
  # 宿主端口来自端口映射（-p），不是容器内的 PORT（现在恒为 80）
  if [[ -z "${PORT}" ]]; then
    PORT="$(docker inspect "${CONTAINER}" --format '{{json .HostConfig.PortBindings}}' 2>/dev/null \
      | python3 -c '
import json, sys
try:
    binds = json.load(sys.stdin) or {}
except Exception:
    sys.exit(0)
for container_port, mapping in binds.items():
    if mapping:
        print(mapping[0].get("HostPort", ""))
        break
' || true)"
  fi
  [[ -z "${SERVED_NAME}" ]] && SERVED_NAME="$(container_env "${CONTAINER}" SERVED_NAME)"
fi

# ---------------------------------------------------------------- model dir
GEMMA4_DEFAULT="/mnt/disk0/models/google/gemma-4-31B-it"
GEV_DEFAULT="/mnt/disk0/models/autotrust/GEV-26B-Decide"

# 没有 docker 时（例如在 ugly_hopper 内运行），按端口探测正在运行的服务。
if [[ -z "${CONTAINER}" && -z "${PORT}" && -z "${MODEL_DIR}" && -z "${BASE_URL}" ]]; then
  for candidate_port in 24396 8095 8094; do
    if curl -fsS --max-time 2 "http://127.0.0.1:${candidate_port}/health" >/dev/null 2>&1; then
      PORT="${candidate_port}"
      break
    fi
  done
fi

# 模型名现在恒为 llm（AGENTS.md 5.5），不能再靠它反推变体。
# 兼容旧用法：如果显式传了带变体的 served name，仍然尝试反推目录。
if [[ -z "${MODEL_DIR}" && -n "${SERVED_NAME}" && "${SERVED_NAME}" != "llm" ]]; then
  case "${SERVED_NAME}" in
    *GEV*|*gev*)
      [[ -f "${GEV_DEFAULT}/config.json" ]] && MODEL_DIR="${GEV_DEFAULT}"
      ;;
    *gemma*|*Gemma*)
      [[ -f "${GEMMA4_DEFAULT}/config.json" ]] && MODEL_DIR="${GEMMA4_DEFAULT}"
      ;;
  esac
fi

# 端口也不再区分变体（宿主端口默认统一 24396），只按目录存在性取默认值。
if [[ -z "${MODEL_DIR}" && -z "${BASE_URL}" ]]; then
  if [[ -f "${GEMMA4_DEFAULT}/config.json" ]]; then
    MODEL_DIR="${GEMMA4_DEFAULT}"
  elif [[ -f "${GEV_DEFAULT}/config.json" ]]; then
    MODEL_DIR="${GEV_DEFAULT}"
  fi
fi

DETECT_JSON=""
if [[ -n "${MODEL_DIR}" ]]; then
  if [[ -f "${MODEL_DIR}/config.json" ]]; then
    DETECT_JSON="$(python3 detect_model.py "${MODEL_DIR}" --json)" || {
      echo "ERROR: 无法识别模型: ${MODEL_DIR}" >&2
      exit 1
    }
  elif [[ -n "${CONTAINER}" ]] && command -v docker >/dev/null 2>&1; then
    # 宿主机路径不可读时，退回到容器内检测
    DETECT_JSON="$(docker exec "${CONTAINER}" python3 /opt/gemma4/detect_model.py /models/gemma4 --json 2>/dev/null || true)"
  else
    fail "模型目录缺少 config.json: ${MODEL_DIR}"
  fi
fi

if [[ -n "${DETECT_JSON}" ]]; then
  KIND="$(python3 -c 'import json,sys; print(json.load(sys.stdin)["kind"])' <<< "${DETECT_JSON}")"
  DETECTED_PORT="$(python3 -c 'import json,sys; print(json.load(sys.stdin)["port"])' <<< "${DETECT_JSON}")"
  # DETECTED_PORT / served_name 只是变体信息，不作为实际端口和模型名：
  # 容器内恒 80，宿主端口默认 24396，模型名恒 llm（AGENTS.md 5.2 / 5.4 / 5.5）。
  DETECTED_HOST_PORT_HINT="${DETECTED_PORT:-}"
fi

PORT="${PORT:-24396}"
SERVED_NAME="${SERVED_NAME:-llm}"
BASE_URL="${BASE_URL:-http://127.0.0.1:${PORT}}"
BASE_URL="${BASE_URL%/}"

echo "=== Gemma 4 family verification ==="
echo "container : ${CONTAINER:-<none>}"
echo "model dir : ${MODEL_DIR:-<not detected>}"
echo "kind      : ${KIND}"
echo "base URL  : ${BASE_URL}"
echo "model name: ${SERVED_NAME}"
echo

# ---------------------------------------------------------------- health
# vLLM 的 /health 可能返回 200 但 body 为空，所以只看 curl 退出码。
echo "==> waiting for /health (up to ${TIMEOUT}s)"
HEALTH_OK=0
DEADLINE=$((SECONDS + TIMEOUT))
while (( SECONDS < DEADLINE )); do
  if curl -fsS --max-time 5 "${BASE_URL}/health" > "${TMP}/health.txt" 2>/dev/null; then
    HEALTH_OK=1
    break
  fi
  sleep 2
done
if [[ "${HEALTH_OK}" != "1" ]]; then
  fail "health check 超时: ${BASE_URL}/health"
fi
HEALTH="$(cat "${TMP}/health.txt" 2>/dev/null || true)"
pass "health: ${HEALTH:-OK}"

# ---------------------------------------------------------------- models
echo "==> /v1/models"
curl -fsS --max-time "${TIMEOUT}" "${BASE_URL}/v1/models" > "${TMP}/models.json" \
  || fail "GET /v1/models 失败"
python3 - "${TMP}/models.json" "${SERVED_NAME}" <<'PY'
import json
import sys

data = json.load(open(sys.argv[1]))
expected = sys.argv[2]
ids = [m.get("id") for m in data.get("data", [])]
if expected not in ids:
    print(f"FAIL: /v1/models 中没有 {expected!r}; 实际: {ids}", file=sys.stderr)
    sys.exit(1)
for m in data.get("data", []):
    if m.get("id") == expected:
        print(f"PASS: model={expected}, max_model_len={m.get('max_model_len')}")
        break
PY

if [[ "${SKIP_CHAT}" == "1" ]]; then
  echo
  echo "All checks passed (chat skipped)."
  exit 0
fi

# ---------------------------------------------------------------- chat
echo "==> /v1/chat/completions"
python3 - "${SERVED_NAME}" > "${TMP}/chat.json" <<'PY'
import json
import sys

print(json.dumps({
    "model": sys.argv[1],
    "messages": [{"role": "user", "content": "你好，请用一句中文介绍你自己。"}],
    "max_tokens": 64,
    "temperature": 0,
}, ensure_ascii=False))
PY
curl -fsS --max-time "${TIMEOUT}" \
  -H 'Content-Type: application/json' \
  --data-binary @"${TMP}/chat.json" \
  "${BASE_URL}/v1/chat/completions" > "${TMP}/chat_resp.json" \
  || fail "chat completion 失败"
python3 - "${TMP}/chat_resp.json" <<'PY'
import json
import sys

data = json.load(open(sys.argv[1]))
try:
    content = data["choices"][0]["message"]["content"]
except (KeyError, IndexError):
    print(f"FAIL: unexpected chat response: {data}", file=sys.stderr)
    sys.exit(1)
if isinstance(content, list):
    content = " ".join(
        part.get("text", "") for part in content if isinstance(part, dict)
    )
content = (content or "").strip()
if not content:
    print("FAIL: chat response content is empty", file=sys.stderr)
    sys.exit(1)
print(f"PASS: chat: {content[:200]}")
PY

# ---------------------------------------------------------------- reasoning
if [[ "${WITH_REASONING}" == "1" ]]; then
  echo "==> reasoning"
  python3 - "${SERVED_NAME}" > "${TMP}/reason.json" <<'PY'
import json
import sys

print(json.dumps({
    "model": sys.argv[1],
    "messages": [{"role": "user", "content": "3 个人 3 天喝 3 桶水，9 个人 9 天喝几桶水？只给最终数字。"}],
    "max_tokens": 16,
    "temperature": 0,
}, ensure_ascii=False))
PY
  curl -fsS --max-time "${TIMEOUT}" \
    -H 'Content-Type: application/json' \
    --data-binary @"${TMP}/reason.json" \
    "${BASE_URL}/v1/chat/completions" > "${TMP}/reason_resp.json" \
    || fail "reasoning 请求失败"
  python3 - "${TMP}/reason_resp.json" <<'PY'
import json
import sys

data = json.load(open(sys.argv[1]))
content = data["choices"][0]["message"].get("content", "")
if isinstance(content, list):
    content = " ".join(
        part.get("text", "") for part in content if isinstance(part, dict)
    )
content = (content or "").strip()
if "27" not in content:
    print(f"FAIL: expected 27, got: {content!r}", file=sys.stderr)
    sys.exit(1)
print(f"PASS: reasoning: {content[:100]}")
PY
fi

# ---------------------------------------------------------------- image
if [[ "${WITH_IMAGE}" == "1" ]]; then
  echo "==> image (336x336 red PNG)"
  python3 - "${TMP}/red.b64" <<'PY'
import base64
import struct
import sys
import zlib

w = h = 336
row = bytes((220, 20, 20)) * w
raw = b"".join(b"\x00" + row for _ in range(h))

def chunk(kind, data):
    return (
        struct.pack(">I", len(data))
        + kind
        + data
        + struct.pack(">I", zlib.crc32(kind + data) & 0xFFFFFFFF)
    )

png = (
    b"\x89PNG\r\n\x1a\n"
    + chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0))
    + chunk(b"IDAT", zlib.compress(raw, 9))
    + chunk(b"IEND", b"")
)
open(sys.argv[1], "w").write(base64.b64encode(png).decode())
PY
  python3 - "${SERVED_NAME}" "${TMP}/red.b64" > "${TMP}/image.json" <<'PY'
import json
import sys

b64 = open(sys.argv[2]).read().strip()
payload = {
    "model": sys.argv[1],
    "messages": [{
        "role": "user",
        "content": [
            {"type": "text", "text": "这张图片主要是什么颜色？只回答颜色。"},
            {"type": "image_url", "image_url": {"url": "data:image/png;base64," + b64}},
        ],
    }],
    "max_tokens": 32,
    "temperature": 0,
}
print(json.dumps(payload, ensure_ascii=False))
PY
  curl -fsS --max-time "${TIMEOUT}" \
    -H 'Content-Type: application/json' \
    --data-binary @"${TMP}/image.json" \
    "${BASE_URL}/v1/chat/completions" > "${TMP}/image_resp.json" \
    || fail "image 请求失败"
  python3 - "${TMP}/image_resp.json" <<'PY'
import json
import sys

data = json.load(open(sys.argv[1]))
content = data["choices"][0]["message"].get("content", "")
if isinstance(content, list):
    content = " ".join(
        part.get("text", "") for part in content if isinstance(part, dict)
    )
content = (content or "").strip()
if not content:
    print("FAIL: image response content is empty", file=sys.stderr)
    sys.exit(1)
lower = content.lower()
if "红" not in lower and "red" not in lower:
    print(f"FAIL: expected red, got: {content[:200]}", file=sys.stderr)
    sys.exit(1)
print(f"PASS: image: {content[:200]}")
PY
fi

echo
echo "All checks passed for ${KIND} at ${BASE_URL}."