#!/usr/bin/env python3
"""
用本地 A30 MiniCPM-V 服务做 "按钮坐标识别" 测试
直接替换之前 AdLLMHandler 的云端调用方式

用法:
    python3 test_local_minicpmv_button.py \
        --image /path/to/screenshot.jpg \
        --api-base http://<A30_IP>:30000/v1 \
        --element close_button
"""
import os, sys, json, base64, argparse, time, pathlib, httpx


PROMPTS = {
    "close_button": """
这是一张手机游戏广告的截图。请识别【关闭按钮】的位置（右上角或右下角的 X、关闭、跳过、跳过3、取消、不感兴趣、以后再说 等）。
只输出 JSON 格式，不要其他内容：
{
  "found": true,
  "x": 横坐标像素 (整数),
  "y": 纵坐标像素 (整数),
  "w": 宽度,
  "h": 高度,
  "label": "识别到的文字",
  "confidence": 0.0~1.0,
  "scene_description": "一句话描述当前画面场景"
}
如果找不到关闭按钮，输出 {"found": false, "reason": "..."}
""".strip(),

    "give_up_button": """
这是一张手机游戏"惊喜福利/免费领取"弹窗的截图，请识别【放弃福利/放弃/不领取/暂不领取/谢谢】按钮，一般在弹窗底部中间或下方。
只输出 JSON，格式同上 (x,y,w,h,confidence, label)。找不到就 found=false。
""".strip(),

    "download_button": """
这是一张手机广告截图，请识别【底部蓝色大按钮（点击下载/点击打开/打开第三方应用/安装/开始游戏/立即体验）】的位置。
只输出 JSON，格式同上。找不到就 found=false。
""".strip(),

    "install_button": """
这是一张 Android 系统安装确认弹窗截图。请识别【安装 Install】按钮位置（一般在弹窗右下角）。
只输出 JSON，格式同上。找不到就 found=false。
""".strip(),
}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--image", required=True)
    ap.add_argument("--api-base", default=os.environ.get("LOCAL_LLM_API", "http://127.0.0.1:30000/v1"))
    ap.add_argument("--model", default="minicpmv")
    ap.add_argument("--element", default="close_button", choices=list(PROMPTS.keys()))
    args = ap.parse_args()

    img = pathlib.Path(args.image).expanduser().resolve()
    if not img.exists():
        print(f"❌ 图片不存在: {img}")
        sys.exit(1)

    with open(img, "rb") as f:
        b64 = base64.b64encode(f.read()).decode()

    url = f"{args.api_base}/chat/completions"
    payload = {
        "model": args.model,
        "temperature": 0.1,
        "max_tokens": 1024,
        "response_format": {"type": "json_object"},
        "messages": [
            {"role": "user", "content": [
                {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}},
                {"type": "text", "text": PROMPTS[args.element]},
            ]}
        ],
    }
    t0 = time.time()
    try:
        r = httpx.post(url, json=payload, timeout=120)
        print(f"HTTP {r.status_code}  ({time.time()-t0:.2f}s)")
        data = r.json()
        content = data["choices"][0]["message"]["content"]
        print(f"\n🤖 模型原始输出:\n{content}\n")
        try:
            parsed = json.loads(content)
            print("✅ 解析为 JSON:")
            print(json.dumps(parsed, ensure_ascii=False, indent=2))
        except Exception as e:
            print(f"⚠️  JSON 解析失败 ({e}), 尝试截取 {{...}}")
            import re
            m = re.search(r"\{.*\}", content, re.S)
            if m:
                print(json.dumps(json.loads(m.group(0)), ensure_ascii=False, indent=2))
    except httpx.HTTPError as e:
        print(f"❌ HTTP 失败: {e}")
        sys.exit(2)


if __name__ == "__main__":
    main()