#!/usr/bin/env python3
"""MiniMax-M3 图像识别 (火山方舟 Coding Plan)

用法:
  vision.py <图片路径或URL> [-p 提问] [--json] [--max-tokens N] [--detail]
  vision.py a.png b.jpg -p "这两张图有什么区别?"        # 多图对比
  cat prompt.txt | vision.py a.png -p -                  # 从 stdin 读提问

Key 优先级: 环境变量 ARK_CP_API_KEY > ~/.config/ark_cp_api_key
"""
import argparse, base64, json, mimetypes, os, sys, time, urllib.request, urllib.error

ENDPOINT = "https://ark.cn-beijing.volces.com/api/coding/v3/chat/completions"
MODEL = "minimax-m3"


def load_key():
    k = os.environ.get("ARK_CP_API_KEY")
    if k:
        return k.strip()
    p = os.path.expanduser("~/.config/ark_cp_api_key")
    if os.path.exists(p):
        return open(p, encoding="utf-8").read().strip()
    sys.exit("错误: 未找到 API key。请设置环境变量 ARK_CP_API_KEY 或写入 ~/.config/ark_cp_api_key")


def to_part(src):
    """本地文件转 base64 data URI; http(s) 链接直接透传"""
    if src.startswith(("http://", "https://")):
        return {"type": "image_url", "image_url": {"url": src}}
    if not os.path.exists(src):
        sys.exit(f"错误: 文件不存在 {src}")
    mime = mimetypes.guess_type(src)[0] or "image/png"
    if not mime.startswith("image/"):
        sys.exit(f"错误: 不是图片文件 {src} ({mime})")
    b64 = base64.b64encode(open(src, "rb").read()).decode()
    return {"type": "image_url", "image_url": {"url": f"data:{mime};base64,{b64}"}}


def main():
    ap = argparse.ArgumentParser(description="MiniMax-M3 图像识别")
    ap.add_argument("images", nargs="+", help="图片路径或 URL, 可多张")
    ap.add_argument("-p", "--prompt", default="详细描述这张图片的内容。",
                    help="提问; 传 - 表示从 stdin 读")
    ap.add_argument("--json", action="store_true", help="要求模型只输出 JSON")
    ap.add_argument("--max-tokens", type=int, default=2048)
    ap.add_argument("--detail", action="store_true", help="输出耗时/token/request_id")
    a = ap.parse_args()

    prompt = sys.stdin.read().strip() if a.prompt == "-" else a.prompt
    if a.json:
        prompt += "\n\n只输出一个 JSON 对象，不要任何解释、不要代码块标记。"

    content = [{"type": "text", "text": prompt}] + [to_part(s) for s in a.images]
    body = json.dumps({"model": MODEL, "max_tokens": a.max_tokens,
                       "messages": [{"role": "user", "content": content}]}).encode()
    req = urllib.request.Request(ENDPOINT, data=body, headers={
        "Authorization": f"Bearer {load_key()}", "x-api-key": load_key(),
        "Content-Type": "application/json"})

    t0 = time.time()
    try:
        d = json.load(urllib.request.urlopen(req, timeout=600))
    except urllib.error.HTTPError as e:
        sys.exit(f"API 错误 {e.code}: {e.read().decode('utf-8', 'replace')[:400]}")
    except Exception as e:
        sys.exit(f"请求失败: {e}")
    dt = time.time() - t0

    text = (d["choices"][0]["message"].get("content") or "").strip()
    if not text:
        sys.exit("模型返回空内容 (可能 max_tokens 太小被推理耗尽, 试试调大 --max-tokens)")
    print(text)

    if a.detail:
        u = d.get("usage", {})
        print(f"\n--- {dt:.1f}s | in {u.get('prompt_tokens')} / out {u.get('completion_tokens')} tok "
              f"| {d.get('id')}", file=sys.stderr)


if __name__ == "__main__":
    main()