#!/usr/bin/env python3
"""
cbm-query-wrapper.py — P3: CBM 查询钩子
封装 CBM CLI 调用，查询结果自动写入织忆。
用法:
  python3 cbm-query-wrapper.py search_graph --project X --name-pattern Y
  python3 cbm-query-wrapper.py get_architecture --project X
"""
import json, os, subprocess, sys, urllib.request, urllib.error

CBM = os.path.expanduser("~/.local/bin/codebase-memory-mcp")
ZHIYI_URL = "http://localhost:7821/api/v1/commit"
ZHIYI_KEY = os.environ.get("ZHIYI_API_KEY", "zhiyi-dev-key-2026")

def zhiyi_write(content, category="distilled"):
    """写入织忆"""
    payload = json.dumps({
        "content": content,
        "category": category,
        "agent_id": "a06"
    }).encode()
    req = urllib.request.Request(ZHIYI_URL, data=payload,
        headers={"X-API-Key": ZHIYI_KEY, "Content-Type": "application/json"})
    try:
        urllib.request.urlopen(req, timeout=5)
        return True
    except urllib.error.URLError as e:
        print(f"[cbm-wrapper] ⚠ 织忆写入失败: {e}", file=sys.stderr)
        return False

def run_cbm(args):
    """执行 CBM CLI 并返回 JSON 结果"""
    cmd = [CBM, "cli"] + args
    try:
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        return result.stdout, result.returncode
    except FileNotFoundError:
        print(f"[cbm-wrapper] ❌ CBM binary not found: {CBM}", file=sys.stderr)
        return "", 1
    except subprocess.TimeoutExpired:
        print("[cbm-wrapper] ⚠ CBM timed out (60s)", file=sys.stderr)
        return "", 1

def main():
    if len(sys.argv) < 2:
        print("用法: cbm-query-wrapper.py <cbm_cli_args...>")
        sys.exit(1)
    
    tool = sys.argv[1]
    args = sys.argv[1:]
    
    # 执行 CBM 查询
    stdout, rc = run_cbm(args)
    print(stdout, end="")
    
    if rc != 0:
        sys.exit(rc)
    
    # 对有价值的结果写入织忆
    try:
        data = json.loads(stdout)
        project = data.get("project", data.get("name", "unknown"))
        nodes = data.get("nodes", data.get("total_nodes", 0))
        edges = data.get("edges") or data.get("total_edges") or 0
        
        if tool in ("index_repository",) and int(nodes) > 0:
            content = f"CBM索引完成: {project} ({nodes}节点/{edges}边)"
            zhiyi_write(content)
        elif tool in ("get_architecture", "list_projects"):
            content = f"CBM架构查询: {project}"
            zhiyi_write(content)
        elif tool in ("search_graph", "trace_path"):
            results = data.get("results") or data.get("nodes") or []
            if results:
                names = [r.get("name","?") for r in results[:5]]
                content = f"CBM查询 '{tool}': 找到 {len(results)} 个结果 ({', '.join(names)})"
                zhiyi_write(content)
    except (json.JSONDecodeError, KeyError, TypeError, ValueError):
        pass  # 非 JSON 输出或不匹配模式就不写入

if __name__ == "__main__":
    main()