"""CharacterEval: checks the platform's protocol module against the official code (docs/复现规范/CharacterEval-核查矩阵.md).

Compares worker/src/benchmarks/charactereval/protocol.ts (through charactereval_protocol_dump.ts) with the Python
reference charactereval_reference.py, whose official functions are the repository's source text (checked here too):

    git clone https://github.com/morecry/CharacterEval repo && git -C repo checkout c3d44a6
    python3 -I tools/experiments/charactereval_spec_check.py repo

S: the pinned sources and the generated files. V01: data hashes and counts. V02: profile reprs and all 8,032 reward-model
inputs (the released frozen replies results/generation_trans_qwen_7b.jsonl as the responses). V03: all 4,531 message
arrays. V04: first-line cut and failure filter on edge cases. V08: aggregation. M: metric table. One line per check;
exit status 1 when any check fails. Runs no model and calls no endpoint.
"""
import ast
import collections
import hashlib
import importlib.util
import json
import os
import random
import subprocess
import sys
import tempfile

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(HERE, "..", ".."))
DUMP_TS = os.path.join(HERE, "charactereval_protocol_dump.ts")

_spec = importlib.util.spec_from_file_location("charactereval_reference", os.path.join(HERE, "charactereval_reference.py"))
ref = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(ref)

results = []


def check(cid, what, ok, detail=""):
    results.append(bool(ok))
    print(f"{'PASS' if ok else 'FAIL'} {cid} {what}{(': ' + detail) if detail else ''}")


def load(repo, rel):
    return json.load(open(os.path.join(repo, rel), encoding="utf-8"))


def first_diff(a, b):
    for i, (x, y) in enumerate(zip(a, b)):
        if x != y:
            return i
    return None if len(a) == len(b) else min(len(a), len(b))


def main():
    if len(sys.argv) != 2:
        raise SystemExit(__doc__)
    repo = os.path.abspath(sys.argv[1])

    # S: pinned sources, verbatim blocks, generated files.
    try:
        head = subprocess.run(["git", "-C", repo, "rev-parse", "HEAD"], capture_output=True, text=True, check=True).stdout.strip()
    except (subprocess.CalledProcessError, FileNotFoundError):
        head = "(not a git checkout)"
    check("S01", "repository commit", head == ref.COMMIT, head)
    for f in ("get_response.py", "transform_format.py", "run_char_rm.py", "compute_score.py"):
        got = ref.sha256_file(os.path.join(repo, f))
        check("S02", f"{f} SHA-256", got == ref.SHA256[f], got[:12])
    blocks = ref.verbatim_blocks()
    bad = [f"{f} {loc}" for f, loc, mode, text in blocks if text != ref.expected_block(ref.source_segment(repo, f, loc), mode)]
    check("S03", f"reference: {len(blocks)} verbatim blocks equal the repository's source text", len(blocks) == 7 and not bad, ", ".join(bad))
    cs = ast.parse(ref.read(os.path.join(repo, "compute_score.py")))
    official_mean = [n for n in ast.walk(cs) if isinstance(n, ast.BinOp) and isinstance(n.op, ast.Div)]
    mine = [n for n in ast.walk(ast.parse(ref.read(os.path.join(HERE, "charactereval_reference.py"))))
            if isinstance(n, ast.BinOp) and isinstance(n.op, ast.Div) and "score_dict" in ast.unparse(n)]
    check("S04", "reference: compute_score mean expression equals compute_score.py (AST)",
          len(official_mean) == 1 and len(mine) == 1 and ast.dump(official_mean[0]) == ast.dump(mine[0]), ast.unparse(official_mean[0]))
    check("S05", "strings.gen.ts is what the generator extracts now",
          ref.read(ref.STRINGS_GEN) == ref.strings_gen(repo))
    check("S06", "profiles.gen.json is what the generator writes now", ref.read(ref.PROFILES_GEN) == ref.profiles_gen(repo))

    # V01: data files and counts.
    for rel in ("data/test_data.jsonl", "data/id2metric.jsonl", "data/character_profiles.json"):
        got = ref.sha256_file(os.path.join(repo, rel))
        check("V01", f"{rel} SHA-256", got == ref.SHA256[rel], got[:12])
    test = load(repo, "data/test_data.jsonl")
    id_metric = load(repo, "data/id2metric.jsonl")
    profiles = ref.load_profiles(repo)
    ids = [d["id"] for d in test]
    scored = [d for d in test if str(d["id"]) in id_metric]
    pairs = sum(len(v) for v in id_metric.values())
    check("V01", "test examples / unique ids = 4,564 / 4,564", len(test) == 4564 and len(set(ids)) == 4564, f"{len(test)} / {len(set(ids))}")
    check("V01", "examples with metrics = 4,531 (33 without)", len(scored) == 4531 and len(test) - len(scored) == 33, f"{len(scored)}")
    check("V01", "(example, metric) pairs = 8,032; every id2metric id is a test id",
          pairs == 8032 and set(id_metric) <= {str(i) for i in ids}, f"{pairs}")
    roles = {d["role"] for d in test}
    check("V01", "characters used / profiles = 77 / 78, every used one has a profile",
          len(roles) == 77 and len(profiles) == 78 and roles <= set(profiles), f"{len(roles)} / {len(profiles)}")

    # Platform side.
    rng = random.Random(20261028)
    probe = next(d for d in test if len(id_metric.get(str(d["id"]), [])) == 2)
    fl_cases = ["", "\n", "\n\n\n", "\nabc", "abc\ndef", "abc\r\ndef", "\r\n", "abc\r", "  角色：你好  \n续写", "a\n", "a
b
c\x85d\x0be",
                "ERROR", "error", " ERROR", "ERROR\n", None]
    order = [m for g in ref.GROUPS.values() for m in g]
    known = []
    for i, m in enumerate(order):
        mu = 1 + 0.25 * i
        known.append([(m, mu - 0.5), (m, mu + 0.5)] + [(m, mu)] * (1 + i))
    known_rows = [r for k in range(max(map(len, known))) for rows in known for r in rows[k:k + 1]]  # interleaved
    known_means = {m: 1 + 0.25 * i for i, m in enumerate(order)}
    known_dims = {"ca": 1.25, "cc": 2.25, "ra": 3.375}
    frozen = load(repo, "results/generation_trans_qwen_7b.jsonl")
    agg = [
        ("known", known_rows),
        ("known-missing-fluency", [r for r in known_rows if r[0] != "Fluency"]),
        ("random-uniform-8032", [(r["metric_en"], rng.uniform(1, 5)) for r in frozen]),
        ("random-bf16-grid-8032", [(r["metric_en"], rng.randrange(0, 257) / 256 * 4 + 1) for r in frozen]),
        ("random-sparse-40", [(rng.choice(order[1:]), rng.uniform(1, 5)) for _ in range(40)]),
    ]
    cases = {"firstLine": fl_cases, "firstLineId": probe["id"],
             "aggregation": [{"name": n, "rows": [{"metricEn": m, "score": s} for m, s in rows]} for n, rows in agg]}
    with tempfile.TemporaryDirectory() as tmp:
        cp, op = os.path.join(tmp, "cases.json"), os.path.join(tmp, "out.json")
        json.dump(cases, open(cp, "w", encoding="utf-8"), ensure_ascii=False)
        run = subprocess.run(["node", "--experimental-strip-types", "--disable-warning=ExperimentalWarning", DUMP_TS, repo, cp, op],
                             capture_output=True, text=True, cwd=ROOT)
        if run.returncode != 0:
            check("TS", "protocol dump ran", False, run.stderr.strip()[-2000:])
            sys.exit(1)
        ts = json.load(open(op, encoding="utf-8"))

    # V02: profile reprs and reward-model inputs.
    gen = json.load(open(ref.PROFILES_GEN, encoding="utf-8"))
    h = ref.sorted_profiles_hash(gen["profiles"])
    check("V02", "profile reprs (sorted, role\\trepr\\n) SHA-256", h == ref.PROFILES_SORTED_SHA256 == gen["sortedReprSha256"]
          and gen["profiles"] == ref.profile_reprs(repo), h[:12])
    outputs = {}
    for r in frozen:
        outputs.setdefault(r["id"], r["model_output"])
    datas = [dict(d, model_output=outputs.get(d["id"])) for d in test]  # the 33 without metrics have no frozen reply
    records = ref.transform(datas, id_metric)
    keys = ("id", "role", "novel_name", "context", "model_output", "metric_en", "metric_zh")
    same = len(records) == len(frozen) and all(all(a[k] == b[k] for k in keys) for a, b in zip(records, frozen))
    check("V02", "reference transform of the frozen replies reproduces generation_trans_qwen_7b.jsonl row for row", same, f"{len(records)} rows")
    py_rm = [(r["id"], r["metric_en"], t) for r, t in zip(records, ref.rm_inputs(records, profiles))]
    ts_rm = [(r["id"], r["metric_en"], r["text"]) for r in ts["rm"]]
    nd = sum(a != b for a, b in zip(py_rm, ts_rm)) + abs(len(py_rm) - len(ts_rm))
    i = first_diff(py_rm, ts_rm)
    check("V02", f"reward-model inputs, Python vs TS, byte for byte ({len(py_rm)} / {len(ts_rm)})",
          len(py_rm) == 8032 and nd == 0, f"{nd} differ" + (f", first at row {i}" if i is not None else ""))

    # V03: message arrays.
    py_msg = {str(d["id"]): ref.target_messages(d, profiles) for d in scored}
    ts_msg = ts["messages"]
    diff = [k for k in py_msg if py_msg[k] != ts_msg.get(k)] + [k for k in ts_msg if k not in py_msg]
    lens = [len(v) for v in py_msg.values()]
    check("V03", f"message arrays, Python vs TS, byte for byte ({len(py_msg)} / {len(ts_msg)}; {min(lens)}–{max(lens)} messages)",
          len(py_msg) == 4531 and not diff, f"{len(diff)} differ" + (f", e.g. id {diff[0]}" if diff else ""))
    no_system = all(m["role"] in ("user", "assistant") for v in ts_msg.values() for m in v)
    alternating = all(all(m["role"] == ("user" if j % 2 == 0 else "assistant") for j, m in enumerate(v)) and len(v) % 2 == 1
                      for v in ts_msg.values())
    check("V03", "no system message; user/assistant alternate and the last message is the user's", no_system and alternating)

    # V04: first line and failure filter.
    bad = []
    for c, t in zip(fl_cases, ts["firstLines"]):
        rows = [[r["metric_en"], r["metric_zh"], r["model_output"]] for r in ref.transform([dict(probe, model_output=c)], id_metric)]
        fl = None if c is None else c.split("\n")[0]
        if t["firstLine"] != fl or t["rows"] != rows or t["failure"] != (not rows):
            bad.append(repr(c))
    check("V04", f"first line and failure filter on {len(fl_cases)} edge cases (leading newline, only newlines, \\r\\n, empty, ERROR, None)",
          len(ts["firstLines"]) == len(fl_cases) and not bad, ", ".join(bad))

    # V08: aggregation.
    worst, bad = 0.0, []
    for (name, rows), t in zip(agg, ts["aggregation"]):
        means = ref.compute_score([{"metric_en": m, m: s} for m, s in rows])
        dims = ref.dimension_averages(means)
        n = collections.Counter(m for m, _ in rows)
        if set(means) != set(t["means"]) or dict(n) != t["n"] or set(dims) != set(t["dims"]):
            bad.append(name)
            continue
        for k in means:
            worst = max(worst, abs(means[k] - t["means"][k]))
        for g in dims:
            if (dims[g] is None) != (t["dims"][g] is None):
                bad.append(f"{name}/{g}")
            elif dims[g] is not None:
                worst = max(worst, abs(dims[g] - t["dims"][g]))
        if name == "known" and (means != known_means or dims != known_dims or t["dims"] != known_dims):
            bad.append("known answer")
        if name == "known-missing-fluency" and not (dims["ca"] is None and t["dims"]["ca"] is None and dims["cc"] == 2.25):
            bad.append("known missing")
    weighted_ca = sum(s for m, s in known_rows if m in ref.GROUPS["ca"]) / sum(1 for m, _ in known_rows if m in ref.GROUPS["ca"])
    check("V08", f"per-metric means and dimension averages, reference vs TS on {len(agg)} score sets (known answer, random)",
          not bad and worst < 1e-9 and weighted_ca != known_dims["ca"], f"max |diff| {worst:.3g}" + (f"; {', '.join(bad)}" if bad else ""))

    # M: metric table.
    pairs_data = ref.metrics_from_data(repo)
    tsm = ts["metrics"]
    check("M01", "12 metrics: TS (metric_en, metric_zh) equal id2metric's distinct pairs",
          len(tsm) == 12 and sorted([m["en"], m["zh"]] for m in tsm) == sorted(pairs_data))
    check("M02", "dimension membership equals the reference (Character Consistency has 5, including Exposure/KE)",
          ts["groups"] == ref.GROUPS and {m["en"]: m["group"] for m in tsm} == {m: g for g, ms in ref.GROUPS.items() for m in ms})
    check("M03", "metric keys, abbreviations and dimensions agree with meta.ts DIMENSIONS", not ts["metaMismatch"], ", ".join(ts["metaMismatch"]))

    print(f"{sum(results)}/{len(results)} checks passed")
    sys.exit(0 if all(results) else 1)


if __name__ == "__main__":
    main()