"""CharacterBench audit, authors' side: runs the authors' own code on the authors' data and on what production ran, and
compares every stage with the platform (characterbench_audit_platform.ts output and the D1 exports).

Stages (each writes counts, the ids of every difference and the reason):

  A data     the repository's files and the samples process.py keeps (norm_session_zh/en in one try block) against the
             samples production holds (D1 cb_items): kept set, messages, character profile, annotation score.
  B prompts  role-play prompt: the paper's Table 11 filled with process.py's dialogue (protocol 2.0) — must match the
             platform for every sample; the code's roleplay_prompt.py is reported as the intended difference. Judge input:
             process_wo_context_{zh,en}_all.py process() for a placeholder reply and for every reply each run stored.
  E scoring  each stored judge output read by all_eval.py's check() and float(), the sample's 10 outputs averaged
             (self-consistency, as the released pred_*_sc), each column = mean / max(golds) × 5 with golds the human
             scores of the samples kept (all_eval.py), AVG = mean of the 13 columns; against what production stored
             per sample, per column and per run, and the ranking.

The authors' functions are read from the repository by source segment and executed unchanged (characterbench_prompt_check
helpers); nothing of the platform is imported.

    node --disable-warning=ExperimentalWarning tools/experiments/characterbench_audit_platform.ts EXPORT REPO platform.json
    python3 -I tools/experiments/characterbench_audit.py REPO TABLE11.json EXPORT platform.json audit.json
"""
import collections
import copy
import hashlib
import math
import json
import os
import statistics
import sys
import tempfile

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from characterbench_prompt_check import data_info, source_of, validate_session  # noqa: E402

FILES = [("memory_consistency", "mc"), ("fact_accuracy", "fa"), ("boundary_consistency", "bck"), ("attribute_consistency_bot", "acb"),
         ("attribute_consistency_human", "ach"), ("behavior_consistency_bot", "bcpb"), ("behavior_consistency_human", "bcph"),
         ("emotion_self_regulation", "es"), ("empathetic_responsiveness", "er"), ("morality_stability", "ms"), ("morality_robustness", "mr"),
         ("human_likeness", "hl"), ("engagement", "eg")]
KEY = dict(FILES)
sha1 = lambda s: hashlib.sha1(s.encode()).hexdigest()


def jl(path):
    with open(path, encoding="utf-8") as f:
        return [json.loads(l) for l in f if l.strip()]


def official(repo):
    """The authors' functions, executed from their source."""
    ns = {}
    exec(source_of(f"{repo}/process.py", {"norm_session_zh", "norm_session_en"}), ns)
    ev = {"np": None}
    exec(source_of(f"{repo}/all_eval.py", {"check"}), ev)
    proc = {}
    for lang in ("zh", "en"):
        p = {"json": json, "deepcopy": copy.deepcopy, "data_info": data_info(repo, repo, lang)}
        exec(source_of(f"{repo}/construct_prompts/process_wo_context_{lang}_all.py", {"process"}), p)
        proc[lang] = p["process"]
    rp = {}
    for node in __import__("ast").parse(open(f"{repo}/roleplay_prompt.py", encoding="utf-8").read()).body:
        rp[node.targets[0].id] = node.value.value
    return ns["norm_session_zh"], ns["norm_session_en"], proc, ev["check"], rp


def judge_inputs(proc, lang, dt, sessions):
    """process() on a list of (key, session-with-response); returns key -> (instruction, output) for the rows it keeps."""
    path = os.path.join(tempfile.mkdtemp(), f"{dt}_test.json")
    json.dump([s for _, s in sessions], open(path, "w", encoding="utf-8"), ensure_ascii=False)
    outs, _ = proc[lang](path, dt)
    # process() appends in input order and only skips rows without a response or with an empty one.
    kept = [k for k, s in sessions if (s.get("response_zh") if lang == "zh" else s["translation_en"].get("response_en")) not in (None, "")]
    assert len(kept) == len(outs), (dt, lang, len(kept), len(outs))
    return {k: (o["instruction"], o["output"]) for k, o in zip(kept, outs)}


def main(repo, table11, export, platform_path, out_path):
    norm_zh, norm_en, proc, check, code_rp = official(repo)
    t11 = json.load(open(table11, encoding="utf-8"))
    plat = json.load(open(platform_path, encoding="utf-8"))
    report = {}

    # ------------------------------------------------------------------ A. data
    raw = {dt: json.load(open(f"{repo}/eval_data/raw_data/{dt}_test.json", encoding="utf-8")) for dt, _ in FILES}
    repo_hashes = {}
    for p in plat["sourceFiles"]:
        repo_hashes[p] = hashlib.sha256(open(os.path.join(repo, p), "rb").read()).hexdigest()
    prod_items = {}
    for r in jl(os.path.join(export, "cb_items.jsonl")):
        it = json.loads(r["item_json"])
        prod_items[f"{it['lang']}|{it['file']}|{it['index']}"] = it
    kept, dropped, norm = {}, [], {}
    for dt, _ in FILES:
        for i, session in enumerate(raw[dt]):
            try:
                szh = validate_session(norm_zh(copy.deepcopy(session)))
                sen = validate_session(norm_en(copy.deepcopy(session)))
            except Exception as e:  # process.py: the same exception on all five attempts drops the sample
                dropped.append({"file": dt, "index": i, "id": session.get("id"), "reason": f"{type(e).__name__}: {e}"})
                continue
            kept[(dt, i)] = session
            norm[("zh", dt, i)], norm[("en", dt, i)] = szh, sen
    official_keys = {f"{l}|{dt}|{i}" for (dt, i) in kept for l in ("zh", "en")}
    field_diffs = collections.defaultdict(list)
    for (l, dt, i), s in norm.items():
        k = f"{l}|{dt}|{i}"
        it = prod_items.get(k)
        if it is None:
            continue
        msgs = [{"role": m.role, "content": m.content} for m in s.messages]
        if msgs != it["messages"]:
            field_diffs["messages"].append(k)
        if s.dialogue_setting["bot_profile"] != it["botProfile"]:
            field_diffs["botProfile"].append(k)
        if float(raw[dt][i]["annotation_score"]) != it["annotation"]:
            field_diffs["annotation"].append(k)
        if raw[dt][i].get("id", i) != it["id"]:
            field_diffs["id"].append(k)
    report["A_data"] = {
        "repository_files": len(repo_hashes),
        "repo_vs_platform_build_hash_mismatch": sorted(p for p in repo_hashes if repo_hashes[p] != plat["sourceFiles"].get(p)),
        "repo_vs_production_sync_hash_mismatch": sorted(p for p in repo_hashes if repo_hashes[p] != plat["prodFileHashes"].get(p)),
        "raw_samples": sum(len(v) for v in raw.values()),
        "raw_per_file": {dt: len(raw[dt]) for dt, _ in FILES},
        "official_kept_per_lang": len(kept),
        "official_dropped": dropped,
        "production_items": len(prod_items),
        "only_official": sorted(official_keys - set(prod_items)),
        "only_production": sorted(set(prod_items) - official_keys),
        "production_dropped": plat["prodDropped"],
        "production_vs_rebuild_item_diffs": plat["itemsProdVsLocal"],
        "production_unused_fields": plat.get("staleFields", {}),
        "field_diffs": {k: v for k, v in field_diffs.items()},
        "fields_compared": ["messages (norm_session_zh/en)", "botProfile (dialogue_setting.bot_profile)", "annotation_score", "id"],
        "agreement": f"{len(official_keys & set(prod_items))}/{len(official_keys | set(prod_items))} samples; {len(norm) - sum(len(set(v)) for v in field_diffs.values())}/{len(norm)} with every field equal",
    }

    # ------------------------------------------------------------------ B. prompts
    by_key = {v["key"]: (sid, v) for sid, v in plat["items"].items()}
    rp_mismatch, code_rp_equal = [], 0
    for (l, dt, i), s in norm.items():
        k = f"{l}|{dt}|{i}"
        dialogue = "\n".join(f"{m.role}: {m.content}" for m in s.messages)
        paper = t11[l].format(character_profile=s.dialogue_setting["bot_profile"], dialogue_context=dialogue)
        code = code_rp["Role_Play_PROMPT_" + l.upper()].format(character_profile=s.dialogue_setting["bot_profile"], user_profile=s.dialogue_setting["user_profile"], dialogue=dialogue)
        if k not in by_key or by_key[k][1]["rp"] != sha1(paper):
            rp_mismatch.append(k)
        code_rp_equal += by_key.get(k, (None, {}))[1].get("rp") == sha1(code)
    # Judge input for a placeholder reply.
    jg_mismatch, jg_total = [], 0
    for l in ("zh", "en"):
        for dt, _ in FILES:
            rows = []
            for (d, i), s in kept.items():
                if d != dt:
                    continue
                s = copy.deepcopy(s)
                s["response_zh"] = "RESPONSE_X"
                s["translation_en"]["response_en"] = "RESPONSE_X"
                rows.append((f"{l}|{dt}|{i}", s))
            for k, (ins, _) in judge_inputs(proc, l, dt, rows).items():
                jg_total += 1
                if k not in by_key or by_key[k][1]["jg"] != sha1(ins):
                    jg_mismatch.append(k)
    report["B_prompts"] = {
        "roleplay_vs_paper_table11": {"compared": len(norm), "mismatch": len(rp_mismatch), "ids": rp_mismatch[:50]},
        "roleplay_vs_code_roleplay_prompt_py": {"compared": len(norm), "equal": code_rp_equal, "note": "intended: protocol 2.0 uses the paper's Table 11, not roleplay_prompt.py"},
        "judge_input_placeholder_reply": {"compared": jg_total, "mismatch": len(jg_mismatch), "ids": jg_mismatch[:50]},
        "templates_production_vs_repository": plat["templatesProdVsLocal"],
    }

    # ------------------------------------------------------------------ B (real replies) and E. scoring
    seed_key = {sid: v["key"] for sid, v in plat["items"].items()}
    runs = {}
    for rid, pr in plat["runs"].items():
        run = json.load(open(os.path.join(export, f"run_{rid}.json"), encoding="utf-8"))
        res = json.loads(run["result_json"])
        sessions = jl(os.path.join(export, f"sessions_{rid}.jsonl"))
        replies, judg, excluded, status = {}, {}, [], collections.Counter()
        finish = collections.Counter()
        for s in sessions:
            status[s["status"]] += 1
            msgs = json.loads(s["msgs"])
            ai = next((m for m in msgs if m["round"] == 2 and m["role"] == "ai"), None)
            k = seed_key[s["seed_id"]]
            if ai and ai["source"] == "target":
                replies[k] = ai["text"].strip()
                finish[ai.get("finish") or "none"] += 1
            elif ai:
                excluded.append(s["seed_id"])
            j = [x for x in json.loads(s["judg"]) if x["status"] == "done"]
            if j:
                judg[k] = (s["seed_id"], j[0], json.loads(s["result_json"]) if s["result_json"] else None)
        # Judge input with the reply each sample actually got.
        real_mismatch, real_total = [], 0
        for l in ("zh", "en"):
            for dt, _ in FILES:
                rows = []
                for (d, i), s in kept.items():
                    k = f"{l}|{dt}|{i}"
                    if d != dt or k not in replies:
                        continue
                    s = copy.deepcopy(s)
                    if l == "zh":
                        s["response_zh"] = replies[k]
                    else:
                        s["translation_en"]["response_en"] = replies[k]
                    rows.append((k, s))
                for k, (ins, _) in judge_inputs(proc, l, dt, rows).items():
                    real_total += 1
                    p = pr["replies"].get(by_key[k][0]) if k in by_key else None
                    if not p or p["jg"] != sha1(ins):
                        real_mismatch.append(k)
        # Scoring with the authors' reading of each output.
        per_sample_diff, crash, unsplit, out_of_range = [], [], [], collections.Counter()
        cols = {l: collections.defaultdict(list) for l in ("zh", "en")}
        golds = {l: collections.defaultdict(list) for l in ("zh", "en")}
        plat_sample = {}
        for k, (sid, j, sres) in judg.items():
            l, dt, i = k.split("|")
            texts = j["raw"].split(" | ")
            parsed = json.loads(j["parsed"])
            if len(texts) != 10:
                unsplit.append(sid)
                continue
            vals = []
            for t in texts:
                try:
                    vals.append(float(check(t)))
                except Exception as e:  # all_eval.py stops here (float() of a non-numeric first line)
                    crash.append({"seed": sid, "text": t[:80], "error": str(e)})
                    vals.append(None)
            mx_col = {"mc": 4, "fa": 4, "bck": 3, "acb": 4, "ach": 4, "bcpb": 4, "bcph": 3, "es": 4, "er": 4, "ms": 2, "mr": 2, "hl": 5, "eg": 5}[KEY[dt]]
            out_of_range[KEY[dt]] += sum(1 for v in vals if v is not None and not (1 <= v <= mx_col))
            if None in vals:
                continue
            score = sum(vals) / len(vals)
            cols[l][dt].append(score)
            golds[l][dt].append(float(check(str(kept[(dt, int(i))]["annotation_score"]))))
            plat_sample[k] = parsed["score"]
            if vals != [float(x) for x in parsed["samples"]] or abs(score - parsed["score"]) > 1e-12:
                per_sample_diff.append({"seed": sid, "official": vals, "platform": parsed["samples"]})
        # Unrounded: the column as the platform aggregates it from the stored session results (each session's combined
        # dimension score), against the authors' formula on the same outputs.
        sess_cols = {l: collections.defaultdict(list) for l in ("zh", "en")}
        for k, (sid, j, sres) in judg.items():
            l, dt, _ = k.split("|")
            if sres and not sres.get("excluded") and KEY[dt] in (sres.get("combined") or {}):
                sess_cols[l][dt].append(sres["combined"][KEY[dt]])
        # Protocol 2.0's aggregation, emulated on the same outputs: the sample's column score rounded to 4 decimals in the
        # judgement record, then to 2 decimals in the session result (half up), then averaged and rounded for display.
        half_up = lambda x, d: math.floor(x * 10 ** d + 0.5 + 1e-9) / 10 ** d
        v20 = {l: {} for l in ("zh", "en")}
        for l in ("zh", "en"):
            for dt, key in FILES:
                m = {"mc": 4, "fa": 4, "bck": 3, "acb": 4, "ach": 4, "bcpb": 4, "bcph": 3, "es": 4, "er": 4, "ms": 2, "mr": 2, "hl": 5, "eg": 5}[key]
                v20[l][key] = statistics.fmean(half_up(round(x / m * 5, 4), 2) for x in cols[l][dt])
            v20[l]["avg"] = statistics.fmean(v20[l][key] for _, key in FILES)
        official_cols, platform_cols, col_diff, platform_unrounded_avg = {}, {}, {}, {}
        for l in ("zh", "en"):
            official_cols[l], platform_cols[l] = {}, {}
            for dt, key in FILES:
                mx = max(golds[l][dt])
                official_cols[l][key] = statistics.fmean(cols[l][dt]) / mx * 5
                platform_cols[l][key] = res["perLang"][l]["dims"][key] if l in res["perLang"] else None
                col_diff[f"{l}.{key}"] = abs(official_cols[l][key] - platform_cols[l][key])
            unrounded = {key: statistics.fmean(sess_cols[l][dt]) for dt, key in FILES if sess_cols[l][dt]}
            platform_unrounded_avg[l] = statistics.fmean(unrounded.values()) if len(unrounded) == 13 else None
            col_diff_unrounded = max(abs(unrounded[key] - official_cols[l][key]) for _, key in FILES) if len(unrounded) == 13 else None
            official_cols[l]["_max_abs_diff_unrounded"] = col_diff_unrounded
            official_cols[l]["avg"] = statistics.fmean(official_cols[l][key] for _, key in FILES)
            platform_cols[l]["avg"] = res["perLang"][l]["overall"]
            col_diff[f"{l}.avg"] = abs(official_cols[l]["avg"] - platform_cols[l]["avg"])
        # The leaderboard shows two decimals: count the cells where the platform's value is not the official one rounded
        # (half away from zero, as the platform's round2).
        r2 = lambda x: math.floor(x * 100 + 0.5 + 1e-9) / 100
        cells = [(l, k) for l in ("zh", "en") for k in list(official_cols[l]) if not k.startswith("_")]
        mism = [f"{l}.{k}: official {official_cols[l][k]:.6f} -> {r2(official_cols[l][k]):.2f}, platform {platform_cols[l][k]}" for l, k in cells if abs(r2(official_cols[l][k]) - platform_cols[l][k]) > 1e-9]
        cells20 = [f"{l}.{k}" for l, k in cells if abs(r2(official_cols[l][k]) - half_up(v20[l][k], 2)) > 1e-9]
        unreadable = sum(json.loads(j["parsed"]).get("unreadable", 0) for _, j, _ in judg.values())
        runs[rid] = {
            "cells_compared": len(cells),
            "cells_differing_at_two_decimals": mism,
            "cells_differing_at_two_decimals_under_2_0_aggregation": cells20,
            "reply_finish_reasons": dict(finish),
            "unreadable_outputs_platform_count": unreadable,
            "model": run["target_model"],
            "protocol": pr["protocol"],
            "snapshot_prompts_vs_current_templates": pr["promptDiffs"],
            "target_params": pr["targetParams"],
            "sessions": dict(status),
            "excluded_pairs": sorted(excluded),
            "judge_input_real_replies": {"compared": real_total, "mismatch": len(real_mismatch), "ids": real_mismatch[:50]},
            "judged_samples": len(judg),
            "outputs_not_split_into_10": unsplit,
            "outputs_where_all_eval_stops": crash[:20],
            "outputs_where_all_eval_stops_count": len(crash),
            "out_of_range_outputs": dict(out_of_range),
            "per_sample_score_diffs": len(per_sample_diff),
            "per_sample_score_diff_examples": per_sample_diff[:10],
            "columns_official": official_cols,
            "columns_platform": platform_cols,
            "platform_unrounded_avg": platform_unrounded_avg,
            "max_abs_column_diff": max(col_diff.values()),
            "max_abs_column_diff_unrounded": {l: official_cols[l].pop("_max_abs_diff_unrounded") for l in ("zh", "en")},
            "max_abs_column_diff_at": max(col_diff, key=col_diff.get),
        }
    for l in ("zh", "en"):
        o = sorted(runs, key=lambda r: -runs[r]["columns_official"][l]["avg"])
        # Both unrounded: displayed values tie at two decimals, and tied rows have no order to compare.
        p = sorted(runs, key=lambda r: -runs[r]["platform_unrounded_avg"][l])
        report.setdefault("E_ranking", {})[l] = {"official": [runs[r]["model"] for r in o], "platform": [runs[r]["model"] for r in p], "same": o == p}
    report["runs"] = runs
    json.dump(report, open(out_path, "w", encoding="utf-8"), ensure_ascii=False, indent=1)
    a = report["A_data"]
    print(f"A data: {a['agreement']}; dropped {len(a['official_dropped'])}; hash mismatches {a['repo_vs_production_sync_hash_mismatch']}")
    b = report["B_prompts"]
    print(f"B prompts: role-play vs Table 11 {b['roleplay_vs_paper_table11']['mismatch']}/{b['roleplay_vs_paper_table11']['compared']} mismatches; judge input {b['judge_input_placeholder_reply']['mismatch']}/{b['judge_input_placeholder_reply']['compared']}")
    for rid, r in runs.items():
        print(f"{r['model']}: judge input (real replies) {r['judge_input_real_replies']['mismatch']}/{r['judge_input_real_replies']['compared']} mismatches; "
              f"per-sample score diffs {r['per_sample_score_diffs']}; all_eval stops on {r['outputs_where_all_eval_stops_count']} outputs; "
              f"max column diff {r['max_abs_column_diff']:.2e} ({r['max_abs_column_diff_at']}); cells differing at 2 decimals {len(r['cells_differing_at_two_decimals'])}/{r['cells_compared']}; unrounded {r['max_abs_column_diff_unrounded']}")
    for l, v in report["E_ranking"].items():
        print(f"ranking {l}: same={v['same']}")


if __name__ == "__main__":
    main(*sys.argv[1:6])