"""CharacterBench: check that the platform builds the same prompts as the paper and the authors' code.

The role-play prompt is the paper's Table 11 (taken from the LaTeX source by characterbench_spec_check.py table11) with
the sample's profile and its dialogue context written out as the official process.py writes it (norm_session_zh/en,
"{role}: {content}" per message). The judge input is built by construct_prompts/process_wo_context_{zh,en}_all.py
process. Writes the SHA-1 of every prompt; compare with the platform's hashes (characterbench_prompt_hashes.ts):

    git clone https://github.com/thu-coai/CharacterBench /tmp/cb && git -C /tmp/cb checkout 66c113c
    python3 -I tools/experiments/characterbench_spec_check.py table11 eprint > table11.json
    python3 -I tools/experiments/characterbench_prompt_check.py /tmp/cb data/benchmarks/characterbench table11.json py.json
    node --disable-warning=ExperimentalWarning tools/experiments/characterbench_prompt_hashes.ts ts.json
    python3 -I tools/experiments/characterbench_prompt_check.py --compare py.json ts.json

The official functions are read from the repository by source segment (no module of the repository is imported). The
only stand-in is pydantic's Session.model_validate, replaced by the checks it makes on this data (message role is one of
the literals, content is a str). The judge input is built with the reply "RESPONSE_X" as a plain string.
"""
import ast
import copy
import hashlib
import json
import os
import sys
import tempfile
from types import SimpleNamespace


def source_of(path, names):
    src = open(path, encoding="utf-8").read()
    tree = ast.parse(src)
    return "\n\n".join(ast.get_source_segment(src, n) for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in names)


def constants(path):
    out = {}
    for node in ast.parse(open(path, encoding="utf-8").read()).body:
        if isinstance(node, ast.Assign) and isinstance(node.value, ast.Constant) and isinstance(node.value.value, str):
            out[node.targets[0].id] = node.value.value
    return out


def validate_session(d):
    if not isinstance(d["dialogue_setting"], dict):
        raise ValueError("dialogue_setting")
    msgs = []
    for m in d["messages"]:
        if m["role"] not in ("user", "assistant", "system") or not isinstance(m["content"], str):
            raise ValueError("message")
        msgs.append(SimpleNamespace(role=m["role"], content=m["content"]))
    return SimpleNamespace(dialogue_setting=d["dialogue_setting"], messages=msgs)


def data_info(repo, data, lang):
    """construct_prompts/data_info_{lang}.py, with its imports resolved to the downloaded prompt files."""
    tree = ast.parse(open(f"{repo}/construct_prompts/data_info_{lang}.py", encoding="utf-8").read())
    env = {}
    for n in tree.body:
        if isinstance(n, ast.ImportFrom):
            c = constants(f"{data}/construct_prompts/eval_prompts_{lang}/{n.module.split('.')[-1]}.py")
            for a in n.names:
                env[a.asname or a.name] = c.get(a.name, "")
    for n in tree.body:
        if isinstance(n, ast.Assign):
            env[n.targets[0].id] = eval(compile(ast.Expression(n.value), "data_info", "eval"), {}, env)
    return env["data_info"]


def build(repo, data, table11_path, out_path):
    norm = {}
    exec(source_of(f"{repo}/process.py", {"norm_session_zh", "norm_session_en"}), norm)
    rp = json.load(open(table11_path, encoding="utf-8"))
    h = lambda s: hashlib.sha1(s.encode()).hexdigest()
    out = {}
    tmp = tempfile.mkdtemp()
    for lang in ["zh", "en"]:
        ns = {"json": json, "deepcopy": copy.deepcopy, "data_info": data_info(repo, data, lang)}
        exec(source_of(f"{repo}/construct_prompts/process_wo_context_{lang}_all.py", {"process"}), ns)
        for f in sorted(os.listdir(f"{data}/eval_data/raw_data")):
            dt = f[: -len("_test.json")]
            rows = json.load(open(f"{data}/eval_data/raw_data/{f}", encoding="utf-8"))
            kept = []
            for i, session in enumerate(rows):
                # process.py normalizes both languages in one try block: a failure in either drops the sample.
                try:
                    szh = validate_session(norm["norm_session_zh"](copy.deepcopy(session)))
                    sen = validate_session(norm["norm_session_en"](copy.deepcopy(session)))
                except Exception:
                    continue
                s = szh if lang == "zh" else sen
                prompt = rp[lang].format(
                    character_profile=s.dialogue_setting["bot_profile"],
                    dialogue_context="\n".join(f"{m.role}: {m.content}" for m in s.messages),
                )
                session = copy.deepcopy(session)
                session["response_zh"] = "RESPONSE_X"
                session["translation_en"]["response_en"] = "RESPONSE_X"
                kept.append((i, prompt, session))
            path = f"{tmp}/{f}"
            json.dump([k[2] for k in kept], open(path, "w", encoding="utf-8"), ensure_ascii=False)
            outs, _ = ns["process"](path, dt)
            for j, (i, prompt, _) in enumerate(kept):
                out[f"{lang}|{dt}|{i}"] = {"rp": h(prompt), "jg": h(outs[j]["instruction"])}
    json.dump(out, open(out_path, "w"))
    print(f"{len(out)} prompts written to {out_path}")


def compare(a_path, b_path):
    a, b = json.load(open(a_path)), json.load(open(b_path))
    only_a, only_b = sorted(set(a) - set(b)), sorted(set(b) - set(a))
    rp = [k for k in a if k in b and a[k]["rp"] != b[k]["rp"]]
    jg = [k for k in a if k in b and a[k]["jg"] != b[k]["jg"]]
    print(f"samples: {len(a)} / {len(b)}; only in first: {only_a[:5]}; only in second: {only_b[:5]}")
    print(f"role-play prompt mismatches: {len(rp)} {rp[:5]}")
    print(f"judge prompt mismatches: {len(jg)} {jg[:5]}")
    sys.exit(1 if only_a or only_b or rp or jg else 0)


if __name__ == "__main__":
    if sys.argv[1] == "--compare":
        compare(sys.argv[2], sys.argv[3])
    else:
        build(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])