"""BaichuanCharRM deployment check (CharacterEval matrix V06, E13–E16), run by the operator of the reward-model endpoint.

Three checks, each reported in one JSON document (send it back unchanged):

  A  weights   every file of the model directory against Hugging Face morecry/BaichuanCharRM at commit
               e953da5f57b86db7adca3a2714b4c3720f5b96a2: SHA-256 for the LFS files (weights, tokenizer.model), git blob
               SHA-1 for the others (config, tokenizer config, and the model's own Python code). Standard library only.
  B  reference the official scoring code (run_char_rm.py of github.com/morecry/CharacterEval at c3d44a6, reproduced
               below) on the 26 probes of charactereval_rm_probes.json: the model directory's own BaichuanCharRM class
               and BaichuanTokenizer, bfloat16, batch size 1, encode(add_special_tokens=False) + [eos], the last 4,096
               tokens kept, score = model(input_ids)[1] * 4 + 1. Needs torch, transformers and one GPU with ~28 GB free.
  C  endpoint  the same 26 probes through the deployed endpoint, in the request the platform sends
               (POST <base>/rewards, {"model": ..., "inputs": [{"messages": [{"role": "user", "content": <text>}]}]},
               answer {"data": [{"index": i, "score": s}]}): once as one batch, once in reverse order, once one by one.

Criteria, fixed before any result: A passes when every file matches. C passes when the three passes return the same
score for every probe and every score is in [1, 5]. B against C passes when every |endpoint − reference| <= 1e-6; when
all differences are at most 0.016 (one bfloat16 step of sigmoid × 4) the result is "within one bf16 step" and the
cause (batching, padding, kernels) must be explained; anything larger fails.

    python3 charactereval_rm_server_check.py --model-dir /path/to/BaichuanCharRM                      # A only
    python3 charactereval_rm_server_check.py --model-dir DIR --reference --device cuda:0             # A + B
    RM_API_KEY=... python3 charactereval_rm_server_check.py --model-dir DIR --reference \\
        --api https://host/v1 --api-model BaichuanCharRM                                              # A + B + C

The probe file must sit next to this script (or be named with --probes); its SHA-256 is checked. The API key is read
from the environment variable RM_API_KEY and never printed. Reads files and calls the endpoint; changes nothing.
Exit status 0 when every check that ran passed.
"""
import argparse
import datetime
import hashlib
import json
import os
import platform
import socket
import sys
import time
import urllib.error
import urllib.request

REVISION = "e953da5f57b86db7adca3a2714b4c3720f5b96a2"
# path: (size in bytes, "sha256" | "gitblob", expected) — huggingface.co/api/models/morecry/BaichuanCharRM/tree/<REVISION>
EXPECTED = {
    "pytorch_model-00001-of-00003.bin": (9973568983, "sha256", "c868e1bdccfe16392538d8020ef3f27ccf74b47cfca3c867617e8a8e0f2aa700"),
    "pytorch_model-00002-of-00003.bin": (9947421360, "sha256", "e0681ed449fb183d836920340c791e3bb3351c5f031011d05f11322d0d77b912"),
    "pytorch_model-00003-of-00003.bin": (6585330193, "sha256", "942a5e7322a87a7fb3c609db97c177f5e8835109a38c8dea01953c5c88e68dab"),
    "pytorch_model.bin.index.json": (23326, "gitblob", "756e175126441f96764bf964461646516fc1ef3c"),
    "tokenizer.model": (2001107, "sha256", "79452955be6b419a65984273a9f08af86042e1c2a75ee3ba989cbf620a133cc2"),
    "config.json": (835, "gitblob", "040a315da9111f5d12ad6350c0d6230c29055a6b"),
    "configuration_baichuan.py": (1490, "gitblob", "c5967ae712619bfbebaec545599c96db0c46a33e"),
    "modeling_baichuan.py": (22388, "gitblob", "e1b03ea793c9e00b6c1c3ef5d0f58503312ff4da"),
    "tokenization_baichuan.py": (8720, "gitblob", "1f0f192c59899e65ed7952b42dd0ce90c0ecbecf"),
    "tokenizer_config.json": (987, "gitblob", "6f6e083da07814405e1eca55dd76213d41123099"),
    "special_tokens_map.json": (431, "gitblob", "16bdf0c6837427c80a810936ef0cae7bf98d80fb"),
}
PROBES_SHA256 = "8ad3fbaaa69a970a444065cc1f6c1a816644174116faae10f0e7b2b6672ec494"
MAX_SEQ_LENGTH = 4096  # run_char_rm.py
EXACT, BF16_STEP = 1e-6, 0.016


def file_hash(path, kind):
    h = hashlib.sha256() if kind == "sha256" else hashlib.sha1(f"blob {os.path.getsize(path)}\0".encode())
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 24), b""):
            h.update(chunk)
    return h.hexdigest()


def check_weights(model_dir):
    files, ok = [], True
    for name, (size, kind, want) in EXPECTED.items():
        path = os.path.join(model_dir, name)
        row = {"file": name, "kind": kind, "expected": want, "expected_size": size}
        if not os.path.isfile(path):
            row.update(present=False, match=False)
        else:
            got = file_hash(path, kind)
            row.update(present=True, size=os.path.getsize(path), actual=got, match=got == want and os.path.getsize(path) == size)
        ok &= row["match"]
        files.append(row)
    others = sorted(f for f in os.listdir(model_dir) if f not in EXPECTED and not f.startswith("."))
    return {"model_dir": os.path.realpath(model_dir), "revision": REVISION, "files": files, "other_files": others,
            "result": "ALL FILES MATCH" if ok else "MISMATCH"}, ok


def reference_scores(model_dir, texts, device):
    """run_char_rm.py at c3d44a6, with the model directory as the package and the probe texts as input_text."""
    import importlib.util
    import torch

    def load(name):
        spec = importlib.util.spec_from_file_location(f"BaichuanCharRM.{name}", os.path.join(model_dir, f"{name}.py"),
                                                      submodule_search_locations=None)
        mod = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = mod
        spec.loader.exec_module(mod)
        return mod

    pkg = type(sys)("BaichuanCharRM")
    pkg.__path__ = [model_dir]
    sys.modules["BaichuanCharRM"] = pkg
    load("configuration_baichuan")
    BaichuanCharRM = load("modeling_baichuan").BaichuanCharRM
    BaichuanTokenizer = load("tokenization_baichuan").BaichuanTokenizer

    # --- run_char_rm.py, lines 23–39 (only the model path, the device and the input list differ) ---
    tokenizer = BaichuanTokenizer.from_pretrained(model_dir)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"
    base_model = BaichuanCharRM.from_pretrained(model_dir, torch_dtype=torch.bfloat16).to(device)
    out = []
    for input_text in texts:
        input_ids = tokenizer.encode(text=input_text, add_special_tokens=False) + [tokenizer.eos_token_id]
        n_tokens = len(input_ids)
        if len(input_ids) > MAX_SEQ_LENGTH:
            input_ids = input_ids[-MAX_SEQ_LENGTH:]
        input_ids = torch.tensor(input_ids).unsqueeze(0).to(device)
        with torch.no_grad():
            score = base_model(input_ids=input_ids)[1].item() * 4 + 1
        out.append({"score": score, "tokens": n_tokens})
    env = {"torch": torch.__version__, "cuda": torch.version.cuda, "device": torch.cuda.get_device_name(device) if torch.cuda.is_available() else None}
    try:
        import transformers
        env["transformers"] = transformers.__version__
    except ImportError:
        pass
    return out, env


def call_endpoint(base, model, key, texts, timeout=600):
    body = json.dumps({"model": model, "inputs": [{"messages": [{"role": "user", "content": t}]} for t in texts]}).encode()
    req = urllib.request.Request(base.rstrip("/") + "/rewards", data=body, method="POST",
                                 headers={"content-type": "application/json", **({"authorization": f"Bearer {key}"} if key else {})})
    t = time.time()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            data = json.load(r)["data"]
    except urllib.error.HTTPError as e:
        raise RuntimeError(f"HTTP {e.code}: {e.read()[:300].decode(errors='replace')}")
    scores = [None] * len(texts)
    for d in data:
        scores[d["index"]] = d["score"]
    if len(data) != len(texts) or any(s is None for s in scores):
        raise RuntimeError(f"{len(data)} scores for {len(texts)} inputs, or indices incomplete")
    return scores, round(time.time() - t, 2)


def check_endpoint(base, model, key, texts):
    batch, t1 = call_endpoint(base, model, key, texts)
    rev, t2 = call_endpoint(base, model, key, texts[::-1])
    rev = rev[::-1]
    single = [call_endpoint(base, model, key, [t])[0][0] for t in texts]
    same = all(a == b == c for a, b, c in zip(batch, rev, single))
    in_range = all(1 <= s <= 5 for s in batch)
    return {"base_url": base, "model": model, "batch": batch, "reversed": rev, "single": single,
            "deterministic": same, "in_range": in_range, "seconds": {"batch": t1, "reversed": t2}}, same and in_range


def compare(ref, api):
    diffs = [abs(a - r) for a, r in zip(api, ref)]
    worst = max(diffs)
    verdict = "EQUAL" if worst <= EXACT else "WITHIN ONE BF16 STEP" if worst <= BF16_STEP else "DIFFERENT"
    return {"max_abs_diff": worst, "exact": sum(d <= EXACT for d in diffs), "within_bf16_step": sum(d <= BF16_STEP for d in diffs),
            "n": len(diffs), "verdict": verdict}, verdict == "EQUAL"


def main():
    ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    ap.add_argument("--model-dir", required=True)
    ap.add_argument("--probes", default=os.path.join(os.path.dirname(os.path.abspath(__file__)), "charactereval_rm_probes.json"))
    ap.add_argument("--reference", action="store_true", help="run the official scoring code (check B)")
    ap.add_argument("--device", default="cuda:0")
    ap.add_argument("--api", help="endpoint base URL, the part before /rewards (check C)")
    ap.add_argument("--api-model", default="BaichuanCharRM")
    args = ap.parse_args()

    raw = open(args.probes, "rb").read()
    if hashlib.sha256(raw).hexdigest() != PROBES_SHA256:
        sys.exit(f"{args.probes}: SHA-256 is not {PROBES_SHA256}; use the probe file published with this script")
    probes = json.loads(raw)["probes"]
    texts = [p["text"] for p in probes]

    report = {"script": "charactereval_rm_server_check.py", "time_utc": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds"),
              "host": socket.gethostname(), "python": platform.python_version(), "probes_sha256": PROBES_SHA256,
              "criteria": {"weights": "every file matches", "endpoint": "same score in all three passes, all in [1, 5]",
                           "endpoint_vs_reference": f"EQUAL when every |diff| <= {EXACT}; WITHIN ONE BF16 STEP when <= {BF16_STEP}"}}
    ok = True
    report["A_weights"], a_ok = check_weights(args.model_dir)
    ok &= a_ok
    ref = None
    if args.reference:
        scores, env = reference_scores(args.model_dir, texts, args.device)
        ref = [s["score"] for s in scores]
        report["B_reference"] = {"environment": env, "probes": [{"n": p["n"], "kind": p["kind"], "id": p["id"], "metric_en": p["metric_en"], **s} for p, s in zip(probes, scores)]}
    if args.api:
        try:
            report["C_endpoint"], c_ok = check_endpoint(args.api, args.api_model, os.environ.get("RM_API_KEY"), texts)
        except Exception as e:  # noqa: BLE001 — reported, not raised
            report["C_endpoint"], c_ok = {"base_url": args.api, "error": str(e)}, False
        ok &= c_ok
        if ref is not None and c_ok:
            report["B_vs_C"], bc_ok = compare(ref, report["C_endpoint"]["batch"])
            ok &= bc_ok
    report["result"] = "PASS" if ok else "FAIL"
    json.dump(report, sys.stdout, ensure_ascii=False, indent=1)
    print()
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()