#!/usr/bin/env python3
"""重建批量测试的算子清单（inventory），并可校验迁移 shape yaml。

对每个 FlagGems 仓库静态扫描 benchmark/ 下的 test_*.py（含 test_ 前缀的
子目录，如 test_FLA/——op_file 记相对路径去前后缀，run_pytest.sh 的
`benchmark/test_${OP_FILE}.py` 拼接对其天然成立；models_benchmark/ 等
非 test_ 前缀路径不在性能批量范围）：

- 每个模块级 test_* 函数出一行 inventory：op（函数名去 test_ 前缀，即
  run_pytest.sh 的 OP）、op_file（文件名去 test_ 前缀/.py 后缀，即 OP_FILE）；
- 标注 marked_skip（函数带无条件 pytest.mark.skip，跑了必 SKIP，批量驱动默认剔除）；
- 提取 op_name 字符串常量（op_name= 关键字实参，以及 *Benchmark 类
  __init__ 里 super().__init__ 的首个字符串位置实参），用于 shape yaml 键校验；
- --verify-collect 时额外跑一次 pytest --collect-only 核实函数确实可被收集
  （AST 见到 ≠ pytest 收得到，import 失败/条件定义都会导致差异）。

产物：
  ops/inventory_<label>.csv        # repo,op,op_file,file,marked_skip,collected,op_names
  --migrate-shapes 时按 op_name 并集过滤输入 yaml、合并上游 core_shapes 底座：
  <out>.yaml + <out>.unmatched.yaml（未匹配任何 op_name 的键，供人工复核后删除）

用法示例（默认扫描两个新仓库）：
  python ops/gen_inventory.py --verify-collect
  # 上游 op_name/core_shapes 变动后，对现有 shape 集重新校验（输入=输出即原地刷新）
  python ops/gen_inventory.py --migrate-shapes ops/shapes_single.yaml \
      --shapes-out ops/shapes_single.yaml
"""

from __future__ import annotations

import argparse
import ast
import csv
import os
import subprocess
import sys
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple

import yaml

OPS_DIR = Path(__file__).resolve().parent

DEFAULT_REPOS = [
    ("main", Path("/workspace/dev/FlagGems")),
    ("vllm", Path("/workspace/dev/FlagGems-vllm")),
]

INVENTORY_FIELDS = ["repo", "op", "op_file", "file", "marked_skip", "collected", "op_names"]


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    p.add_argument("--repo", action="append", default=None, metavar="LABEL=PATH",
                   help="要扫描的仓库（可多次），默认 main=/workspace/dev/FlagGems "
                        "与 vllm=/workspace/dev/FlagGems-vllm")
    p.add_argument("--verify-collect", action="store_true",
                   help="用 pytest --collect-only 核实每个测试函数可被收集（慢，"
                        "需要完整运行环境；结果写进 collected 列）")
    p.add_argument("--out-dir", type=Path, default=OPS_DIR,
                   help="inventory CSV 输出目录（默认 ops/）")
    p.add_argument("--migrate-shapes", type=Path, default=None,
                   help="旧 shape yaml 路径；按扫描出的 op_name 并集过滤后迁移")
    p.add_argument("--shapes-out", type=Path, default=None,
                   help="迁移后 yaml 输出路径（--migrate-shapes 时必填）")
    return p.parse_args()


def _resolve_repos(raw: Optional[List[str]]) -> List[Tuple[str, Path]]:
    if not raw:
        return [(label, path) for label, path in DEFAULT_REPOS if path.is_dir()]
    repos = []
    for item in raw:
        label, _, path = item.partition("=")
        if not path:
            raise SystemExit(f"--repo 需要 LABEL=PATH 形式，收到: {item!r}")
        repos.append((label, Path(path)))
    return repos


def _is_unconditional_skip(dec: ast.AST) -> bool:
    """pytest.mark.skip（非 skipif）——带不带 reason 实参都算无条件。"""
    target = dec.func if isinstance(dec, ast.Call) else dec
    if not (isinstance(target, ast.Attribute) and target.attr == "skip"):
        return False
    mark = target.value
    return (isinstance(mark, ast.Attribute) and mark.attr == "mark"
            and isinstance(mark.value, ast.Name) and mark.value.id == "pytest")


def _op_names_in(node: ast.AST) -> Set[str]:
    """节点范围内的 op_name 字符串常量，覆盖三种上游写法：
    1. 任意调用的 op_name= 关键字实参；
    2. super().__init__ 的首个字符串位置实参（FusedDeepseekV4... 模式）；
    3. *Benchmark 类实例化的首个字符串位置实参（ScaledMMBenchmark("scaled_mm",
       ...) 这类经形参转发、常量提取够不到 super() 调用的模式）。
    误报只会让 shape yaml 多留一个无人读取的键，无害。"""
    names: Set[str] = set()
    for sub in ast.walk(node):
        if not isinstance(sub, ast.Call):
            continue
        for kw in sub.keywords:
            if kw.arg == "op_name" and isinstance(kw.value, ast.Constant) \
                    and isinstance(kw.value.value, str):
                names.add(kw.value.value)
        first_str = (sub.args[0].value
                     if sub.args and isinstance(sub.args[0], ast.Constant)
                     and isinstance(sub.args[0].value, str) else None)
        if first_str is None:
            continue
        func = sub.func
        if (isinstance(func, ast.Attribute) and func.attr == "__init__"
                and isinstance(func.value, ast.Call)
                and isinstance(func.value.func, ast.Name)
                and func.value.func.id == "super"):
            names.add(first_str)
        callee = func.id if isinstance(func, ast.Name) else (
            func.attr if isinstance(func, ast.Attribute) else "")
        if callee.endswith("Benchmark"):
            names.add(first_str)
    return names


def scan_repo(label: str, root: Path) -> Tuple[List[Dict[str, str]], Set[str]]:
    """返回 (inventory 行, 该仓库全部 op_name 集合)。"""
    bench = root / "benchmark"
    rows: List[Dict[str, str]] = []
    op_names_all: Set[str] = set()
    for path in sorted(bench.rglob("test_*.py")):
        rel = path.relative_to(bench)
        # 仅收 test_ 前缀链路（根级 test_x.py 与 test_XXX/ 子目录），保证
        # "benchmark/test_" + op_file + ".py" 能原样重建路径。
        if not str(rel).startswith("test_"):
            continue
        try:
            tree = ast.parse(path.read_text())
        except SyntaxError as exc:
            print(f"[gen] warning: {path} 解析失败，跳过: {exc}", file=sys.stderr)
            continue
        op_names_all |= _op_names_in(tree)
        op_file = str(rel)[len("test_"):-len(".py")]
        for node in tree.body:
            if not isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
                continue
            if not node.name.startswith("test_"):
                continue
            marked = any(_is_unconditional_skip(d) for d in node.decorator_list)
            func_ops = sorted(_op_names_in(node))
            rows.append({
                "repo": label,
                "op": node.name[len("test_"):],
                "op_file": op_file,
                "file": f"benchmark/{rel}",
                "marked_skip": "yes" if marked else "",
                "collected": "",
                "op_names": ";".join(func_ops),
            })
    return rows, op_names_all


def verify_collect(label: str, root: Path, rows: List[Dict[str, str]]) -> None:
    """跑 pytest --collect-only 核实测试函数可被收集，写 collected 列。"""
    env = os.environ.copy()
    env.setdefault("GEMS_VENDOR", "nvidia")  # 跳过 nvidia-smi 子进程探测（可能挂死）
    proc = subprocess.run(
        [sys.executable, "-m", "pytest", "--collect-only", "-q",
         "--continue-on-collection-errors", "benchmark/"],
        cwd=str(root), text=True, capture_output=True, env=env, timeout=1200,
    )
    collected_funcs: Set[Tuple[str, str]] = set()
    for line in proc.stdout.splitlines():
        line = line.strip()
        if "::" not in line or not line.startswith("benchmark/"):
            continue
        file_part, _, rest = line.partition("::")
        func = rest.split("::")[0].split("[")[0]
        collected_funcs.add((file_part, func))
    if not collected_funcs:
        print(f"[gen] warning: {label} collect 无结果 (rc={proc.returncode})，"
              f"stderr 尾部: {proc.stderr[-500:]}", file=sys.stderr)
        return
    for row in rows:
        key = (row["file"], f"test_{row['op']}")
        row["collected"] = "yes" if key in collected_funcs else "no"
    n_missing = sum(1 for r in rows if r["collected"] == "no")
    if n_missing:
        print(f"[gen] {label}: {n_missing} 个 AST 可见但 pytest 未收集到的函数"
              f"（import 失败/条件定义），inventory 中 collected=no")


def carry_over_collected(out_dir: Path, label: str,
                         rows: List[Dict[str, str]]) -> None:
    """未跑 --verify-collect 时，从已有 CSV 继承 collected 列，避免重写清单时
    把上次核实的结果冲掉。"""
    path = out_dir / f"inventory_{label}.csv"
    if not path.is_file():
        return
    with path.open(newline="") as f:
        prev = {(r["op"], r["op_file"]): r.get("collected", "")
                for r in csv.DictReader(f)}
    for row in rows:
        row["collected"] = prev.get((row["op"], row["op_file"]), "")


def write_inventory(out_dir: Path, label: str, rows: List[Dict[str, str]]) -> Path:
    out = out_dir / f"inventory_{label}.csv"
    tmp = out.with_name(out.name + ".tmp")
    with tmp.open("w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=INVENTORY_FIELDS)
        writer.writeheader()
        writer.writerows(rows)
    os.replace(tmp, out)
    return out


def migrate_shapes(old_path: Path, out_path: Path, valid_op_names: Set[str],
                   repos: List[Tuple[str, Path]]) -> None:
    """迁移旧 shape yaml，并以上游 core_shapes.yaml 为底座合并。

    上游 set_shapes 的回退链是 op_name 键 → MRO 类名键（BlasBenchmark 等）→
    基类 DEFAULT_SHAPES（1 维）。--shape_file 是整体替换而非叠加，若我们的
    yaml 缺少类名级条目，BLAS 这类需要 (B,M,N,K) 的算子会跌到 1 维默认值上
    解包崩溃。因此把各仓库 core_shapes.yaml 里我们没有的键（含类名键与
    op 级键）原样并入，迁移条目优先，先并 main 后并 vllm（同名键先到先得）。"""
    data = yaml.safe_load(old_path.read_text()) or {}
    core_maps = []
    for label, root in repos:
        core = root / "benchmark" / "core_shapes.yaml"
        if core.is_file():
            core_maps.append((label, yaml.safe_load(core.read_text()) or {}))
    core_keys = {k for _, m in core_maps for k in m}
    kept = {k: v for k, v in data.items() if k in valid_op_names}
    # core_shapes 来源的键（类名键等）不算"未匹配"——它们每次都从上游取新版，
    # 这样输入=输出的在地刷新也能同步上游 core_shapes 的变动。
    dropped = {k: v for k, v in data.items()
               if k not in valid_op_names and k not in core_keys}
    n_curated = len(kept)
    merged_from = []
    for label, m in core_maps:
        n_before = len(kept)
        for key, val in m.items():
            kept.setdefault(key, val)
        merged_from.append(f"{label}:+{len(kept) - n_before}")
    header = (f"# 由 gen_inventory.py 从 {old_path.name} 迁移：保留新上游仍存在的"
              f" op_name 键 {n_curated}/{len(data)}，"
              f"并合并上游 core_shapes.yaml 缺失键（{' '.join(merged_from)}）。\n"
              "# 顶层键 = FlagGems op_name 或 Benchmark 类名（上游回退链需要）；\n"
              "# 两者都无键的算子使用其基类 DEFAULT_SHAPES。\n")
    out_path.write_text(header + yaml.safe_dump(
        kept, sort_keys=True, default_flow_style=None, allow_unicode=True))
    if dropped:
        unmatched = out_path.with_suffix(".unmatched.yaml")
        unmatched.write_text(
            f"# {old_path.name} 中未匹配新上游任何 op_name 的键（{len(dropped)} 个），"
            "供人工复核后手动挪回。\n"
            + yaml.safe_dump(dropped, sort_keys=True, default_flow_style=None,
                             allow_unicode=True))
        print(f"[gen] shapes: 保留 {len(kept)}，剔除 {len(dropped)} -> {unmatched}")
    else:
        print(f"[gen] shapes: 全部 {len(kept)} 个键有效")
    # 覆盖率：有效 op_name 里有多少没有 shape 条目（用上游默认 shape，仅提示）
    uncovered = sorted(valid_op_names - set(kept))
    print(f"[gen] shapes: 新上游 {len(valid_op_names)} 个 op_name 中 "
          f"{len(uncovered)} 个无 shape 条目（将用上游默认 shape）")


def main() -> None:
    args = parse_args()
    repos = _resolve_repos(args.repo)
    if not repos:
        raise SystemExit("没有可扫描的仓库")
    args.out_dir.mkdir(parents=True, exist_ok=True)

    all_op_names: Set[str] = set()
    for label, root in repos:
        rows, op_names = scan_repo(label, root)
        all_op_names |= op_names
        if args.verify_collect:
            verify_collect(label, root, rows)
        else:
            carry_over_collected(args.out_dir, label, rows)
        out = write_inventory(args.out_dir, label, rows)
        n_skip = sum(1 for r in rows if r["marked_skip"])
        print(f"[gen] {label}: {len(rows)} 个测试函数 -> {out}"
              f"（其中 {n_skip} 个带无条件 skip 标记）")

    if args.migrate_shapes:
        if not args.shapes_out:
            raise SystemExit("--migrate-shapes 需要同时给 --shapes-out")
        migrate_shapes(args.migrate_shapes, args.shapes_out, all_op_names, repos)


if __name__ == "__main__":
    main()