#!/usr/bin/env python3
"""调用本地 qwen-mcp 的 qwen_vision 工具识别图片。

用法：
    python3 /Users/admin/workspace/tools/call_qwen_vision.py <图片路径> [prompt]

示例：
    python3 /Users/admin/workspace/tools/call_qwen_vision.py /tmp/test.png
    python3 /Users/admin/workspace/tools/call_qwen_vision.py /tmp/test.png "描述图片内容"
"""
import argparse
import asyncio
import json
import os
import sys

SERVER_SCRIPT = "/Users/admin/workspace/tools/run_mcp_qwen.sh"


async def call_vision(image_path: str, prompt: str) -> str:
    if not os.path.isfile(image_path):
        raise FileNotFoundError(f"图片不存在：{image_path}")

    proc = await asyncio.create_subprocess_exec(
        SERVER_SCRIPT,
        stdin=asyncio.subprocess.PIPE,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE,
    )

    pending: dict[int, asyncio.Future] = {}
    next_id = 1

    async def read_messages():
        while True:
            line = await proc.stdout.readline()
            if not line:
                break
            try:
                msg = json.loads(line.decode().strip())
            except json.JSONDecodeError:
                continue
            if "id" in msg:
                fut = pending.pop(msg["id"], None)
                if fut is not None and not fut.done():
                    fut.set_result(msg)

    async def send(method: str, params: dict) -> dict:
        nonlocal next_id
        req_id = next_id
        next_id += 1
        msg = {"jsonrpc": "2.0", "id": req_id, "method": method, "params": params}
        proc.stdin.write((json.dumps(msg, ensure_ascii=False) + "\n").encode())
        await proc.stdin.drain()
        fut = asyncio.get_event_loop().create_future()
        pending[req_id] = fut
        return await asyncio.wait_for(fut, timeout=120.0)

    reader = asyncio.create_task(read_messages())
    try:
        init_resp = await send("initialize", {
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "clientInfo": {"name": "call_qwen_vision", "version": "1.0.0"},
        })
        if "error" in init_resp:
            raise RuntimeError(f"initialize failed: {init_resp['error']}")

        await send("notifications/initialized", {})

        tool_resp = await send("tools/call", {
            "name": "qwen_vision",
            "arguments": {"prompt": prompt, "image_url": image_path},
        })
        if "error" in tool_resp:
            raise RuntimeError(f"tools/call failed: {tool_resp['error']}")

        result = tool_resp.get("result", {})
        if result.get("isError"):
            raise RuntimeError(f"tool error: {result}")

        parts = []
        for content in result.get("content", []):
            text = content.get("text")
            if text is not None:
                parts.append(text)
        return "\n".join(parts)
    finally:
        proc.stdin.close()
        await proc.wait()
        reader.cancel()
        try:
            await reader
        except asyncio.CancelledError:
            pass


def main() -> int:
    parser = argparse.ArgumentParser(description="调用 qwen_vision 识别图片")
    parser.add_argument("image", help="图片路径")
    parser.add_argument("prompt", nargs="?", default="请描述这张图片里有什么。", help="识别提示词")
    args = parser.parse_args()

    try:
        result = asyncio.run(call_vision(args.image, args.prompt))
        print(result)
        return 0
    except Exception as exc:
        print(f"错误: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())