#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
taxonomy-lookup.py — 实质 / 形式 外部 JSON 词表查询工具

阶段二 子 Agent B 用这个查实质/形式 路径, 避免把词表 (911 + 565 条路径) 塞进 context.

用法:
    taxonomy-lookup.py --dim 实质 --list-l2
    taxonomy-lookup.py --dim 实质 --subtree /理念/知识
    taxonomy-lookup.py --dim 实质 --match "近景人物特写"
    taxonomy-lookup.py --dim 实质 --match "年轻女性 卧室 床上"     # 多 token 自动拆词聚合
    taxonomy-lookup.py --dim 实质 --narrow "AI生成 提示词结构"     # 层级下钻 (扁平 match miss 时降级)
    taxonomy-lookup.py --dim 实质 --validate /理念/知识/化学

词表文件位置 (按优先级):
    1. $TAXONOMY_DIR / 分类库导出_{实质|形式}_*.json     (env override)
    2. procedure-dsl/data/分类库导出_*.json               (DSL 本地副本, 跨平台干净)
    3. <agent-repo>/examples/process_pipeline/script/resource/分类库导出_*.json
                                                          (项目 canonical 词表位置)
    4. <agent-repo>/examples/process_pipeline/test_script/procedure_dsl_agent/data/
       分类库导出_*.json                                  (另一 agent 用的同副本)
    5. ~/Downloads/分类库导出_*.json                      (本地下载兜底, 跨平台)
"""
import argparse
import glob
import json
import os
import sys
from collections import defaultdict
from pathlib import Path

DIM_PATTERNS = {
    '实质': '分类库导出_实质_*.json',
    '形式': '分类库导出_形式_*.json',
}

# Windows 控制台 UTF-8
for _s in (sys.stdout, sys.stderr):
    if hasattr(_s, 'reconfigure'):
        try:
            _s.reconfigure(encoding='utf-8', errors='replace')
        except Exception:
            pass


# DSL_ROOT = procedure-dsl/  (本脚本位于 procedure-dsl/spec/tools/)
# 词表 bundled 在 spec/taxonomy/ 里, skill 自包含.
DSL_ROOT = Path(__file__).resolve().parent.parent.parent
SPEC_DIR = DSL_ROOT / 'spec'
PROCESS_PIPELINE = DSL_ROOT.parents[2]                                  # legacy fallback 用

DEFAULT_SEARCH_PATHS = [
    os.environ.get('TAXONOMY_DIR', ''),
    str(SPEC_DIR / 'taxonomy'),                                         # ← canonical: spec/ 自包含
    str(DSL_ROOT / 'data'),                                             # legacy fallback
    str(PROCESS_PIPELINE / 'script' / 'resource'),                      # legacy fallback
    str(PROCESS_PIPELINE / 'test_script' / 'procedure_dsl_agent' / 'data'),
    str(Path.home() / 'Downloads'),
]


def find_taxonomy_file(dim: str) -> str:
    pattern = DIM_PATTERNS[dim]
    for d in DEFAULT_SEARCH_PATHS:
        if not d:
            continue
        matches = sorted(glob.glob(os.path.join(d, pattern)), reverse=True)
        if matches:
            return matches[0]
    raise FileNotFoundError(f'未找到 {dim} 词表 (pattern: {pattern})')


def load(dim: str) -> dict:
    path = find_taxonomy_file(dim)
    with open(path, 'r', encoding='utf-8') as f:
        return json.load(f)['data']


def cmd_list_l2(data: dict):
    """列二级路径 + 该子树叶子数, 帮助 Agent 选 domain."""
    counts = defaultdict(int)
    for k in data.keys():
        parts = k.split('/')
        if len(parts) >= 3:
            counts['/'.join(parts[:3])] += 1
    print(f'共 {len(data)} 条路径, {len(counts)} 个二级 domain:')
    print()
    for l2, n in sorted(counts.items(), key=lambda x: -x[1]):
        print(f'  {n:>5}  {l2}')


def cmd_subtree(data: dict, prefix: str, max_alias: int = 6):
    """返回某子树的全部叶子 + 部分 alias."""
    if not prefix.startswith('/'):
        prefix = '/' + prefix
    keys = sorted([k for k in data.keys() if k.startswith(prefix)])
    if not keys:
        print(f'(空) 没有路径以 {prefix} 开头')
        return
    print(f'{prefix} 下 {len(keys)} 条叶子路径:')
    print()
    for k in keys:
        aliases = data[k]
        print(f'  {k}')
        if aliases:
            shown = aliases[:max_alias]
            tail = f' ... (+{len(aliases) - max_alias})' if len(aliases) > max_alias else ''
            print(f'    └─ {" / ".join(shown)}{tail}')


def _split_tokens(query: str) -> list:
    """拆 query 为有效 token: 按空白切, 过滤掉单字 (单字噪声过大).

    全是单字的退化 query 会走 fallback: 用整个 query 当单 token 跑.
    这样 `'特写'` 这种单 token 短查询行为跟旧版完全一致.
    """
    raw = [t for t in query.split() if t]
    tokens = [t for t in raw if len(t) >= 2]
    if not tokens and query.strip():
        tokens = [query.strip()]
    return tokens


def _score_for_token(path: str, aliases: list, token: str):
    """单 token 打分 (原 cmd_match 内联逻辑, 抽成独立函数复用).

    Returns:
        (score: int, matched_aliases: list[str])
    """
    score = 0
    matched_aliases = []
    tlow = token.lower()
    # 路径名匹配 (优先级低)
    for part in path.split('/'):
        if not part:
            continue
        if token == part:
            score += 30
        elif token in part:
            score += 10
        elif tlow in part.lower():
            score += 5
    # alias 匹配 (优先级高)
    for alias in aliases:
        if token == alias:
            score += 50
            matched_aliases.append(alias)
        elif token in alias:
            score += 20
            matched_aliases.append(alias)
        elif tlow in alias.lower():
            score += 5
            matched_aliases.append(alias)
    return score, matched_aliases


def cmd_match(data: dict, query: str, top: int = 10):
    """字符串匹配 (无 LLM): 多 token 拆词 → 各自打分 → 聚合排序.

    单 token (无空白) 行为跟旧版完全一致.
    多 token (含空白) 拆词后聚合: 每 path 累加每 token 的 score, 命中 ≥2 token
    给 coverage bonus (避免单 token 大量重复打高分压死多 token 弱命中).
    """
    tokens = _split_tokens(query)
    if not tokens:
        print(f'(空 query)')
        return

    if len(tokens) == 1:
        # 单 token: 原行为
        candidates = []
        single = tokens[0]
        for path, aliases in data.items():
            score, matched = _score_for_token(path, aliases, single)
            if score > 0:
                candidates.append((score, path, matched or aliases[:3], [single]))
        candidates.sort(key=lambda x: (-x[0], len(x[1])))
        _print_candidates(query, candidates, top, tokens=tokens)
        return

    # 多 token: 聚合
    agg = {}  # path -> {score, matched_aliases, matched_tokens, aliases}
    for tok in tokens:
        for path, aliases in data.items():
            s, m = _score_for_token(path, aliases, tok)
            if s > 0:
                entry = agg.setdefault(path, {
                    'score': 0,
                    'matched_aliases': [],
                    'matched_tokens': set(),
                    'aliases': aliases,
                })
                entry['score'] += s
                for a in m:
                    if a not in entry['matched_aliases']:
                        entry['matched_aliases'].append(a)
                entry['matched_tokens'].add(tok)

    candidates = []
    for path, info in agg.items():
        cov = len(info['matched_tokens'])
        score = info['score']
        if cov >= 2:
            score = int(score * (1 + 0.3 * (cov - 1)))    # 多 token 覆盖 bonus
        ms = info['matched_aliases'] or info['aliases'][:3]
        candidates.append((score, path, ms, sorted(info['matched_tokens'])))
    candidates.sort(key=lambda x: (-x[0], len(x[1])))
    _print_candidates(query, candidates, top, tokens=tokens)


def _print_candidates(query: str, candidates: list, top: int, tokens: list = None):
    """match 结果统一打印."""
    if not candidates:
        prefix = f'(无匹配) 查询: {query!r}'
        if tokens and len(tokens) > 1:
            prefix += f'  拆词={tokens}'
        print(prefix)
        print('  提示: 试 --narrow <query> 走层级下钻 (按子树整体语义打分, 容忍 alias 完全没覆盖).')
        return
    header = f'查询: {query!r} — 共 {len(candidates)} 候选, top {min(top, len(candidates))}:'
    if tokens and len(tokens) > 1:
        header = f'查询: {query!r} (拆 {len(tokens)} token: {tokens}) — 共 {len(candidates)} 候选, top {min(top, len(candidates))}:'
    print(header)
    print()
    for score, path, ms, toks in candidates[:top]:
        ax = ' / '.join(ms[:5])
        tok_info = f'   [命中 token: {",".join(toks)}]' if toks and len(toks) > 1 else ''
        print(f'  [score {score:>3}] {path}{tok_info}')
        if ax:
            print(f'                └─ {ax}')


def cmd_narrow(data: dict, query: str, depth: int = 3, top_per_level: int = 3):
    """层级化下钻: 从 L2 起按子树整体语义打分, 每层选 top-K 下钻直到 leaf 或 depth 上限.

    用于扁平 --match 完全 miss 时的兜底: alias 没覆盖但路径段语义沾边的 case
    (e.g. query 'AI生成 提示词' 在 alias 里可能无, 但 `/理念/方法/生成式` 路径段语义对).

    打分: 子树内所有 path + alias 文本里 token 命中次数 (hits) + 命中的不同 token 数 (cov).
    排序: 先按 cov 高优先 (覆盖多 token 的子树更对路), 再按 hits.
    """
    tokens = _split_tokens(query)
    if not tokens:
        print('(空 query)')
        return

    def score_subtree(prefix: str):
        """返回 (hits, covered_token_count) — prefix 子树整体语义信号强度."""
        hits = 0
        covered = set()
        for path, aliases in data.items():
            if not (path == prefix or path.startswith(prefix + '/')):
                continue
            text = (path + ' ' + ' '.join(aliases)).lower()
            for tok in tokens:
                t = tok.lower()
                if t in text:
                    hits += text.count(t)
                    covered.add(tok)
        return hits, len(covered)

    def n_descendants(prefix: str) -> int:
        return sum(1 for k in data if k == prefix or k.startswith(prefix + '/'))

    print(f'层级下钻: query={query!r}  tokens={tokens}  depth={depth}  top_per_level={top_per_level}')
    print()

    # 起点: 所有 L2 (路径前两段)
    l2_set = sorted({'/' + '/'.join(k.split('/')[1:3]) for k in data if k.count('/') >= 2})
    scored = [(p, *score_subtree(p)) for p in l2_set]
    scored = [c for c in scored if c[1] > 0]
    scored.sort(key=lambda x: (-x[2], -x[1]))   # 先 covered 后 hits
    current = scored[:top_per_level]

    if not current:
        print('[L2] 所有 L2 子树都无任何 token 命中. 词表可能无相关语义, 建议换 query.')
        return

    print(f'[L2] 候选 {len(scored)} 个, 取 top {len(current)}:')
    for prefix, hits, cov in current:
        leaf_mark = ' ★叶子' if prefix in data else ''
        print(f'  [hits {hits:>4} / cov {cov}/{len(tokens)}]  {prefix}{leaf_mark}   ({n_descendants(prefix)} 子节点)')
    print()

    # 逐层下钻
    for level_i in range(depth):
        next_set = set()
        for prefix, _, _ in current:
            for k in data:
                if k.startswith(prefix + '/'):
                    rest = k[len(prefix) + 1:].split('/')
                    if rest[0]:
                        next_set.add(prefix + '/' + rest[0])
        if not next_set:
            print(f'[L{3 + level_i}] (已到叶子层, 停止下钻)')
            break
        next_scored = [(p, *score_subtree(p)) for p in next_set]
        next_scored = [c for c in next_scored if c[1] > 0]
        next_scored.sort(key=lambda x: (-x[2], -x[1]))
        if not next_scored:
            print(f'[L{3 + level_i}] (无任何子节点命中, 停止下钻)')
            break
        current = next_scored[:top_per_level]
        print(f'[L{3 + level_i}] 候选 {len(next_scored)} 个, 取 top {len(current)}:')
        for prefix, hits, cov in current:
            leaf_mark = ' ★叶子' if prefix in data else ''
            print(f'  [hits {hits:>4} / cov {cov}/{len(tokens)}]  {prefix}{leaf_mark}   ({n_descendants(prefix)} 子节点)')
        print()

    print('提示: 末层 top 候选若仍非叶子, 用 --subtree <path> 列叶子, 再 --validate <leaf> 确认.')


def cmd_validate(data: dict, path: str) -> bool:
    """校验路径是否在词表中; 路径不存在但是 ancestor 时报 PARTIAL."""
    if not path.startswith('/'):
        path = '/' + path
    if path in data:
        n = len(data[path])
        print(f'OK       {path}  ({n} aliases)')
        return True
    descendants = [k for k in data.keys() if k.startswith(path + '/')]
    if descendants:
        print(f'PARTIAL  {path} 不是叶子, 是 ancestor — 含 {len(descendants)} 条后代路径')
        for d in descendants[:5]:
            print(f'         · {d}')
        if len(descendants) > 5:
            print(f'         ... (+{len(descendants) - 5})')
        return False
    print(f'MISS     {path}  (词表中不存在)')
    return False


def main():
    p = argparse.ArgumentParser(
        description='实质/形式 外部 JSON 词表查询',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument('--dim', required=True, choices=['实质', '形式'],
                   help='词表维度')
    g = p.add_mutually_exclusive_group(required=True)
    g.add_argument('--list-l2', action='store_true',
                   help='列二级路径 + 路径数')
    g.add_argument('--subtree', metavar='PATH',
                   help='返回某子树的全部叶子 + alias')
    g.add_argument('--match', metavar='QUERY',
                   help='匹配候选路径 (多 token 自动拆词聚合; 单 token 行为同旧版)')
    g.add_argument('--narrow', metavar='QUERY',
                   help='层级化下钻 (按子树语义信号打分逐层选 top-K, 兜底 --match miss)')
    g.add_argument('--validate', metavar='PATH',
                   help='校验路径是否在词表中')
    p.add_argument('--top', type=int, default=10,
                   help='match 输出 top N (default: 10)')
    p.add_argument('--max-alias', type=int, default=6,
                   help='subtree 每条路径展示 alias 数 (default: 6)')
    p.add_argument('--depth', type=int, default=3,
                   help='narrow 下钻层数 (default: 3, 即 L2→L3→L4→L5)')
    p.add_argument('--top-per-level', type=int, default=3,
                   help='narrow 每层保留候选数 (default: 3)')
    args = p.parse_args()

    try:
        data = load(args.dim)
    except FileNotFoundError as e:
        print(f'ERROR: {e}', file=sys.stderr)
        sys.exit(2)

    if args.list_l2:
        cmd_list_l2(data)
    elif args.subtree:
        cmd_subtree(data, args.subtree, max_alias=args.max_alias)
    elif args.match:
        cmd_match(data, args.match, top=args.top)
    elif args.narrow:
        cmd_narrow(data, args.narrow, depth=args.depth, top_per_level=args.top_per_level)
    elif args.validate:
        sys.exit(0 if cmd_validate(data, args.validate) else 1)


if __name__ == '__main__':
    main()