#!/usr/bin/env python3
"""
check_inputs.py — TLUSTY 测试执行前输入校验（冷启动链 / 种子链）

复刻 docs/testing_workflow_2026_08_11.md §五 与 references/tlusty-stages.md 的"关键不变量"。
退出码：0=全过，1=有违规。

用法：
    python3 check_inputs.py <inputs_dir> --chain cold --teff 60000
    python3 check_inputs.py <inputs_dir> --chain seed --teff 55000
"""
import argparse
import os
import re
import sys

# NaN 模式（与 conv_check.rs:18 NAN_RE_PATTERN 一致）
NAN_RE = re.compile(r"(?i)(\bnan\b|\binf(?:inity)?\b|\*{3,}|[eE]\+(?:3\d{2}|[4-9]\d{2,}))")


def read_lines(path):
    with open(path, "r", errors="replace") as f:
        return f.read().splitlines()


def check_dot5_l2(path, expect_lte, expect_ltgray):
    """校验 .5 第2行 LTE/LTGRAY 两个标志。行首允许空格。"""
    lines = read_lines(path)
    if len(lines) < 2:
        return f"{path}: 文件不足2行"
    l2 = lines[1].split("!")[0].split()
    if len(l2) < 2:
        return f"{path}: L2 字段不足2个: {lines[1]!r}"
    lte, ltgray = l2[0].upper(), l2[1].upper()
    if lte != expect_lte or ltgray != expect_ltgray:
        return f"{path}: L2 期望 {expect_lte} {expect_ltgray}, 实际 {lte} {ltgray} ({lines[1].strip()!r})"
    return None


def check_dot5_nst_name(path):
    """校验 .5 第3行 'nst' 文件名。
    约定：生产/测试默认 nst 名都是字面 'nst'；run.sh 会在执行时把 {stage}.nst 复制成 nst。
    所以 'nst' 不算违规（运行时会有）；只有引用了非标准名且该名在 inputs/ 找不到时才报错。"""
    lines = read_lines(path)
    if len(lines) < 3:
        return None
    l3 = lines[2]
    m = re.search(r"'([^']+)'", l3)
    if not m:
        return None
    nst_name = m.group(1)
    if nst_name == "nst":
        return None  # 标准名，run.sh 运行时把 {stage}.nst → nst，无需 inputs/ 里有
    base = os.path.dirname(path)
    candidates = [os.path.join(base, nst_name), os.path.join(base, "..", nst_name)]
    if not any(os.path.exists(c) for c in candidates):
        return f"{path}: L3 指定 nst 文件 {nst_name!r} 未找到（非标准名，run.sh 不会自动改名）"
    return None


def check_nst_niter(path):
    """校验 nst 第1行含 NITER=。"""
    lines = read_lines(path)
    if not lines:
        return f"{path}: 空文件"
    if "NITER" not in lines[0].upper():
        return f"{path}: L1 缺 NITER=: {lines[0][:80]!r}"
    return None


def check_file_exists(path, what):
    return None if os.path.exists(path) else f"缺输入: {what} ({path})"


def check_no_nan(path, what):
    if not os.path.exists(path):
        return f"缺输入: {what} ({path})"
    n = 0
    with open(path, "r", errors="replace") as f:
        for line in f:
            if NAN_RE.search(line):
                n += 1
    return None if n == 0 else f"{what} 含 {n} 行 NaN/Inf/坏值 ({path})"


def main():
    ap = argparse.ArgumentParser(description="TLUSTY 测试执行前输入校验")
    ap.add_argument("inputs_dir", help="inputs/ 目录路径")
    ap.add_argument("--chain", choices=["cold", "seed"], required=True)
    ap.add_argument("--teff", type=float, default=None, help="Teff（K），用于上下文，非必填")
    args = ap.parse_args()

    d = args.inputs_dir
    violations = []

    def add(v):
        if v:
            violations.append(v)

    if args.chain == "cold":
        # 必须有 lte/nc/nl 的 .5 与 .nst
        for f in ["lte.5", "nc.5", "nl.5", "lte.nst", "nc.nst", "nl.nst"]:
            add(check_file_exists(os.path.join(d, f), f))
        # L2: lte=T T, nc/nl=F F
        if os.path.exists(os.path.join(d, "lte.5")):
            add(check_dot5_l2(os.path.join(d, "lte.5"), "T", "T"))
        for s in ("nc", "nl"):
            p = os.path.join(d, f"{s}.5")
            if os.path.exists(p):
                add(check_dot5_l2(p, "F", "F"))
        # nst 名一致
        for s in ("lte", "nc", "nl"):
            p = os.path.join(d, f"{s}.5")
            if os.path.exists(p):
                add(check_dot5_nst_name(p))
        # nst L1 含 NITER
        for s in ("lte", "nc", "nl"):
            p = os.path.join(d, f"{s}.nst")
            if os.path.exists(p):
                add(check_nst_niter(p))
        # 冷启动链不应有 fort.8
        if os.path.exists(os.path.join(d, "fort.8")):
            violations.append("冷启动链不应有 inputs/fort.8（lte 阶段 ltgray=T 会删它；若有种子请用 --chain seed）")

    else:  # seed
        for f in ["seed_nc.5", "seed_nc.nst", "fort.8"]:
            add(check_file_exists(os.path.join(d, f), f))
        add(check_no_nan(os.path.join(d, "fort.8"), "fort.8 种子"))
        if os.path.exists(os.path.join(d, "seed_nc.5")):
            add(check_dot5_l2(os.path.join(d, "seed_nc.5"), "F", "F"))
            add(check_dot5_nst_name(os.path.join(d, "seed_nc.5")))
        if os.path.exists(os.path.join(d, "seed_nc.nst")):
            add(check_nst_niter(os.path.join(d, "seed_nc.nst")))

    # data 软链：run.sh / new_test.sh 会在执行时创建，inputs/ 阶段可能还没有。
    # 这里只做提示（note），不算硬违规。
    notes = []
    data_link = None
    for cand in [os.path.join(d, "..", "run", "data"), os.path.join(d, "data")]:
        if os.path.islink(cand) or os.path.isdir(cand):
            data_link = cand
            break
    if data_link is None:
        notes.append("data 软链（run/data）暂不存在——run.sh 执行时会自动创建并指向 assets/data/，可忽略")

    # 汇报
    print(f"=== 输入校验: {d} (chain={args.chain}" +
          (f", teff={args.teff}" if args.teff else "") + ") ===")
    for n in notes:
        print(f"  ℹ {n}")
    if not violations:
        print("✅ 全部通过")
        return 0
    for v in violations:
        print(f"  ❌ {v}")
    print(f"\n共 {len(violations)} 项违规。请修正后再执行（见 references/tlusty-stages.md）。")
    return 1


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