#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""output 已有 test-design+cpp 的资料一致性校验脚本（复用前强制）

功能定位：压测 aclnn 接口时，如果 output/{api}/ 下已存在 test-design 与 cpp，则对照 aclnn 资料（md 文档）
判断接口支持的范围与 test-design 里面记录的范围是否一致：
  - 一致（退出码 0）   → 不再生成 CPP，直接使用 output 目录下的 cpp 进行压测
  - 不一致（退出码 2） → 增量更新 test-design（新增缺失的 dtype/format/attr 取值信息，不推翻已有设计）
                         → 重新生成对应场景 CPP → 依次压测
  - 无法判定（退出码 3，md 缺失/解析失败）→ 按不一致处理，走完整重新推导（Step 1）

校验维度：
  1. dtype 支持列表：md 参数说明（含约束说明段落）中每个 tensor/scalar 的 dtype 集合
     vs 03_参数定义.yaml dtype_with_ranges——md 有而 design 无 → 需新增；design 有而 md 无 → 需复核
  2. format：03 yaml format 字段 vs md 提及的 format（ND/NCHW/NHWC 等）
  3. attr 取值范围/枚举：md 中 attr 参数的取值列表（取值：/支持 A、B、C 等模式）
     vs 03/05 yaml 是否记录（design 文本中检索枚举 token）——md 有而 design 无 → 需新增 attr 取值信息

用法:
  python scripts/check_design_consistency.py --api aclnnAbs
  python scripts/check_design_consistency.py --api aclnnAbs --repos ../../CANN-aclnn-api-reference ../../ops-math
  python scripts/check_design_consistency.py --api aclnnAbs --json     # JSON 输出（Agent 消费）
"""

import argparse
import glob
import json
import os
import re
import sys

import yaml

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
    sys.stderr.reconfigure(encoding="utf-8", errors="replace")

# dtype 别名 → 规范名（与 03 yaml dtype 取值一致）；正则按别名长度降序拼接避免子串误匹配
DTYPE_ALIASES = [
    ("bfloat16", "BFLOAT16"), ("bf16", "BFLOAT16"),
    ("float16", "FLOAT16"), ("fp16", "FLOAT16"), ("half", "FLOAT16"),
    ("float32", "FLOAT"), ("fp32", "FLOAT"), ("float", "FLOAT"),
    ("double", "DOUBLE"), ("fp64", "DOUBLE"), ("float64", "DOUBLE"),
    ("int64", "INT64"), ("int32", "INT32"), ("int16", "INT16"), ("int8", "INT8"),
    ("int", "INT32"),
    ("uint64", "UINT64"), ("uint32", "UINT32"), ("uint16", "UINT16"), ("uint8", "UINT8"),
    ("bool", "BOOL"), ("boolean", "BOOL"),
    ("complex64", "COMPLEX64"), ("complex32", "COMPLEX32"),
]
_DTYPE_RE = re.compile(
    r"\b(" + "|".join(sorted((a for a, _ in DTYPE_ALIASES), key=len, reverse=True)) + r")\b",
    re.IGNORECASE)
_ALIAS_MAP = {a.lower(): n for a, n in DTYPE_ALIASES}

_FORMAT_RE = re.compile(r"\b(ND|NCHW|NHWC|NCDHW|NDHWC|NC1HWC0|FZ(?:_[A-Z0-9]+)?)\b")

# attr 取值模式：取值：A/B/C、取值：A、B、C、支持 xxx、yyy、enum 风格大写标识
_ATTR_VALUE_RE = re.compile(
    r"(?:取值|枚举值|可选值|value)\s*[：:]\s*([^\n|]+)"
    r"|支持\s*([A-Za-z][A-Za-z0-9_]*(?:\s*[、,，/|\s]\s*[A-Za-z][A-Za-z0-9_]*)+)")
_ENUM_TOKEN_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_]{1,31}$")

DEFAULT_REPOS = ["CANN-aclnn-api-reference", "ops-math", "ops-nn", "ops-cv", "ops-transformer"]


def log(tag, msg):
    print("[{}] {}".format(tag, msg))


def find_md(api, repos):
    cands = []
    for repo in repos:
        root = repo if os.path.isabs(repo) else os.path.join(os.getcwd(), repo)
        if not os.path.isdir(root):
            continue
        for p in glob.glob(os.path.join(root, "**", api + "*.md"), recursive=True):
            # experimental 目录排除（严谨：aclnn 资料不允许使用 experimental 目录下的接口文档）
            if any(seg.lower() == "experimental" for seg in p.replace("\\", "/").split("/")):
                continue
            cands.append(p)
    if not cands:
        return None

    def prio(p):
        b = os.path.basename(p)[:-3]
        if b == api:
            return 0
        if b.split("&")[0] == api:
            return 1
        return 2

    return sorted(cands, key=prio)[0]


def parse_md_tables(md_text):
    """解析 md 全部表格 → [(行首参数名, 整行文本)]；兼容 <term> 标签。"""
    rows = []
    for line in md_text.split("\n"):
        s = line.strip()
        if not s.startswith("|"):
            continue
        cells = [re.sub(r"<[^>]+>", "", c).strip() for c in s.strip("|").split("|")]
        if not cells or not cells[0] or set(cells[0]) <= set("-: "):
            continue
        rows.append((cells[0], " | ".join(cells)))
    return rows


def extract_dtypes(text):
    """从文本提取规范 dtype 集合。"""
    found = set()
    for m in _DTYPE_RE.finditer(text):
        found.add(_ALIAS_MAP[m.group(1).lower()])
    return found


def extract_formats(text):
    return set(_FORMAT_RE.findall(text))


def extract_attr_values(text):
    """从 attr 参数行文本提取取值枚举 token 集合（启发式）。"""
    tokens = set()
    for m in _ATTR_VALUE_RE.finditer(text):
        seg = m.group(1) or m.group(2) or ""
        for tok in re.split(r"[、,，/|;；\s]+", seg):
            tok = tok.strip("`*")
            if _ENUM_TOKEN_RE.match(tok) and tok.upper() not in ("ND", "AND", "OR", "NONE", "TRUE", "FALSE"):
                tokens.add(tok)
    return tokens


def main():
    parser = argparse.ArgumentParser(
        description="output 已有 test-design+cpp 的资料一致性校验（md 支持范围 vs test-design 记录范围）")
    parser.add_argument("--api", required=True, help="aclnn 接口名，如 aclnnAbs")
    parser.add_argument("--design_dir", default=None,
                        help="test-design 目录（默认 output/{api}/test-design）")
    parser.add_argument("--repos", nargs="*", default=None,
                        help="aclnn 资料（md）查找根目录列表，默认当前目录下 CANN-aclnn-api-reference 与 ops-* 仓库")
    parser.add_argument("--json", action="store_true", help="以 JSON 输出结果（Agent 消费）")
    args = parser.parse_args()

    design_dir = args.design_dir or os.path.join("output", args.api, "test-design")
    param_yaml = os.path.join(design_dir, "03_参数定义.yaml")
    if not os.path.isfile(param_yaml):
        log("ERROR", "缺少 {}（output 目录无 test-design，无需一致性校验，走正常生成流程）".format(param_yaml))
        return 3

    # 1) 载入 design 范围
    design = yaml.safe_load(open(param_yaml, encoding="utf-8")) or {}
    params = design.get("parameters", []) or []
    design_dtypes = {}
    design_formats = {}
    attr_params = []
    for p in params:
        name = str(p.get("name", "")).strip()
        if not name:
            continue
        dtypes = {str(d.get("dtype", "")).upper() for d in (p.get("dtype_with_ranges") or []) if d.get("dtype")}
        design_dtypes[name] = dtypes
        if p.get("format"):
            design_formats[name] = str(p["format"]).upper()
        if p.get("type") in ("aclScalar", "aclBoolArray", "aclIntArray") or \
                (p.get("type") in ("int8_t", "int64_t", "uint64_t", "bool", "double", "float") and p.get("io_type") == "input"):
            attr_params.append(name)

    design_text_all = ""
    for fn in ("03_参数定义.yaml", "05_约束定义.yaml"):
        fp = os.path.join(design_dir, fn)
        if os.path.isfile(fp):
            design_text_all += open(fp, encoding="utf-8", errors="ignore").read()

    # 2) 定位并解析 md 资料
    repos = args.repos if args.repos else DEFAULT_REPOS
    md_path = find_md(args.api, repos)
    if not md_path:
        log("CHECK", "未找到 {} 的 md 资料（查找根目录: {}）→ 无法判定，按不一致处理走重新推导".format(
            args.api, ", ".join(repos)))
        result = {"api": args.api, "verdict": "unknown", "reason": "md 资料未找到",
                  "action": "按不一致处理：走完整重新推导（Step 1）"}
        print(json.dumps(result, ensure_ascii=False, indent=2) if args.json else "")
        return 3
    md_text = open(md_path, encoding="utf-8", errors="ignore").read()
    md_rows = parse_md_tables(md_text)

    def md_text_for(name):
        """该参数的候选文本：表格行 + 全文中含参数名的行（仅该行，避免相邻表格行误引入）。"""
        texts = [row_text for pname, row_text in md_rows
                 if pname == name or pname.startswith(name + "Optional")]
        for ln in md_text.split("\n"):
            if re.search(r"\b" + re.escape(name) + r"\b", ln):
                texts.append(ln)
        return "\n".join(texts[:10])

    # 3) 逐参数比对 dtype / format / attr 取值
    missing_dtype, extra_dtype, format_mismatch, attr_missing, no_signal = [], [], [], [], []
    for name in design_dtypes:
        text = md_text_for(name)
        if not text:
            no_signal.append(name)
            continue
        md_dt = extract_dtypes(text)
        ym_dt = design_dtypes[name] or set()
        if md_dt:
            miss = md_dt - ym_dt
            extra = ym_dt - md_dt
            if miss:
                missing_dtype.append({"param": name, "md有_design缺": sorted(miss), "md集合": sorted(md_dt)})
            if extra:
                extra_dtype.append({"param": name, "design有_md未列": sorted(extra)})
        # format（md 明确提及非 ND format 时比对）
        md_fmt = extract_formats(text)
        ym_fmt = design_formats.get(name, "ND")
        explicit_fmt = {f for f in md_fmt if f != "ND"} - {"NC1HWC0"}  # NC1HWC0 常为内部布局描述
        if explicit_fmt and ym_fmt not in explicit_fmt and ym_fmt != "ND":
            format_mismatch.append({"param": name, "md提及": sorted(explicit_fmt), "design": ym_fmt})
        # attr 取值（仅 attr 类参数）
        if name in attr_params:
            vals = extract_attr_values(text)
            if vals:
                miss_vals = {v for v in vals if v not in design_text_all}
                if miss_vals:
                    attr_missing.append({"param": name, "md取值_design未记录": sorted(miss_vals)})

    inconsistent = bool(missing_dtype or format_mismatch or attr_missing)
    needs_review = bool(extra_dtype)

    if inconsistent:
        verdict, action = "inconsistent", (
            "增量更新 test-design（新增缺失信息，不推翻已有设计）→ 重新生成对应场景 CPP → 依次压测")
    elif needs_review:
        verdict, action = "review", (
            "design 记录的 dtype 多于 md 资料列举（可能为资料滞后或设计冗余），需人工复核；"
            "复核确认后按需更新 test-design 并重新生成 CPP")
    else:
        verdict, action = "consistent", "不再生成 CPP，直接使用 output 目录下的 cpp 进行压测"

    result = {
        "api": args.api,
        "md": os.path.relpath(md_path, os.getcwd()) if os.path.isabs(md_path) else md_path,
        "design_dir": design_dir,
        "verdict": verdict,
        "action": action,
        "missing_dtype": missing_dtype,
        "extra_dtype_review": extra_dtype,
        "format_mismatch": format_mismatch,
        "attr_value_missing": attr_missing,
        "md_no_dtype_signal_params": no_signal,
        "checked_params": len(design_dtypes),
    }

    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        log("CHECK", "api={} | md={} | 校验参数 {} 个".format(
            args.api, result["md"], result["checked_params"]))
        for item in missing_dtype:
            log("GAP", "参数 {} 缺 dtype: {}（md 支持）".format(item["param"], ", ".join(item["md有_design缺"])))
        for item in format_mismatch:
            log("GAP", "参数 {} format 不一致: md={} design={}".format(
                item["param"], "/".join(item["md提及"]), item["design"]))
        for item in attr_missing:
            log("GAP", "attr {} 取值 design 未记录: {}".format(
                item["param"], ", ".join(item["md取值_design未记录"])))
        for item in extra_dtype:
            log("WARN", "参数 {} design 多出 dtype（md 未列，需复核）: {}".format(
                item["param"], ", ".join(item["design有_md未列"])))
        if no_signal:
            log("NOTE", "md 未提取到 dtype 信号的参数（跳过比对）: {}".format(", ".join(no_signal)))
        log("RESULT", "verdict={} → {}".format(verdict, action))
    return 0 if verdict == "consistent" else (2 if verdict in ("inconsistent", "review") else 3)


if __name__ == "__main__":
    sys.exit(main())