"""CharacterEval: Python reference implementation of the official pipeline, and the generators of the platform's constants.

Source: github.com/morecry/CharacterEval at commit c3d44a6 (docs/复现规范/CharacterEval-核查矩阵.md, rows E01–E20).
The blocks between `# >>> verbatim` and `# <<< verbatim` are the repository's own source text, pasted mechanically from
the source segments (Python ast) of the named files; charactereval_spec_check.py compares every block with the
repository again. Around them, only the I/O is replaced: the official scripts read and write fixed files and load models,
here the same logic takes its inputs as arguments. No module of the repository is imported or run.

    python3 -I tools/experiments/charactereval_reference.py generate <repo>

writes worker/src/benchmarks/charactereval/strings.gen.ts (every string literal the platform needs, extracted with
Python ast from get_response.py, run_char_rm.py, transform_format.py, and the metric names from data/id2metric.jsonl)
and profiles.gen.json (the Python `str(profile)` of every character in data/character_profiles.json).
"""
import ast
import copy
import hashlib
import json
import os
import sys

ROOT = os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", ".."))
GEN_DIR = os.path.join(ROOT, "worker", "src", "benchmarks", "charactereval")
STRINGS_GEN = os.path.join(GEN_DIR, "strings.gen.ts")
PROFILES_GEN = os.path.join(GEN_DIR, "profiles.gen.json")

REPO_URL = "https://github.com/morecry/CharacterEval"
COMMIT = "c3d44a6fc1790cc8c4b2fd7c01f0c72930655e0c"
SHA256 = {
    "get_response.py": "ac5030abbefc5d9979e65a2d5f47490cfac5c7ca5195c7737124ae78fed7a9e5",
    "transform_format.py": "c6cfe4a42338147878cc37359d34e80b851bafd0521b6b85bc171090977c8b2c",
    "run_char_rm.py": "85bc7ebdf6f283dc563fe7bfb8b98cdf7799b9cbb5a9c443b78f3e0c742180be",
    "compute_score.py": "cfcf291f85ee759c114aa583b76934a35e71799f4c5b06eef0ab7a28b0d031c0",
    "data/test_data.jsonl": "6d9723a1484f1d2739bcf388c7960abe1933b7a5a1fddae0170cc9ea7c66dcd8",
    "data/id2metric.jsonl": "df80ee716124a7deacc6f24ee372a4600b719cdaa79f01cefcdfb286359f1801",
    "data/character_profiles.json": "d9b4c9f7ab63b89b3c348ae88b43c87872f0f6a5cb06d7f31cbe64a630e6434e",
}
# Matrix E06/V02: SHA-256 of the profile reprs sorted by role name and joined as "role\trepr\n".
PROFILES_SORTED_SHA256 = "198d40264a5109b1ed974da4274fa5d1bc711f01df4d5c62c5af53bb14d370c2"

# Matrix E18: the paper's dimension averages are unweighted means of these metrics' means (no code defines them).
GROUPS = {
    "ca": ["Fluency", "Coherence", "Consistency"],
    "cc": ["Exposure", "Accuracy", "Hallucination", "Behavior", "Utterance"],
    "ra": ["Humanlikeness", "Communication_skills", "Diversity", "Empathy"],
}


# ---------------------------------------------------------------------------------------------------------------------
# Official logic. get_response.py (generation messages).

# >>> verbatim get_response.py def:concat_messages all0
def concat_messages(conversations, role, system):
    history = []
    first_query = system
    if conversations[0]['from'] == role:
        first_response = f"好的！现在我来扮演{role}。" + "我首先发话：" + conversations[0]['value']
    else:
        first_response = f"好的！现在我来扮演{role}。"
    
    history.append({"role": "user", "content": first_query})
    history.append({"role": "assistant", "content": first_response})
    
    for i in range(len(conversations)):
        if conversations[i]['from'] == role:
            if i ==0:
                continue
            else:
                assert conversations[i-1]['from'] != role
                query = f"{conversations[i-1]['from']}：" + conversations[i-1]['value']
                response = f"{conversations[i]['from']}：" + conversations[i]['value']
            history.append({"role": "user", "content": query})
            history.append({"role": "assistant", "content": response})
    assert conversations[-1]['from'] != role
    
    query = f"{conversations[-1]['from']}：" + conversations[-1]['value']
    return history, query
# <<< verbatim

# >>> verbatim get_response.py def:make_inputs all0
def make_inputs(context):
    dialogues= context.split('\n') 
    inputs = []  
    for dial in dialogues:
        role = dial.split("：")[0]
        dial = "：".join(dial.split("：")[1:])
        inputs.append({"from":role,"value":dial})
    return inputs
# <<< verbatim


def target_messages(data, role_informations):
    """get_response_chatglm without the model: the history it passes to model.chat, then the query as the last user
    message (ChatGLM3's chat appends {"role": "user", "content": query} to the history)."""
    context = data['context']
    role = data['role']

    role_information = role_informations[role]
    # >>> verbatim get_response.py def:get_response_chatglm/assign:role_system first4
    role_system = f'''{role_information}
现在请你扮演一个角色扮演专家。请你根据上述信息扮演{role}进行对话。
'''
    # <<< verbatim

    # >>> verbatim get_response.py def:get_response_chatglm/assign:messages,query first4
    messages,query = concat_messages(make_inputs(context), role, role_system)
    # <<< verbatim
    return messages + [{"role": "user", "content": query}]


# transform_format.py (failure filter, first line, one row per assigned metric).
def transform(datas, id_metric):
    datas = copy.deepcopy(datas)
    results = []

    # >>> verbatim transform_format.py for:datas all4
    for data in datas:
        if data['model_output'] is not None and data['model_output'] != "ERROR":
            model_output = data['model_output'].split("\n")[0] # Prevent continuous generation
            data['model_output'] = model_output
            if str(data['id']) in id_metric:
                for x in id_metric[str(data['id'])]:
                    data['metric_en']= x[0]
                    data['metric_zh']= x[1]
                    tmp = copy.deepcopy(data)
                    results.append(tmp)
    # <<< verbatim
    return results


# run_char_rm.py (reward-model input text). format_input reads the module-level dict character_profile.
character_profile = {}

# >>> verbatim run_char_rm.py def:format_input all0
def format_input(example):
    input_text = "<RoleInfo>\n\n" \
        + str(character_profile[example['role']]) + "\n\n<Context>\n\n" + example['context'] + "\n\n<Response>\n\n" + example['model_output'] + "\n\n<Dimension>\n\n" + example["metric_zh"]
    return input_text
# <<< verbatim


def rm_inputs(records, profiles):
    global character_profile
    character_profile = profiles
    return [format_input(r) for r in records]


# compute_score.py (per-metric arithmetic mean over all records; records carry the score under their metric_en).
def compute_score(records):
    score_dict = {}
    # >>> verbatim compute_score.py for:records all4
    for record in records:
        if record['metric_en'] not in score_dict:
            score_dict[record['metric_en']] = []
        score_dict[record['metric_en']].append(record[record['metric_en']])
    # <<< verbatim
    means = {}
    for metric in score_dict:
        means[metric] = sum(score_dict[metric]) / len(score_dict[metric])  # AST-checked against compute_score.py's print
    return means


def dimension_averages(means):
    """Matrix E18 (platform reading of the paper's Avg. columns): unweighted mean of the group's metric means, None when
    one of them has no record."""
    out = {}
    for g, ms in GROUPS.items():
        out[g] = None if any(m not in means for m in ms) else sum(means[m] for m in ms) / len(ms)
    return out


# ---------------------------------------------------------------------------------------------------------------------
# Locating source segments (used to paste and to check the verbatim blocks, and by the generator).

def read(path):
    return open(path, encoding="utf-8").read()


def sha256_file(path):
    return hashlib.sha256(open(path, "rb").read()).hexdigest()


def locate(tree, locator):
    """def:NAME, def:NAME/assign:TARGET (TARGET is a name or comma-joined tuple of names), for:ITER (module level)."""
    node = tree
    for step in locator.split("/"):
        kind, name = step.split(":", 1)
        found = None
        for n in node.body:
            if kind == "def" and isinstance(n, ast.FunctionDef) and n.name == name:
                found = n
            elif kind == "for" and isinstance(n, ast.For) and isinstance(n.iter, ast.Name) and n.iter.id == name:
                found = n
            elif kind == "assign" and isinstance(n, ast.Assign) and len(n.targets) == 1:
                t = n.targets[0]
                names = [t.id] if isinstance(t, ast.Name) else [e.id for e in getattr(t, "elts", []) if isinstance(e, ast.Name)]
                if ",".join(names) == name:
                    found = n
        if found is None:
            raise LookupError(f"{locator}: {step} not found")
        node = found
    return node


def source_segment(repo, file, locator):
    src = read(os.path.join(repo, file))
    return ast.get_source_segment(src, locate(ast.parse(src), locator))


def expected_block(segment, mode):
    """mode allN: every line indented by N; firstN: only the first line (the rest are inside a string literal)."""
    n = int(mode[-1])
    lines = segment.split("\n")
    if mode.startswith("all"):
        return "\n".join((" " * n + l) if l else l for l in lines)
    return "\n".join([" " * n + lines[0]] + lines[1:])


def verbatim_blocks(path=os.path.abspath(__file__)):
    """[(file, locator, mode, text)] for every marked block of this file."""
    out, cur, buf = [], None, []
    for line in read(path).split("\n"):
        s = line.strip()
        if s.startswith("# >>> verbatim "):
            cur, buf = s[len("# >>> verbatim "):].split(" "), []
        elif s == "# <<< verbatim" and cur:
            out.append((cur[0], cur[1], cur[2], "\n".join(buf)))
            cur = None
        elif cur is not None:
            buf.append(line)
    return out


# ---------------------------------------------------------------------------------------------------------------------
# Generators.

def segments(node):
    """A string expression as a list of literal and expression parts, in source order: f-string pieces, '+' operands."""
    if isinstance(node, ast.BinOp) and isinstance(node.op, ast.Add):
        return segments(node.left) + segments(node.right)
    if isinstance(node, ast.JoinedStr):
        out = []
        for v in node.values:
            if isinstance(v, ast.Constant):
                out.append({"lit": v.value})
            else:
                assert isinstance(v, ast.FormattedValue) and v.conversion == -1 and v.format_spec is None
                out.append({"expr": ast.unparse(v.value)})
        return out
    if isinstance(node, ast.Constant) and isinstance(node.value, str):
        return [{"lit": node.value}]
    return [{"expr": ast.unparse(node)}]


def assigns(fn, name):
    """Values assigned to `name` inside fn, in source order."""
    hits = [n for n in ast.walk(fn) if isinstance(n, ast.Assign) and len(n.targets) == 1
            and isinstance(n.targets[0], ast.Name) and n.targets[0].id == name]
    return [n.value for n in sorted(hits, key=lambda n: (n.lineno, n.col_offset))]


def method_arg(fn, method, receiver=None):
    """The single string argument of every `x.method("...")` call in fn (receiver: the constant receiver, for join)."""
    out = []
    for n in ast.walk(fn):
        if isinstance(n, ast.Call) and isinstance(n.func, ast.Attribute) and n.func.attr == method:
            if receiver is not None:
                if isinstance(n.func.value, ast.Constant):
                    out.append(n.func.value.value)
            elif n.args and isinstance(n.args[0], ast.Constant):
                out.append(n.args[0].value)
    return out


def one(values, what):
    key = lambda v: ast.dump(v) if isinstance(v, ast.AST) else json.dumps(v)
    if len(set(map(key, values))) != 1:
        raise SystemExit(f"{what}: expected one value, found {values!r}")
    return values[0]


def extract_strings(repo):
    gr = ast.parse(read(os.path.join(repo, "get_response.py")))
    concat = locate(gr, "def:concat_messages")
    make = locate(gr, "def:make_inputs")
    chat = locate(gr, "def:get_response_chatglm")
    first = assigns(concat, "first_response")  # [character speaks first, otherwise]
    query = assigns(concat, "query")  # [turn in the loop, final query]
    tf = ast.parse(read(os.path.join(repo, "transform_format.py")))
    loop = locate(tf, "for:datas")
    cond = loop.body[0].test  # data['model_output'] is not None and data['model_output'] != "ERROR"
    sentinel = [c.comparators[0].value for c in cond.values if isinstance(c, ast.Compare) and isinstance(c.ops[0], ast.NotEq)]
    rm = ast.parse(read(os.path.join(repo, "run_char_rm.py")))
    fmt = locate(rm, "def:format_input")
    return {
        "SYSTEM": ("get_response.py get_response_chatglm: role_system", segments(one(assigns(chat, "role_system"), "role_system"))),
        "FIRST_RESPONSE_ROLE_FIRST": ("get_response.py concat_messages: first_response when conversations[0]['from'] == role", segments(first[0])),
        "FIRST_RESPONSE": ("get_response.py concat_messages: first_response otherwise", segments(first[1])),
        "TURN_QUERY": ("get_response.py concat_messages: query in the loop", segments(query[0])),
        "TURN_RESPONSE": ("get_response.py concat_messages: response in the loop", segments(one(assigns(concat, "response"), "response"))),
        "FINAL_QUERY": ("get_response.py concat_messages: final query", segments(query[1])),
        "CONTEXT_LINE_SEP": ("get_response.py make_inputs: context.split(...)", one([a for a in method_arg(make, "split") if a == "\n"], "line sep")),
        "SPEAKER_SEP": ("get_response.py make_inputs: dial.split(...)", one([a for a in method_arg(make, "split") if a != "\n"], "speaker sep")),
        "SPEAKER_JOIN": ("get_response.py make_inputs: \"...\".join(...)", one(method_arg(make, "join", receiver=True), "join")),
        "FAILURE_SENTINEL": ("transform_format.py: data['model_output'] != ...", one(sentinel, "sentinel")),
        "FIRST_LINE_SEP": ("transform_format.py: data['model_output'].split(...)[0]", one(method_arg(loop, "split"), "first line sep")),
        "RM_INPUT": ("run_char_rm.py format_input: input_text", segments(one(assigns(fmt, "input_text"), "input_text"))),
    }


def metrics_from_data(repo):
    """The distinct [metric_en, metric_zh] pairs of id2metric, in order of first appearance."""
    seen = []
    for pairs in json.load(open(os.path.join(repo, "data", "id2metric.jsonl"), encoding="utf-8")).values():
        for p in pairs:
            if list(p) not in seen:
                seen.append(list(p))
    return seen


def strings_gen(repo):
    files = ["get_response.py", "transform_format.py", "run_char_rm.py", "data/id2metric.jsonl"]
    js = lambda v: json.dumps(v, ensure_ascii=False)
    out = [
        "// GENERATED by tools/experiments/charactereval_reference.py generate <repo>; do not edit by hand.",
        f"// Every string below is extracted with Python ast from {REPO_URL} at commit {COMMIT}",
        "// (or read from data/id2metric.jsonl); a template is its literal and expression parts in source order.",
        "",
        "export type Seg = { lit: string } | { expr: string };",
        "",
        "export const SOURCE = {",
        f"  repo: {js(REPO_URL)},",
        f"  commit: {js(COMMIT)},",
        "  sha256: {",
    ]
    for f in files:
        out.append(f"    {js(f)}: {js(sha256_file(os.path.join(repo, f)))},")
    out += ["  },", "} as const;", ""]
    for name, (where, value) in extract_strings(repo).items():
        out.append(f"/** {where} */")
        out.append(f"export const {name}{': Seg[]' if isinstance(value, list) else ''} = {js(value)};")
    out.append("/** data/id2metric.jsonl: distinct [metric_en, metric_zh] pairs, in order of first appearance. */")
    out.append(f"export const METRICS_FROM_DATA: [string, string][] = {js(metrics_from_data(repo))};")
    return "\n".join(out) + "\n"


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


def profile_reprs(repo):
    return {role: str(p) for role, p in load_profiles(repo).items()}


def sorted_profiles_hash(reprs):
    return hashlib.sha256("".join(f"{r}\t{reprs[r]}\n" for r in sorted(reprs)).encode("utf-8")).hexdigest()


def profiles_gen(repo):
    reprs = profile_reprs(repo)
    doc = {
        "generated": "tools/experiments/charactereval_reference.py generate <repo>; do not edit by hand. "
                     "Each value is Python str(profile), which get_response.py and run_char_rm.py insert as the role information.",
        "source": {"repo": REPO_URL, "commit": COMMIT, "file": "data/character_profiles.json",
                   "sha256": sha256_file(os.path.join(repo, "data", "character_profiles.json"))},
        "sortedReprSha256": sorted_profiles_hash(reprs),
        "profiles": reprs,
    }
    return json.dumps(doc, ensure_ascii=False, indent=1) + "\n"


def main():
    if len(sys.argv) != 3 or sys.argv[1] != "generate":
        raise SystemExit(__doc__)
    repo = sys.argv[2]
    for f in ("get_response.py", "transform_format.py", "run_char_rm.py", "data/id2metric.jsonl", "data/character_profiles.json"):
        if sha256_file(os.path.join(repo, f)) != SHA256[f]:
            raise SystemExit(f"{f}: SHA-256 differs from commit {COMMIT[:7]}")
    open(STRINGS_GEN, "w", encoding="utf-8").write(strings_gen(repo))
    open(PROFILES_GEN, "w", encoding="utf-8").write(profiles_gen(repo))
    print(f"wrote {STRINGS_GEN}\nwrote {PROFILES_GEN}")


if __name__ == "__main__":
    main()