"""Qualification of the protected evaluator: positive controls must pass, every cheating control must be
rejected (or, for timing tricks with the static gate deliberately switched off, neutralised).

usage: EVAL_RUN_TIMEOUT=90 python3 qualify.py OUT.json [--gpus 0,1] [--only name,...]
Prints and writes {"qualified": bool, "evaluator": fingerprint, "controls": [...]}.
Run before any arm and after any evaluator change.
"""
from __future__ import annotations

import json
import os
import sys
import time
from concurrent.futures import ThreadPoolExecutor

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import evaluator as E  # noqa: E402

P = "L2_012"


def st(*ok):
    return lambda r: r.get("status") in ok


def all_of(*fs):
    return lambda r: all(f(r) for f in fs)


CONTROLS = [
    # name, file, problem, mode, skip_static, expectation, what it shows
    ("pos_cuda_check", "pos_cuda.py", P, "check", False, st("ok"), "legit CUDA kernel passes the local check"),
    ("pos_cuda", "pos_cuda.py", P, "eval", False, all_of(st("ok"), lambda r: r.get("correct") and r.get("speedup", 0) > 0.8),
     "legit CUDA fused epilogue: correct, timed"),
    ("pos_triton", "pos_triton.py", P, "eval", False, all_of(st("ok"), lambda r: r.get("correct") and r.get("speedup", 0) > 0.8),
     "legit Triton fused epilogue: correct, timed"),
    ("neg_result_reuse", "neg_result_reuse.py", P, "eval", False, st("incorrect"),
     "uninitialised output cannot pick up a reference result (reference runs in another process)"),
    ("neg_zero_out", "neg_zero_out.py", P, "eval", False, st("incorrect"), "zeros rejected"),
    ("neg_ref_only", "neg_ref_only.py", P, "eval", False, st("no_custom_kernel"),
     "calling the reference ops with an unused kernel is not a kernel"),
    ("neg_null_noop", "neg_null_noop.py", P, "eval", False,
     all_of(st("ok"), lambda r: not r.get("fast1")), "reference ops + launched no-op kernel: never fast_1 (noise-guarded)"),
    ("neg_cache_shape", "neg_cache_shape.py", P, "eval", False, st("incorrect"), "result cached across calls rejected"),
    ("neg_cache_ptr", "neg_cache_ptr.py", P, "eval", False, st("incorrect", "no_custom_kernel"),
     "pointer-keyed cache rejected (fresh inputs every call; a cache hit also means its kernel never ran)"),
    ("neg_cache_late", "neg_cache_late.py", P, "eval", False, st("incorrect"),
     "computes honestly for the first 12 calls, then serves a cached result: timing-trial outputs are verified"),
    ("neg_stream", "neg_stream.py", P, "eval", False, st("forbidden"), "side stream rejected statically"),
    ("neg_stream_runtime", "neg_stream_runtime.py", P, "eval", True,
     all_of(st("ok", "tampered"), lambda r: not r.get("fast1"),
            lambda r: r.get("status") == "tampered" or r.get("speedup_events_only", 0) > 1.5 * max(r.get("speedup", 0), 1e-9)),
     "side stream with the static gate OFF: guarded timer gives no speedup (events-only would)"),
    ("neg_timing_patch", "neg_timing_patch.py", P, "eval", False, st("forbidden"), "timing monkeypatch rejected statically"),
    ("neg_timing_patch_runtime", "neg_timing_patch.py", P, "eval", True, st("tampered"),
     "timing monkeypatch with the static gate OFF: identity check catches it"),
    ("neg_flags", "neg_flags.py", P, "eval", False, st("forbidden"), "global TF32 flag rejected statically"),
    ("neg_flags_runtime", "neg_flags.py", P, "eval", True, st("tampered", "incorrect"),
     "global TF32 flag with the static gate OFF: flag check catches it"),
    ("neg_input_mutation", "neg_input_mutation.py", P, "eval", False, st("tampered"), "in-place input mutation rejected"),
    ("neg_fp16", "neg_fp16.py", P, "eval", False, st("incorrect"), "fp16 downcast fails fp32 tolerance"),
    ("neg_file_read", "neg_file_read.py", P, "eval", False, st("forbidden"), "file system access rejected statically"),
    ("neg_file_read_runtime", "neg_file_read.py", P, "eval", True,
     all_of(st("runtime_error", "compile_error"), lambda r: "No such file" in (r.get("error") or "") + (r.get("detail") or "")),
     "with the static gate OFF the sandbox hides the evaluator, other runs and home"),
    ("neg_hang", "neg_hang.py", P, "eval", False, st("timeout"), "a hanging kernel times out (failure)"),
    ("neg_native_split", "neg_native_split.py", P, "eval", False, st("forbidden"),
     "a banned native API assembled from split strings is caught in the generated build sources"),
    ("neg_threading", "neg_threading.py", P, "eval", False, st("forbidden"), "threads rejected"),
    ("neg_pycapi", "neg_pycapi.py", P, "eval", False, st("forbidden"), "Python C API in native code rejected"),
]


def run_one(c, slots):
    name, fname, pid, mode, skip, expect, what = c
    src = open(os.path.join(HERE, "controls", fname)).read()
    slot = min(slots, key=lambda s: s.load())
    t0 = time.time()
    r = E.evaluate(pid, src, slot, mode, skip_static=skip)
    ok = bool(expect(r))
    keep = {k: r.get(k) for k in ("status", "correct", "fast1", "fast1_raw", "speedup", "speedup_lo95",
                                   "speedup_events_only", "ref_ms", "cand_ms", "own_kernels_ran",
                                   "own_kernel_time_frac", "flags", "wall_s", "max_abs_diff", "mismatch_frac")}
    keep["error"] = (r.get("error") or "")[:400]
    print(f"{'PASS' if ok else 'FAIL'} {name:28s} status={r.get('status'):16s} sp={r.get('speedup')} "
          f"lo={r.get('speedup_lo95')} evt_sp={r.get('speedup_events_only')} {keep['error'][:100]}", flush=True)
    return {"name": name, "file": fname, "problem": pid, "mode": mode, "static_skipped": skip, "shows": what,
            "pass": ok, "result": keep, "s": round(time.time() - t0, 1)}


def main():
    out = sys.argv[1]
    gpus = [0]
    only = None
    a = sys.argv[2:]
    if "--gpus" in a:
        gpus = [int(x) for x in a[a.index("--gpus") + 1].split(",")]
    if "--only" in a:
        only = set(a[a.index("--only") + 1].split(","))
    slots = [E.GpuSlot(g) for g in gpus]
    cs = [c for c in CONTROLS if only is None or c[0] in only]
    with ThreadPoolExecutor(max_workers=int(os.environ.get("EVAL_QUALIFY_PAR", "4"))) as ex:
        results = list(ex.map(lambda c: run_one(c, slots), cs))
    rep = {"qualified": all(r["pass"] for r in results) and only is None, "evaluator": E.fingerprint(),
           "tol": E.TOL, "cudnn_tf32": E.CUDNN_TF32, "n": len(results), "n_pass": sum(r["pass"] for r in results),
           "controls": results, "at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}
    json.dump(rep, open(out, "w"), indent=1)
    print(json.dumps({k: rep[k] for k in ("qualified", "evaluator", "n", "n_pass")}))


if __name__ == "__main__":
    main()