"""End-to-end check of a local CharacterEval run against the official logic (tools/experiments/charactereval_reference.py).

For every session of the run: the request the engine sent to the model under test must equal the reference message
array; the reward-model inputs it sent must equal the reference transform + format_input of the recorded reply, one
per assigned metric; the session's scores must equal the mock's scores of those inputs; and the run's metric means and
dimension averages must equal compute_score + E18 on the same scores (to the decimals shown).

    python3 check_ce_run.py <repo> <requests.jsonl> <base api url> <token file> <run id>
"""
import hashlib
import importlib.util
import json
import sys
import urllib.request

repo, log, base, token_file, run_id = sys.argv[1:6]
spec = importlib.util.spec_from_file_location('ref', '<platform>/tools/experiments/charactereval_reference.py')
ref = importlib.util.module_from_spec(spec)
spec.loader.exec_module(ref)

token = open(token_file).read().strip()


def get(path):
    req = urllib.request.Request(base + path, headers={'Cookie': f'rpb_session={token}'})
    return json.load(urllib.request.urlopen(req))


tests = {str(t['id']): t for t in json.load(open(f'{repo}/data/test_data.jsonl'))}
id2metric = json.load(open(f'{repo}/data/id2metric.jsonl'))
profiles = json.load(open(f'{repo}/data/character_profiles.json'))
reqs = [json.loads(l) for l in open(log, encoding='utf-8')]
chats = [r['body'] for r in reqs if r['path'].endswith('/chat/completions')]
reward_texts = [x['messages'][0]['content'] for r in reqs if r['path'].endswith('/rewards') for x in r['body']['inputs']]
mock = lambda t: 1 + (int(hashlib.sha256(t.encode()).hexdigest(), 16) % 400) / 100

run = get(f'/runs/{run_id}')
sessions = get(f'/runs/{run_id}/sessions')
bad = 0
records = []
for s in sessions:
    d = get(f'/sessions/{s["id"]}')
    sid = d['seed_id'][3:]
    data = tests[sid]
    want = ref.target_messages(data, profiles)
    sent = [c for c in chats if [{'role': m['role'], 'content': m['content']} for m in c['messages']] == want]
    reply = next(m for m in d['messages'] if m['round'] == 2)
    raw = reply.get('raw_text') or reply['text']
    rows = ref.transform([{**data, 'model_output': raw}], id2metric)
    texts = ref.rm_inputs(rows, profiles)
    missing = [t for t in texts if t not in reward_texts]
    scores = {r['metric_en']: mock(t) for r, t in zip(rows, texts)}
    keymap = {m: k for k, m in [('flu', 'Fluency'), ('coh', 'Coherence'), ('cons', 'Consistency'), ('ke', 'Exposure'), ('ka', 'Accuracy'), ('kh', 'Hallucination'), ('pb', 'Behavior'), ('pu', 'Utterance'), ('hl', 'Humanlikeness'), ('cs', 'Communication_skills'), ('ed', 'Diversity'), ('emp', 'Empathy')]}
    got = d['result']['combined']
    diff = [m for m, v in scores.items() if abs(got.get(keymap[m], float('nan')) - v) > 1e-12]
    one_line = all(chr(10) not in r['model_output'] for r in rows)
    ok = len(sent) >= 1 and not missing and not diff and len(got) == len(scores)
    bad += not ok
    print(('PASS' if ok else 'FAIL'), d['seed_id'], f'request={len(sent)} rm_inputs={len(texts)} missing={len(missing)} score_diff={diff} one_line={one_line}')
    for r, t in zip(rows, texts):
        records.append({**r, r['metric_en']: mock(t)})

means = ref.compute_score(records)
dims = ref.dimension_averages(means)
res = run['result']['perLang']['zh']
keymap = {'Fluency': 'flu', 'Coherence': 'coh', 'Consistency': 'cons', 'Exposure': 'ke', 'Accuracy': 'ka', 'Hallucination': 'kh', 'Behavior': 'pb', 'Utterance': 'pu', 'Humanlikeness': 'hl', 'Communication_skills': 'cs', 'Diversity': 'ed', 'Empathy': 'emp'}
for m, v in means.items():
    ok = abs(res['dims'][keymap[m]] - round(v, 3)) < 1e-9
    bad += not ok
    print(('PASS' if ok else 'FAIL'), 'metric', m, round(v, 3), res['dims'][keymap[m]])
for g, v in dims.items():
    if v is None:
        continue
    ok = abs(res['groups'][g] - round(v, 3)) < 1e-9
    bad += not ok
    print(('PASS' if ok else 'FAIL'), 'dimension', g, round(v, 3), res['groups'][g])
print('overall (character consistency)', res['overall'], 'reference', None if dims['cc'] is None else round(dims['cc'], 3))
print('ALL PASS' if not bad else f'{bad} FAIL')
sys.exit(1 if bad else 0)