#!/usr/bin/env python3
"""Replay the pinned streaming cache with generated, disposable expert shards.

Requires Python, NumPy, GCC-compatible cc and OpenMP. The fault case truncates
only a generated copy inside TemporaryDirectory. No released model is used.
Source revision: 81bb6c23ccee2154bc35cad922588ab66494a9a9.
"""
import hashlib
import os
from pathlib import Path
import shutil
import subprocess
import sys
import tempfile

HASHES = {
    "include/k3/k3.h": "799b92a3c811707d82e9d4539673792497e03be7486bfec45a0bd95aa429bd97",
    "src/core/k3_ops.c": "8094ada8bace36ce79ef0a279ec20864eb0ae3a387a0447d0fae45fbde7d65c6",
    "src/io/k3_load.c": "4133aff10966a26f11f7ffcc84b0d7121f34e2709a223d310edf5db43a26e01d",
    "src/io/k3_load.h": "b02d3d27af260ec35d97911656da44710c634711aa49e1b2d88e5023e5dd6dab",
    "src/io/k3_st.c": "9ff95a8e5802ffebe5c5d71def2d3831ce986056baad813c5dc6b67b66230ddd",
    "src/io/k3_st.h": "3b3085414deeb97c473783b6f56ebdc8bc64d766320f40c6569ae04e9aadfe47",
    "src/io/k3_portable_io.h": "188d9fb106848f787cab0ba0f9aa490b46e3d9603d48e3276da558f2ce52e98a",
    "src/cache/k3_cache.c": "91db22eb80bca940518df793d8b4990ab8eeee1a6ec5327bd7304eb5598018f9",
    "src/cache/k3_cache.h": "3b78981ace96a799b0e46fb2d8025daa37e3afe4fa75e1e2dfe24ddaadb05641",
    "third_party/json.h": "ab4fc64d6177e78abfd560fc5d94e6d8f414a14b11a48698c78479621050c619",
    "tools/_paths.py": "8300db925218524c0cdbe994cf8ffea2e0d71860f0cbc09d4b9a97b087f00fbd",
    "tools/make_cache_fixture.py": "0ddbb16cdd1689881ca8509e70aba8523d6f204045e203da56d5349e08c7df83",
    "tests/unit/test_cache.c": "7595ff4482bd4052a5a9d28947a01274f7fe6814ee33c31d17571c5e941667df"
}


def run(args, env=None):
    print("$ " + " ".join(map(str, args)), flush=True)
    result = subprocess.run(list(map(str, args)), env=env, check=False)
    print(f"exit={result.returncode}", flush=True)
    if result.returncode:
        raise SystemExit(result.returncode)


def main():
    if len(sys.argv) != 2:
        raise SystemExit("usage: python replay.py /path/to/kimi-k3-in-c")
    root = Path(sys.argv[1]).resolve()
    for relative, expected in HASHES.items():
        actual = hashlib.sha256((root / relative).read_bytes()).hexdigest()
        if actual != expected:
            raise SystemExit(f"source digest mismatch: {relative}")
    print(f"Verified {len(HASHES)} pinned source digests", flush=True)
    run([sys.executable, "--version"])
    run([sys.executable, "-c", "import numpy; print('NumPy', numpy.__version__)"])
    run(["cc", "--version"])
    flags = ["-O3", "-std=gnu99", "-Wall", "-Wextra", "-Wno-unused-parameter",
             "-fopenmp", "-pthread", "-ffp-contract=off"]
    includes = ["-I" + str(root / p) for p in
                ["include/k3", "src/io", "src/cache", "third_party"]]
    implementation = [root / p for p in ["src/cache/k3_cache.c", "src/io/k3_load.c",
                                       "src/io/k3_st.c", "src/core/k3_ops.c"]]
    env = os.environ.copy()
    env["OMP_NUM_THREADS"] = "4"
    env.pop("K3_NOPREFETCH", None)
    env.pop("K3_EXPERT_CHUNK", None)
    with tempfile.TemporaryDirectory(prefix="k3-cache-field-note-") as tmp:
        tmp = Path(tmp)
        fixture, fault = tmp / "fixture", tmp / "fault-copy"
        run([sys.executable, root / "tools/make_cache_fixture.py", fixture, "24"], env)
        shutil.copytree(fixture, fault)
        upstream, probe = tmp / "cache-test", tmp / "cache-probe"
        run(["cc", *flags, *includes, root / "tests/unit/test_cache.c",
             *implementation, "-lm", "-o", upstream], env)
        run([upstream, fixture], env)
        split_env = dict(env, K3_EXPERT_CHUNK="4096")
        run([upstream, fixture], split_env)
        run(["cc", *flags, *includes, Path(__file__).with_name("cache-probe.c"),
             *implementation, "-lm", "-o", probe], env)
        run([probe, fixture, fault], env)


if __name__ == "__main__":
    main()