#!/usr/bin/env python3
"""模拟 OpenCode MCP 客户端测试 szyjiyi 连接"""
import json
import subprocess
import sys
import time

def send_message(proc, msg):
    """发送 JSON-RPC 消息"""
    line = json.dumps(msg) + "\n"
    proc.stdin.write(line)
    proc.stdin.flush()
    print(f">>> SENT: {msg['method']}", file=sys.stderr)

def read_message(proc):
    """读取 JSON-RPC 响应"""
    line = proc.stdout.readline()
    if not line:
        return None
    try:
        msg = json.loads(line)
        print(f"<<< RECV: {msg.get('result', {}).get('serverInfo', msg.get('error', 'response'))}", file=sys.stderr)
        return msg
    except json.JSONDecodeError as e:
        print(f"<<< ERROR: Failed to parse: {line}", file=sys.stderr)
        return None

def main():
    # 启动 szyjiyi MCP 服务器
    proc = subprocess.Popen(
        [sys.executable, "-m", "szyjiyi", "--project-dir", r"D:\Szy_Augmen_Tools\Szy-jiyimemory"],
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        encoding='utf-8',  # 关键修复：明确指定 UTF-8 编码
        bufsize=1
    )
    
    time.sleep(0.5)  # 等待服务器启动
    
    try:
        # 1. Initialize
        send_message(proc, {
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "capabilities": {},
                "clientInfo": {"name": "test-client", "version": "1.0"}
            }
        })
        resp = read_message(proc)
        if not resp or "error" in resp:
            print(f"FAIL: Initialize failed: {resp}", file=sys.stderr)
            return 1
        
        # 2. List tools
        send_message(proc, {
            "jsonrpc": "2.0",
            "id": 2,
            "method": "tools/list",
            "params": {}
        })
        resp = read_message(proc)
        if not resp or "error" in resp:
            print(f"FAIL: tools/list failed: {resp}", file=sys.stderr)
            return 1
        
        # 3. Call status (模拟 OpenCode 首次调用)
        send_message(proc, {
            "jsonrpc": "2.0",
            "id": 3,
            "method": "tools/call",
            "params": {
                "name": "status",
                "arguments": {}
            }
        })
        resp = read_message(proc)
        if not resp or "error" in resp:
            print(f"FAIL: status call failed: {resp}", file=sys.stderr)
            return 1
        
        # 4. Call recall (模拟实际使用)
        send_message(proc, {
            "jsonrpc": "2.0",
            "id": 4,
            "method": "tools/call",
            "params": {
                "name": "recall",
                "arguments": {
                    "query": "test",
                    "tags": ["preference"],
                    "scope": "user",
                    "top_k": 5
                }
            }
        })
        resp = read_message(proc)
        if not resp or "error" in resp:
            print(f"FAIL: recall call failed: {resp}", file=sys.stderr)
            return 1
        
        print("\nSUCCESS: All MCP operations completed successfully", file=sys.stderr)
        
        # 检查服务器是否还活着
        if proc.poll() is not None:
            print(f"WARNING: Server process exited with code {proc.returncode}", file=sys.stderr)
            stderr_output = proc.stderr.read()
            if stderr_output:
                print(f"Server stderr:\n{stderr_output}", file=sys.stderr)
            return 1
        
        return 0
        
    finally:
        proc.terminate()
        proc.wait(timeout=2)

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