#!/usr/bin/env python3
"""
diff_configs.py — A/B 测试配置对照工具。

给定 2 个或多个 content 目录（或 inputs 目录），解析各自的 nst（L1/L2 的 KEY=VAL）
与 .5 关键字段（L1 TEFF/GRAV、L2 LTE/LTGRAY、L3 nst 名），输出并排对照表，高亮差异。
用于 A/B 测试时快速确认"只有想改的参数变了，其余一致"。

用法：
    python3 diff_configs.py dirA/inputs dirB/inputs dirC/inputs
    python3 diff_configs.py dirA dirB --teff 50000     # 传 content 目录也行（自动找 inputs/）
"""
import argparse
import os
import re
import sys

# nst 一行里逗号分隔的 KEY=VAL（L1 与 L2 都是这种格式）
KV_RE = re.compile(r"(\w+)=([^,]+)")


def parse_nst(path):
    """返回 nst 的 {KEY: VAL} 字典（合并 L1+L2+escape-hatch 行）。"""
    kv = {}
    if not os.path.exists(path):
        return kv
    for line in open(path, "r", errors="replace"):
        for m in KV_RE.finditer(line):
            kv[m.group(1).upper()] = m.group(2).strip()
    return kv


def parse_dot5(path):
    """返回 .5 关键字段。"""
    info = {}
    if not os.path.exists(path):
        return info
    lines = open(path, "r", errors="replace").read().splitlines()
    if len(lines) >= 1:
        t = lines[0].split("!")[0].split()
        info["TEFF"] = t[0] if len(t) > 0 else "?"
        info["GRAV"] = t[1] if len(t) > 1 else "?"
    if len(lines) >= 2:
        t = lines[1].split("!")[0].split()
        info["LTE/LTGRAY"] = f"{t[0]}/{t[1]}" if len(t) >= 2 else "?"
    if len(lines) >= 3:
        m = re.search(r"'([^']+)'", lines[2])
        info["nst名"] = m.group(1) if m else "?"
    return info


def resolve_inputs(d):
    if os.path.basename(d) == "inputs" and os.path.isdir(d):
        return d
    cand = os.path.join(d, "inputs")
    return cand if os.path.isdir(cand) else d


def main():
    ap = argparse.ArgumentParser(description="A/B 配置对照（nst 关键字段 + .5 关键行）")
    ap.add_argument("dirs", nargs="+", help="inputs/ 或 content 目录")
    ap.add_argument("--teff", default=None, help="仅用于显示标题")
    args = ap.parse_args()

    dirs = [resolve_inputs(d) for d in args.dirs]
    names = [os.path.basename(os.path.dirname(d)) if os.path.basename(d) == "inputs"
             else os.path.basename(d) for d in dirs]

    # 收集所有出现过的 key（nst + .5 字段），分两组
    nst_keys_order = ["ND", "NLAMBD", "VTB", "ISPODF", "DDNU", "CNU1", "CHMAX", "ITEK", "NITER",
                      "ORELAX", "IDLTE", "IACC", "ICHANG", "IELCOR", "DPSILG"]
    dot5_keys = ["TEFF", "GRAV", "LTE/LTGRAY", "nst名"]

    nst_maps = []
    dot5_maps = []
    for d in dirs:
        # nst: 各阶段都有，优先 nc.nst（A/B 通常比 nc/nl）；也读 lte/nl
        nst_path = None
        for stage in ("nc", "nl", "lte", "seed_nc"):
            p = os.path.join(d, f"{stage}.nst")
            if os.path.exists(p):
                nst_path = p
                break
        nst_maps.append((parse_nst(nst_path) if nst_path else {}, nst_path))
        # .5：优先 nc.5
        dot5_path = None
        for stage in ("nc", "nl", "lte", "seed_nc"):
            p = os.path.join(d, f"{stage}.5")
            if os.path.exists(p):
                dot5_path = p
                break
        dot5_maps.append((parse_dot5(dot5_path) if dot5_path else {}, dot5_path))

    # 收集所有 nst key（含非标准的）
    all_nst_keys = list(nst_keys_order)
    seen = set(k.upper() for k in all_nst_keys)
    for nm, _ in nst_maps:
        for k in nm:
            if k not in seen:
                all_nst_keys.append(k)
                seen.add(k)

    def print_table(title, keys, maps, getter):
        print(f"\n=== {title} ===")
        hdr = "  字段".ljust(16) + "".join(n.ljust(22) for n in names)
        print(hdr)
        print("  " + "-" * (len(hdr) - 2))
        for k in keys:
            vals = [getter(nm, k) for nm, _ in maps]
            # 高亮：若不全相同，标记 *
            differ = len(set(repr(v) for v in vals)) > 1
            mark = " * " if differ else "   "
            row = mark + k.ljust(13) + "".join((v if v else "—").ljust(22) for v in vals)
            print(row)

    print(f"\n配置对照（{len(names)} 组）" + (f"，Teff={args.teff}" if args.teff else ""))
    print("  * = 该行各组不一致（A/B 测试应只有你想改的字段带 *）")
    for d, n, (nm, np_), (dm, dp_) in zip(dirs, names, nst_maps, dot5_maps):
        print(f"  [{n}]  nst={os.path.basename(np_) or '(无)'}  .5={os.path.basename(dp_) or '(无)'}")

    print_table(".5 关键字段", dot5_keys, dot5_maps, lambda m, k: m.get(k, ""))
    print_table("nst 关键参数（nc/nl 阶段）", all_nst_keys, nst_maps, lambda m, k: m.get(k.upper(), ""))

    # 源文件路径（便于核对）
    print("\n源文件：")
    for d, n, (_, np_), (_, dp_) in zip(dirs, names, nst_maps, dot5_maps):
        print(f"  [{n}] {np_}")
    return 0


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