"""CoSER: compares the prompts printed in the paper's LaTeX tables with the strings the platform sends.

The paper's prompts are read from Tables/prompts_agent.tex, prompts_agent_2.tex, prompts_eval.tex and prompts_eval_2.tex of
the arXiv e-print; the platform's from worker/src/benchmarks/coser/gca/strings.gen.ts (itself generated from the official
code with Python ast and compared with a fresh extraction by coser_spec_check.py G01). Neither side is typed by hand.

LaTeX markup is removed by fixed rules; line breaks and spaces are collapsed, because the paper's table cells do not keep
them; the placeholder names of the paper and the code are kept as written. The two token sequences are compared with
difflib and every difference is printed. Each difference must be listed in EXPECTED with what it means; a difference not
listed there, or an entry of EXPECTED that no longer occurs, fails the check. Annotations of the paper's table such as
"(if available)" and "(for CoSER models)" are not prompt text and are removed; they are listed in ANNOTATIONS and each must
occur. Exit 1 on any failure. No model is called.

    python3 -I tools/experiments/coser_prompt_check.py <e-print dir of 2502.09082v3> [--json out.json]

The report reads its Appendix A from tools/experiments/coser_prompt_check.result.json, written with --json.
"""
import argparse
import difflib
import hashlib
import json
import os
import re
import subprocess
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
STRINGS = os.path.join(HERE, "..", "..", "worker", "src", "benchmarks", "coser", "gca", "strings.gen.ts")

# Sources of the check (arXiv:2502.09082v3 e-print).
TABLE_SHA256 = {
    "Tables/prompts_agent.tex": "54fc52663e9c",
    "Tables/prompts_agent_2.tex": "c95d562ecefc",
    "Tables/prompts_eval.tex": "ba199552a708",
    "Tables/prompts_eval_2.tex": "7b36a67c86b1",
}

# The current official code's next-speaker string: compared for the record, not sent and not counted in the totals.
NOT_SENT = "next speaker (current code's string, 624efc5, not sent)"

ANNOTATIONS = [
    r"\textit{\textbf{(if available)}}",
    r"\textbf{(if  retrieval augmented)}",
    r"\textit{\textbf{(for CoSER models)}}",
    r"\textit{\textbf{(for other models, with output examples)}}",
    r"\textbf{\textit{(if k$>$0)}}",
]


def read(root, path):
    raw = open(os.path.join(root, path), "rb").read()
    want = TABLE_SHA256[path]
    got = hashlib.sha256(raw).hexdigest()
    if not got.startswith(want):
        raise SystemExit(f"{path}: SHA-256 {got} does not start with {want}")
    return raw.decode("utf-8")


def cell(tex, label, end=r"\\"):
    """The text of the table cell that follows `label &`, up to the row end (a line holding only `\\\\`, with or without a rule)."""
    i = tex.index(label)
    start = tex.index("&", i) + 1
    m = re.compile(r"\n\s*\\\\\s*(\\midrule|\\hline|\n|$)").search(tex, start)
    return tex[start:m.start()]


def detex(s, found):
    for a in ANNOTATIONS:
        if a in s:
            found.add(a)
            s = s.replace(a, " ")
    s = re.sub(r"(?m)(?<!\\)%.*$", "", s)
    s = s.replace(r"\qquad", " ").replace(r"\quad", " ")
    s = s.replace(r"\{", "{").replace(r"\}", "}").replace(r"\_", "_").replace(r"\#", "#").replace(r"\&", "&").replace(r"\%", "%")
    s = s.replace("$<$", "<").replace("$>$", ">")
    s = re.sub(r"\\text(bf|it)\{([^{}]*)\}", r"\2", s)
    s = re.sub(r"\\text(bf|it)\{([^{}]*)\}", r"\2", s)
    return s


def tokens(s):
    return s.split()


def compose(parts, values=None):
    out = []
    for p in parts:
        if isinstance(p, str):
            out.append(p)
        else:
            out.append((values or {}).get(p["v"], "{" + p["v"] + "}"))
    return "".join(out)


def platform_strings():
    code = (
        "import('" + STRINGS.replace("\\", "\\\\") + "').then((m) => { process.stdout.write(JSON.stringify(m)); });"
    )
    r = subprocess.run(["node", "--experimental-strip-types", "--no-warnings", "-e", code], capture_output=True, text=True, check=True)
    return json.loads(r.stdout)


def pairs(root):
    agent, agent2 = read(root, "Tables/prompts_agent.tex"), read(root, "Tables/prompts_agent_2.tex")
    ev, ev2 = read(root, "Tables/prompts_eval.tex"), read(root, "Tables/prompts_eval_2.tex")
    st = platform_strings()
    found = set()
    cp, ep, np_, np_paper, j = st["CHARACTER_PROMPT"], st["ENVIRONMENT_PROMPT"], st["NSP_PROMPT"], st["PAPER_ERA_NSP_PROMPT"], st["JUDGE"]
    inference = cell(agent, r"\textbf{Role-playing Instruction \quad (Fixed Template for Inference)}")
    # The platform's template in the paper's reading order: the optional sections written out, the output format of non-CoSER models.
    char = compose(cp["prompt"], {"other_character_profiles_str": compose(cp["othersSection"]), "motivation": compose(cp["motivationSection"]) + "===Relevant Background Information===\n{retrieved_knowledge}\n\n"})
    char += compose(cp["requirements"], {"output_format": cp["outputFormat"]["thought"] + " " + cp["outputFormat"]["thoughtWithExample"]})
    out = [("character template", detex(inference, found), char)]
    # The identity phrase is drawn at random in the code; the paper prints one of the four.
    env = compose(ep["prompt"], {"random.choice(model_roles)": ep["modelRoles"][3], "major_characters": "{major_characters}"})
    out.append(("environment", detex(cell(agent2, r"\textbf{Environment Model}"), found), env))
    nsp_vals = {"ENVIRONMENT": ep["environment"]}
    nsp_cell = detex(cell(agent2, r"\textbf{Next Sentence Prediction}"), found)
    out.append(("next speaker (the string sent, 93c1474)", nsp_cell, compose(np_paper["prompt"], nsp_vals)))
    out.append((NOT_SENT, nsp_cell, compose(np_["prompt"], nsp_vals)))
    template = j["template"].replace("{additional_instructions}", j["additionalInstructions"])
    out.append(("judge template", detex(cell(ev, r"\textbf{Template}"), found), template))
    labels = {"Anthropomorphism": r"\textbf{Anthropo-}\textbf{morphism}", "Character Fidelity": r"\textbf{Character \qquad Fidelity}", "Storyline Quality": None, "Storyline Consistency": None}
    for dim, label in labels.items():
        if label is None:
            m = re.search(r"\\textbf\{Storyline[^&]*" + dim.split()[1][:4] + r"[^&]*&", ev2)
            start = m.group(0)[:-1].rstrip()
        else:
            start = label
        text = detex(cell(ev2, start), found)
        m = re.match(r"\s*\(intro\)\s*(.*?)\s*\(rubrics\)\s*(.*)$", text, re.S)
        if not m:
            raise SystemExit(f"{dim}: no (intro)/(rubrics) in the cell: {text[:200]!r}")
        out.append((f"{dim}: intro", m.group(1), j["details"][dim]["dimension_brief"]))
        out.append((f"{dim}: rubrics", m.group(2), j["details"][dim]["dimension_criteria"]))
    return out, found


def diffs(a, b):
    """Differences between the token lists: (paper tokens, platform tokens), with a little context."""
    sm = difflib.SequenceMatcher(None, a, b, autojunk=False)
    out = []
    same = 0
    for op, i1, i2, j1, j2 in sm.get_opcodes():
        if op == "equal":
            same += i2 - i1
        else:
            out.append((" ".join(a[i1:i2]), " ".join(b[j1:j2])))
    return out, same


# Every difference between the paper's table and the platform's string, and what it is.
EXPECTED = [
    {"prompt": NOT_SENT, "paper": "",
     "platform": "===Output Format=== You must structure your response as follows: 1. Reasoning: Explain your thought process. 2. Next Speaker: Select exactly one character name from the list: {all_characters}.",
     "meaning": "added to the official code in December 2025 (958f962, “add support for thinking models”); not in the paper's table and not sent by protocol 1.1 (matrix C12)"},
    {"prompt": "judge template", "paper": "{dimension_intro}.", "platform": "{dimension_brief}.", "meaning": "placeholder name only"},
    {"prompt": "judge template", "paper": "Please note that the first {k} messages in the simulated conversation are the same as the reference. Focus your evaluation only on the content after these messages.", "platform": "",
     "meaning": "printed for k > 0 only; the main results and the platform use k = 0, where the code's additional_instructions is empty"},
    {"prompt": "judge template", "paper": "{dimension_rubrics}", "platform": "{dimension_criteria}", "meaning": "placeholder name only"},
    {"prompt": "judge template", "paper": "},},", "platform": "}, },", "meaning": "space inside the JSON example, not kept by the table"},
]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("eprint")
    ap.add_argument("--json")
    args = ap.parse_args()
    ps, found = pairs(args.eprint)
    missing = [a for a in ANNOTATIONS if a not in found]
    seen, unexpected, report = set(), [], []
    total_paper, total_same = 0, 0
    for name, paper, plat in ps:
        a, b = tokens(paper), tokens(plat)
        ds, same = diffs(a, b)
        if name != NOT_SENT:
            total_paper += len(a)
            total_same += same
        report.append({"prompt": name, "paper_tokens": len(a), "platform_tokens": len(b), "equal_tokens": same, "platform_sha256": hashlib.sha256(plat.encode("utf-8")).hexdigest(), "differences": [{"paper": x, "platform": y} for x, y in ds]})
        print(f"{name}: paper {len(a)} tokens, platform {len(b)}, {same} equal, {len(ds)} differences")
        for x, y in ds:
            e = next((e for e in EXPECTED if e["prompt"] == name and e["paper"] == x and e["platform"] == y), None)
            if e:
                seen.add((name, x, y))
                print(f"   listed: paper {x!r} | platform {y[:160]!r} | {e['meaning']}")
            else:
                unexpected.append((name, x, y))
                print(f"   NOT LISTED: paper {x!r} | platform {y[:400]!r}")
    stale = [e for e in EXPECTED if (e["prompt"], e["paper"], e["platform"]) not in seen]
    for e in stale:
        print(f"   STALE entry of the expected list: {e['prompt']}: {e['paper']!r} | {e['platform']!r}")
    if missing:
        print(f"annotations not found in the tables: {missing}")
    ok = not unexpected and not stale and not missing
    sent = [r for r in report if r["prompt"] != NOT_SENT]
    print(f"{'PASS' if ok else 'FAIL'} P-PROMPTS {len(sent)} prompts sent, {total_same}/{total_paper} paper tokens equal to the platform's in order; "
          f"{sum(len(r['differences']) for r in sent)} differences in them, {len(unexpected)} not listed, {len(stale)} stale entries")
    if args.json:
        with open(args.json, "w", encoding="utf-8") as f:
            json.dump({"tables_sha256_prefix": TABLE_SHA256, "prompts": report, "expected": EXPECTED}, f, ensure_ascii=False, indent=1)
    return 0 if ok else 1


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