#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Reqable MCP 客户端(stdio 传输)
用法:
    python3 reqable_mcp.py list                 # 列出所有工具
    python3 reqable_mcp.py call <tool> '<json>'  # 调用工具,input 为 JSON
    python3 reqable_mcp.py m3u8                 # 快捷:搜出所有 m3u8 请求
"""

import json
import subprocess
import sys

MCP_BIN = "/Applications/Reqable.app/Contents/Helpers/mcp-server"
# Reqable 实际监听端口(adb 占了 9000,实际用 9002)
REQABLE_PORT = 9002


class McpClient:
    def __init__(self):
        self.proc = subprocess.Popen(
            [MCP_BIN, "--port", str(REQABLE_PORT)],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,
            text=True,
            bufsize=1,
        )
        self._id = 0

    def _send(self, obj):
        self.proc.stdin.write(json.dumps(obj) + "\n")
        self.proc.stdin.flush()

    def _recv(self, expect_id):
        """读到匹配 id 的 response 为止,跳过 notification"""
        while True:
            line = self.proc.stdout.readline()
            if not line:
                raise RuntimeError("mcp-server 关闭了管道")
            try:
                msg = json.loads(line)
            except json.JSONDecodeError:
                continue
            if msg.get("id") == expect_id:
                if "error" in msg:
                    raise RuntimeError(f"MCP error: {msg['error']}")
                return msg.get("result")

    def request(self, method, params=None):
        self._id += 1
        self._send({"jsonrpc": "2.0", "id": self._id,
                    "method": method, "params": params or {}})
        return self._recv(self._id)

    def initialize(self):
        r = self.request("initialize", {
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "clientInfo": {"name": "hls-whale", "version": "0.1"},
        })
        self._send({"jsonrpc": "2.0", "method": "notifications/initialized"})
        return r

    def list_tools(self):
        return self.request("tools/list")["tools"]

    def call(self, name, arguments):
        return self.request("tools/call", {"name": name, "arguments": arguments})

    def close(self):
        try:
            self.proc.stdin.close()
            self.proc.wait(timeout=3)
        except Exception:
            self.proc.kill()


def main():
    if len(sys.argv) < 2:
        print(__doc__)
        sys.exit(1)
    cmd = sys.argv[1]
    c = McpClient()
    info = c.initialize()
    print(f"[MCP] {info['serverInfo']['name']} {info['serverInfo'].get('version','')}",
          file=sys.stderr)

    try:
        if cmd == "list":
            for t in c.list_tools():
                desc = (t.get("description") or "").split("\n")[0][:80]
                print(f"{t['name']:40s} {desc}")
        elif cmd == "call":
            name = sys.argv[2]
            args = json.loads(sys.argv[3]) if len(sys.argv) > 3 else {}
            r = c.call(name, args)
            for block in r.get("content", []):
                print(block.get("text", ""))
            if r.get("isError"):
                sys.exit(2)
        elif cmd == "m3u8":
            # 快捷:找出搜索类工具并搜 m3u8
            tools = {t["name"]: t for t in c.list_tools()}
            search = next((n for n in tools if "search" in n or "filter" in n or "find" in n), None)
            if not search:
                print("没找到搜索工具,工具列表:", list(tools), file=sys.stderr)
                sys.exit(1)
            print(f"[用 {search}] 搜 m3u8", file=sys.stderr)
            r = c.call(search, {"keyword": "m3u8"} if "keyword" in json.dumps(tools[search]) else {"query": "m3u8"})
            for block in r.get("content", []):
                print(block.get("text", ""))
        else:
            print(__doc__)
    finally:
        c.close()


if __name__ == "__main__":
    main()