#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
TDC VM 工具链 ⑤ 查询 CLI  —— 对齐 AGENTS.md §4.2 验收标准

验收要求能秒答三个问题:
  * 「这个字段/寄存器谁写的」 -> who <rN>
  * 「这段代码在哪(属于哪个函数)」 -> at @pc
  * 「有多少个子程序/闭包」 -> funcs
另附 grep 正则查询。

用法(先跑 tools/tdc_lift.py 生成 <name>.lifted.txt):
    .venv/bin/python tools/tdc_query.py <lifted.txt> who r7
    .venv/bin/python tools/tdc_query.py <lifted.txt> at @9285
    .venv/bin/python tools/tdc_query.py <lifted.txt> funcs
    .venv/bin/python tools/tdc_query.py <lifted.txt> grep ">>> 11"
"""
from __future__ import annotations

import re
import sys
from typing import List, Tuple


def parse(path: str):
    """切分 lifted 文件 -> [(func_name, [ (pc或None, 行文本) ])]"""
    funcs: List[Tuple[str, List[Tuple[int, str]]]] = []
    cur_name, cur_lines = None, []
    cur_pc = [-1]
    for raw in open(path, encoding="utf-8"):
        line = raw.rstrip("\n")
        m = re.match(r"\s*function\s+([A-Za-z0-9_]+)\(\)\s*\{", line)
        if m:
            if cur_name is not None:
                funcs.append((cur_name, cur_lines))
            cur_name, cur_lines = m.group(1), []
            continue
        if line.strip() == "}" and cur_name is not None:
            funcs.append((cur_name, cur_lines))
            cur_name, cur_lines = None, []
            continue
        if cur_name is not None:
            lm = re.match(r"\s*@(\d+):", line)
            if lm:
                cur_pc[0] = int(lm.group(1))
            cur_lines.append((cur_pc[0], line))
    if cur_name is not None:
        funcs.append((cur_name, cur_lines))
    return funcs


def main() -> int:
    if len(sys.argv) < 3:
        print(__doc__)
        return 2
    path, cmd = sys.argv[1], sys.argv[2]
    arg = sys.argv[3] if len(sys.argv) > 3 else ""
    funcs = parse(path)
    print(f"[{path}] 解析出 {len(funcs)} 个函数")

    if cmd == "funcs":
        for name, lines in funcs:
            pcs = [p for p, _ in lines if p >= 0]
            print(f"  {name:16s} 行数={len(lines):5d}  pc范围=@{min(pcs) if pcs else '-'}..@{max(pcs) if pcs else '-'}")
        return 0

    if cmd == "who":
        # 谁写了 rN (赋值/属性写)
        pat = re.compile(rf"(^|[^\w]){re.escape(arg)}\s*(=[^=]|\[[^\]]*\]\s*=)")
        hits = 0
        for name, lines in funcs:
            for pc, line in lines:
                if pat.search(line):
                    print(f"  {name} @{pc if pc>=0 else '?'}: {line.strip()}")
                    hits += 1
        if not hits:
            print(f"  未发现写 {arg} 的语句")
        return 0

    if cmd == "at":
        want = int(arg.lstrip("@"))
        for name, lines in funcs:
            pcs = [p for p, _ in lines if p >= 0]
            if not pcs or pcs[0] > want:
                continue
            nxt = [p for p in pcs if p >= want]
            if nxt and (want <= max(pcs)):
                print(f"  @{want} 属于函数 {name}  (该函数 pc 范围 @{min(pcs)}..@{max(pcs)})")
                # 打印附近 8 行
                for pc, line in lines:
                    if 0 <= pc and want <= pc <= want + 20:
                        print("     " + line.strip())
                break
        return 0

    if cmd == "grep":
        pat = re.compile(arg)
        for name, lines in funcs:
            for pc, line in lines:
                if pat.search(line):
                    print(f"  {name} @{pc if pc>=0 else '?'}: {line.strip()}")
        return 0

    print(f"未知命令: {cmd}")
    return 2


if __name__ == "__main__":
    raise SystemExit(main())