#!/usr/bin/env bash
# A30 一键：下载 MiniCPM-V 2.6 int4 量化模型 + 启动推理服务
# 用法: chmod +x run_minicpmv_a30.sh && ./run_minicpmv_a30.sh <download|start|stop|restart|logs|test>
set -euo pipefail

VENV_DIR="/data/volc_minicpm/venv"
MODEL_DIR="/data/volc_minicpm/models/MiniCPM-V-2_6-int4"
WORKDIR="$(cd "$(dirname "$0")" && pwd)"
SERVER_PY="${WORKDIR}/minicpmv_server.py"
LOG_DIR="/data/volc_minicpm/logs"
PID_FILE="/data/volc_minicpm/minicpmv.pid"
PORT="${PORT:-30000}"
HOST="${HOST:-0.0.0.0}"
DTYPE="${DTYPE:-bf16}"

mkdir -p "$LOG_DIR"

ACT() {
    if [[ -f "$VENV_DIR/bin/activate" ]]; then
        source "$VENV_DIR/bin/activate"
    else
        echo "⚠️  找不到 venv: $VENV_DIR, 用系统 python3"
    fi
}

cmd_download() {
    ACT
    echo "📥 下载 MiniCPM-V 2.6 int4 量化模型 (约 6GB, ModelScope 国内加速) ..."
    python3 - <<'PY'
import os, sys
dst = os.environ.get("MODEL_DIR", "/data/volc_minicpm/models/MiniCPM-V-2_6-int4")
os.makedirs(os.path.dirname(dst), exist_ok=True)
if os.path.isdir(dst) and any(f.endswith((".safetensors", ".bin")) for f in os.listdir(dst)):
    total = sum(os.path.getsize(os.path.join(r,f)) for r,_,fs in os.walk(dst) for f in fs)/1024**3
    print(f"✅ 模型已存在，跳过下载 (总大小 {total:.1f}GB)")
    sys.exit(0)
try:
    from modelscope import snapshot_download
except Exception:
    print("安装 modelscope ...")
    os.system("pip install -q modelscope")
    from modelscope import snapshot_download
local = snapshot_download(
    "openbmb/MiniCPM-V-2_6-int4",
    cache_dir=os.path.join(os.path.dirname(dst), ".ms_cache"),
    local_dir=dst,
)
print(f"✅ 下载完成: {local}")
PY
}

cmd_start() {
    if [[ -f "$PID_FILE" ]] && kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then
        echo "⚠️  服务已在运行 PID=$(cat "$PID_FILE"), 先 stop 再 start"
        exit 0
    fi
    ACT
    if [[ ! -d "$MODEL_DIR" ]] || ! ls "$MODEL_DIR"/*.safetensors >/dev/null 2>&1 && ! ls "$MODEL_DIR"/*.bin >/dev/null 2>&1; then
        echo "❌ 模型不存在, 先执行: $0 download"
        exit 1
    fi
    cd "$WORKDIR"
    export MINICPMV_MODEL_DIR="$MODEL_DIR"
    export MINICPMV_DTYPE="$DTYPE"
    nohup python3 "$SERVER_PY" --host "$HOST" --port "$PORT" >> "$LOG_DIR/minicpmv.log" 2>&1 &
    echo $! > "$PID_FILE"
    echo "🚀 服务启动 PID=$(cat "$PID_FILE"), 日志=$LOG_DIR/minicpmv.log"
    echo "   正在加载模型 (首次约 30-60 秒) ..."
    for i in {1..24}; do
        sleep 5
        if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
            echo "✅ 服务就绪 → http://$HOST:$PORT"
            echo "   OpenAI 兼容 API Base = http://127.0.0.1:$PORT/v1"
            curl -s "http://127.0.0.1:$PORT/health"
            exit 0
        fi
        if ! kill -0 "$(cat "$PID_FILE")" 2>/dev/null; then
            echo "❌ 进程已退出, 查看日志 tail -30 $LOG_DIR/minicpmv.log"
            tail -30 "$LOG_DIR/minicpmv.log"
            exit 1
        fi
        echo -n "."
    done
    echo "⏰ 等待超时, 手动检查"
}

cmd_stop() {
    if [[ -f "$PID_FILE" ]]; then
        PID=$(cat "$PID_FILE")
        if kill -0 "$PID" 2>/dev/null; then
            kill "$PID" 2>/dev/null || true
            sleep 2
            kill -9 "$PID" 2>/dev/null || true
            echo "🛑 已停止 PID=$PID"
        else
            echo "⚠️  PID $PID 不存在"
        fi
        rm -f "$PID_FILE"
    else
        # 兜底: 按名字杀
        pkill -9 -f "minicpmv_server.py" 2>/dev/null || true
        echo "🛑 已尝试按名字停止 minicpmv_server.py"
    fi
}

cmd_logs() {
    tail -f "$LOG_DIR/minicpmv.log"
}

cmd_test() {
    ACT
    python3 - <<'PY'
import os, sys, json, base64, pathlib, httpx
PORT = int(os.environ.get("PORT", "30000"))
BASE = f"http://127.0.0.1:{PORT}/v1"
try:
    models = httpx.get(f"{BASE}/models", timeout=5).json()
    print("✅ 服务存活, 模型:", models)
except Exception as e:
    print("❌ 服务没起来:", e)
    sys.exit(1)

# 1) 纯文本测试
print("\n🧪 测试1: 纯文本问答 ...")
r = httpx.post(f"{BASE}/chat/completions", timeout=60, json={
    "model": "minicpmv",
    "messages": [{"role": "user", "content": "你好，简单介绍一下你自己。"}],
})
print("HTTP", r.status_code)
data = r.json()
print("🤖 回答:", data["choices"][0]["message"]["content"][:200])

# 2) 如果本地有图片测图片
cands = [
    "/data/volc_minicpm/models/test.jpg",
    "/root/test.jpg",
    "/tmp/test.png",
    os.path.expanduser("~/debug_screenshots/latest.jpg"),
]
img_path = next((p for p in cands if os.path.exists(p)), None)
if img_path:
    print(f"\n🧪 测试2: 图片理解 (文件={img_path}) ...")
    with open(img_path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()
    r2 = httpx.post(f"{BASE}/chat/completions", timeout=120, json={
        "model": "minicpmv",
        "response_format": {"type": "json_object"},
        "messages": [{"role": "user", "content": [
            {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
            {"type": "text", "text": "这张图片里是什么场景？如果有按钮/关闭/跳过/安装/下载 等UI元素，请指出位置和置信度。JSON格式 keys=scene_summary, ui_elements[{label,confidence,description}]"},
        ]}],
    })
    print("HTTP", r2.status_code)
    d = r2.json()
    txt = d["choices"][0]["message"]["content"]
    try:
        print("🤖 结构化回答:", json.dumps(json.loads(txt), ensure_ascii=False, indent=2))
    except Exception:
        print("🤖 原始文本:", txt[:500])
else:
    print("\nℹ️  没找到本地测试图片，跳过图片测试。随便放一张图片到 /root/test.jpg 再跑 test。")
PY
}

case="${1:-start}"
case "$case" in
    download) cmd_download ;;
    start)    cmd_start ;;
    stop)     cmd_stop ;;
    restart)  cmd_stop; sleep 2; cmd_start ;;
    logs)     cmd_logs ;;
    test)     cmd_test ;;
    all)      cmd_download && cmd_start && sleep 3 && cmd_test ;;
    *)
        echo "用法: $0 <download|start|stop|restart|logs|test|all>"
        exit 1
        ;;
esac