"""CharacterJudge deployment check, run by the operator on the server that serves `character-judge`.

Proves that the process answering the platform's judge requests loads the released checkpoint: it finds the running
vLLM process, reads the model directory from its command line, and hashes every file the model is built from (SHA-256
for the weight shards, git blob SHA-1 for the small files, as Hugging Face publishes them), against the values Hugging
Face publishes for thu-coai/CharacterJudge at commit f2272fa1d19b9666b4fc1762041cd60e62d38244. Python 3 standard library
only; reads files, changes nothing. Secrets on the command line (API keys, tokens) are masked in the output.

    python3 characterjudge_server_check.py                 # finds the vLLM process serving character-judge
    python3 characterjudge_server_check.py /path/to/model  # or names the model directory

Prints one JSON document; send it back unchanged. Exit status 0 when every file matches.
"""
import datetime
import hashlib
import json
import os
import platform
import re
import socket
import subprocess
import sys

REVISION = "f2272fa1d19b9666b4fc1762041cd60e62d38244"
# path: (size in bytes, "sha256" | "gitblob", expected hash) — from huggingface.co/api/models/thu-coai/CharacterJudge/tree/<REVISION>
EXPECTED = {
    "model-00001-of-00004.safetensors": (4877660776, "sha256", "91c59e05b566e87b0199ca7af05f07ef820ba11c9812796a4063c62447abc3dd"),
    "model-00002-of-00004.safetensors": (4932751008, "sha256", "b2bd7dd0649504c302b057d3a4270af8b8bf411e5b768a6a38c6893c2a369fb2"),
    "model-00003-of-00004.safetensors": (4330865200, "sha256", "d683914d74a951bd8e9326c8eae93b9d673cec24a0ae479eba3d7da82bbb7f29"),
    "model-00004-of-00004.safetensors": (1089994880, "sha256", "f2839e43c4cd6befe61e0744d33494ee4d7611d6350f8e75dc49b298c4105b18"),
    "model.safetensors.index.json": (27752, "gitblob", "6ca5084b88f1a95fef37f1b94c6e87ff69422bb7"),
    "config.json": (740, "gitblob", "2c069af4be91b9176d82e218fbfb330aa40c6f7c"),
    "generation_config.json": (243, "gitblob", "c37d48219e693bb003eb59b5403509236469f3ab"),
    "tokenizer.json": (7028043, "gitblob", "0891f5d360e079340e2a28cd114d92db60c918a6"),
    "tokenizer_config.json": (1534, "gitblob", "a4c13118a13bccf6b9d04b95d6bb696ddc42e500"),
    "vocab.json": (2776833, "gitblob", "4783fe10ac3adce15ac8f358ef5462739852c569"),
    "merges.txt": (1671853, "gitblob", "31349551d90c7606f325fe0f11bbb8bd5fa0d7c7"),
    "added_tokens.json": (80, "gitblob", "409be9f93d4810d569098a16c4bf30c977edf44a"),
    "special_tokens_map.json": (367, "gitblob", "2c53454aff27c6d22873cc28c884c2b2d5d07693"),
}
SECRET_FLAGS = re.compile(r"^--?(api[-_]?key|token|hf[-_]?token|password|secret)", re.I)


def masked(argv):
    out, hide = [], False
    for a in argv:
        if hide:
            out.append("***")
            hide = False
        elif SECRET_FLAGS.match(a):
            if "=" in a:
                out.append(a.split("=", 1)[0] + "=***")
            else:
                out.append(a)
                hide = True
        else:
            out.append(re.sub(r"(sk-|hf_)[A-Za-z0-9_\-]{8,}", r"\1***", a))
    return out


def vllm_processes():
    procs = []
    for pid in os.listdir("/proc") if os.path.isdir("/proc") else []:
        if not pid.isdigit():
            continue
        try:
            argv = open(f"/proc/{pid}/cmdline", "rb").read().split(b"\0")
        except OSError:
            continue
        argv = [a.decode(errors="replace") for a in argv if a]
        if any("vllm" in a for a in argv) and any(a in ("serve", "vllm.entrypoints.openai.api_server") or a.endswith("api_server") for a in argv):
            procs.append({"pid": int(pid), "argv": argv})
    if not procs:  # no /proc (or another layout): fall back to ps
        try:
            for line in subprocess.run(["ps", "-eo", "pid,args"], capture_output=True, text=True).stdout.splitlines()[1:]:
                pid, _, args = line.strip().partition(" ")
                if "vllm" in args and ("serve" in args or "api_server" in args):
                    procs.append({"pid": int(pid), "argv": args.split()})
        except OSError:
            pass
    return procs


def model_dir_of(argv):
    for i, a in enumerate(argv):
        if a == "--model" and i + 1 < len(argv):
            return argv[i + 1]
        if a.startswith("--model="):
            return a.split("=", 1)[1]
        if a == "serve" and i + 1 < len(argv) and not argv[i + 1].startswith("-"):
            return argv[i + 1]
    return None


def served_names(argv):
    for i, a in enumerate(argv):
        if a == "--served-model-name":
            return [x for x in argv[i + 1:] if not x.startswith("-")][:4]
        if a.startswith("--served-model-name="):
            return a.split("=", 1)[1].split(",")
    return []


def resolve(model):
    """A local directory, or a Hugging Face repo id resolved to its snapshot in the local cache."""
    if model and os.path.isdir(model):
        return os.path.realpath(model)
    if model and "/" in model:
        home = os.environ.get("HF_HOME") or os.path.join(os.path.expanduser("~"), ".cache", "huggingface")
        snaps = os.path.join(os.environ.get("HF_HUB_CACHE") or os.path.join(home, "hub"), "models--" + model.replace("/", "--"), "snapshots")
        if os.path.isdir(snaps):
            cands = [os.path.join(snaps, d) for d in os.listdir(snaps)]
            exact = [c for c in cands if os.path.basename(c) == REVISION]
            return (exact or sorted(cands, key=os.path.getmtime))[-1]
    return None


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


def main():
    report = {
        "checked_at": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds"),
        "host": socket.gethostname(),
        "platform": platform.platform(),
        "python": sys.version.split()[0],
        "reference": f"huggingface.co/thu-coai/CharacterJudge@{REVISION}",
    }
    try:
        import vllm  # noqa: PLC0415

        report["vllm_version"] = vllm.__version__
    except Exception as e:  # not importable from this interpreter (e.g. another venv)
        report["vllm_version"] = f"not importable here: {type(e).__name__}"
    procs = vllm_processes()
    report["vllm_processes"] = [{"pid": p["pid"], "argv": masked(p["argv"]), "model_arg": model_dir_of(p["argv"]), "served_model_names": served_names(p["argv"])} for p in procs]
    arg = sys.argv[1] if len(sys.argv) > 1 else None
    if not arg:
        judge = [p for p in report["vllm_processes"] if "character-judge" in p["served_model_names"]] or report["vllm_processes"]
        arg = judge[0]["model_arg"] if judge else None
    model_dir = resolve(arg)
    report["model_arg"] = arg
    report["model_dir"] = model_dir
    files, ok = {}, bool(model_dir)
    for name, (size, kind, expected) in EXPECTED.items():
        path = os.path.join(model_dir, name) if model_dir else None
        if not path or not os.path.exists(path):
            files[name] = {"status": "MISSING"}
            ok = False
            continue
        real = os.path.realpath(path)
        got = digest(real, kind)
        match = got == expected and os.path.getsize(real) == size
        ok &= match
        files[name] = {"status": "MATCH" if match else "DIFFERENT", "kind": kind, "size": os.path.getsize(real), "hash": got, "expected": expected}
    report["files"] = files
    report["extra_weight_files"] = sorted(f for f in (os.listdir(model_dir) if model_dir else []) if re.search(r"\.(safetensors|bin|pt|pth|gguf)$", f) and f not in EXPECTED)
    report["result"] = "ALL FILES MATCH" if ok else "MISMATCH OR MISSING FILES"
    print(json.dumps(report, ensure_ascii=False, indent=1))
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()