#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
graph-search.py — 图感知搜索

流程：
    1. qmd vsearch 获取语义匹配结果（Top 10）
    2. 对每个结果，扩展相邻节点：
       - 默认：读取各页面 relations 字段（实时，慢但零状态）
       - --use-graph：读取预构建的 knowledge-graph.json（快，需先 build）
    3. 去重后返回（直接关联 + 关系扩展）

用法：
    python tools/scripts/graph-search.py "查询词" [--count 15] [--depth 1] [--use-graph]

依赖：
    - qmd 已安装（通过 node）
    - --use-graph 需先运行 knowledge-graph.py build
"""
import argparse
import json
import re
import subprocess
import sys
from pathlib import Path

try:
    import yaml
except ImportError:
    yaml = None

QMD = r'node "C:\Users\hhhh2024\AppData\Roaming\npm\node_modules\@tobilu\qmd\dist\cli\qmd.js"'
WIKI = Path(r"D:\Applications\app\kepano-obsidian-main\wiki")
GRAPH_DB = Path(r"D:\Applications\app\kepano-obsidian-main\tools\data\knowledge-graph.json")


def qmd_search(query: str, count: int = 10) -> list[str]:
    """调用 qmd vsearch 获取匹配的页面名列表"""
    cmd = f'{QMD} vsearch "{query}" -c wiki -n {count}'
    result = subprocess.run(cmd, capture_output=True, text=True, shell=True, timeout=30)
    lines = [line.strip() for line in result.stdout.splitlines() if line.strip()]
    return lines


def get_relations(page_name: str) -> list[dict]:
    """读取页面的 relations 字段"""
    fp = WIKI / f"{page_name}.md"
    if not fp.exists():
        return []
    content = fp.read_text(encoding="utf-8")
    m = re.match(r"^---\n(.*?)\n---", content, re.DOTALL)
    if not m:
        return []
    if not yaml:
        return []
    try:
        front = yaml.safe_load(m.group(1)) or {}
    except:
        return []
    return front.get("relations", [])


def expand_nodes(nodes: list[str], depth: int = 1) -> list[str]:
    """从起始节点出发，沿 relations 扩展相邻节点（实时读取页面 frontmatter）"""
    expanded = list(nodes)
    frontier = list(nodes)
    for _ in range(depth):
        next_frontier = []
        for node in frontier:
            rels = get_relations(node)
            for r in rels:
                target = r.get("target", "").strip("[]")
                if target and target not in expanded:
                    expanded.append(target)
                    next_frontier.append(target)
        frontier = next_frontier
        if not frontier:
            break
    return expanded


def expand_nodes_via_graph(nodes: list[str], depth: int = 1, min_conf: int = 3) -> list[str]:
    """通过预构建的 knowledge-graph.json 扩展（更快，无 IO）"""
    if not GRAPH_DB.exists():
        return list(nodes)
    graph = json.loads(GRAPH_DB.read_text(encoding="utf-8"))
    adj: dict[str, list] = {}
    for e in graph["edges"]:
        adj.setdefault(e["source"], []).append((e["target"], e.get("confidence", 3)))
        adj.setdefault(e["target"], []).append((e["source"], e.get("confidence", 3)))

    expanded = list(nodes)
    frontier = list(nodes)
    for _ in range(depth):
        next_frontier = []
        for node in frontier:
            for target, conf in adj.get(node, []):
                if conf >= min_conf and target not in expanded:
                    expanded.append(target)
                    next_frontier.append(target)
        frontier = next_frontier
        if not frontier:
            break
    return expanded


def main():
    parser = argparse.ArgumentParser(description="Graph-aware search for wiki pages")
    parser.add_argument("query", help="Search query")
    parser.add_argument("--count", type=int, default=15, help="Max results")
    parser.add_argument("--depth", type=int, default=1, help="Graph expansion depth")
    parser.add_argument("--use-graph", action="store_true",
                        help="Use prebuilt knowledge-graph.json instead of live frontmatter")
    args = parser.parse_args()

    print(f"=== Searching: {args.query} ===")
    results = qmd_search(args.query, count=max(10, args.count))
    if not results:
        print("  No results from qmd")
        return

    print(f"\nDirect matches ({len(results)}):")
    for r in results:
        print(f"  [[{r}]]")

    if args.use_graph and not GRAPH_DB.exists():
        print(f"\n  (note: {GRAPH_DB.name} not found, run `knowledge-graph.py build`; "
              f"falling back to live frontmatter)")

    if args.use_graph and GRAPH_DB.exists():
        expanded = expand_nodes_via_graph(results, depth=args.depth)
        source = "knowledge-graph.json"
    else:
        expanded = expand_nodes(results, depth=args.depth)
        source = "live frontmatter"

    new = [e for e in expanded if e not in results]
    if new:
        print(f"\nGraph-expanded (via {source}, depth={args.depth}):")
        for n in new[:args.count]:
            print(f"  [[{n}]]")
        if len(new) > args.count:
            print(f"  ... and {len(new) - args.count} more")

    print(f"\nTotal unique: {len(expanded)}")


if __name__ == "__main__":
    main()