"""Kernel evaluator worker: one process per phase, driven by evaluator.py over a line protocol.

Modes (argv[1]), config JSON path argv[2]:
  ref_outputs   TRUSTED (never sees candidate code): reference outputs for the given seeds -> <ref_dir>/<tag>.pt
  ref_time      TRUSTED: reference timing on the timing seeds, interleaved with the candidate by the parent
  cand_compile  SANDBOXED, no GPU device: import the candidate (load_inline builds into the work dir)
  cand_run      SANDBOXED, one GPU: correctness outputs on the given seeds, profiler kernel check, optional timing

Line protocol for timing modes: the worker prints "READY {json}" after warmup, then reads commands from stdin:
"GO <n>" -> runs n trials and prints "DONE {json}"; "EXIT" -> finishes and prints "RESULT {json}".
Every mode ends with exactly one "RESULT {json}" line on stdout.
"""
from __future__ import annotations

import json
import os
import sys
import time
import traceback

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


def emit(tag: str, obj) -> None:
    sys.stdout.write(f"{tag} {json.dumps(obj, default=str)}\n")
    sys.stdout.flush()


def torch_flags(torch) -> dict:
    b = torch.backends
    return {
        "cuda.matmul.allow_tf32": b.cuda.matmul.allow_tf32,
        "cudnn.allow_tf32": b.cudnn.allow_tf32,
        "cudnn.benchmark": b.cudnn.benchmark,
        "cudnn.deterministic": b.cudnn.deterministic,
        "cudnn.enabled": b.cudnn.enabled,
        "fp16_reduced": b.cuda.matmul.allow_fp16_reduced_precision_reduction,
        "bf16_reduced": b.cuda.matmul.allow_bf16_reduced_precision_reduction,
        "f32_matmul_precision": torch.get_float32_matmul_precision(),
        "deterministic_algos": torch.are_deterministic_algorithms_enabled(),
    }


def set_flags(torch, cfg) -> None:
    # Same explicit numerics settings in the reference and candidate processes (torch defaults unless the
    # evaluator config says otherwise); the candidate cannot change them (static gate + runtime check).
    torch.backends.cuda.matmul.allow_tf32 = bool(cfg.get("matmul_tf32", False))
    torch.backends.cudnn.allow_tf32 = bool(cfg.get("cudnn_tf32", False))
    torch.backends.cudnn.benchmark = False


FAST_HDR = "#include <torch/types.h>\n#include <ATen/cuda/CUDAContext.h>"
PCH_DIR = "/usr/local/eval-pch"


def build_accel() -> None:
    """Compile-time accelerations that never change semantics (trusted harness, installed before the candidate
    is imported, identical in the compile and run phases so the build hash matches):
    - g++ half (main.cpp): a precompiled torch/extension.h is offered via an include dir; GCC uses it only if
      valid for the exact flags.
    - nvcc half: <torch/extension.h> in cuda_sources is first tried as <torch/types.h> + CUDAContext (3x faster
      through nvcc); if that variant fails to build, the original sources are compiled unchanged."""
    import torch.utils.cpp_extension as ce
    orig = ce.load_inline
    ce._check_and_build_extension_h_precompiler_headers = lambda *a, **k: None
    ce.remove_extension_h_precompiler_headers = lambda *a, **k: None
    work = os.environ.get("HOME", "/tmp")

    def load_inline(name, cpp_sources, cuda_sources=None, *args, **kw):
        kw["use_pch"] = False
        if os.path.isdir(PCH_DIR):
            kw["extra_include_paths"] = [PCH_DIR] + list(kw.get("extra_include_paths") or [])
        if cuda_sources is not None and not args and not kw.get("no_implicit_headers"):
            cpp = [cpp_sources] if isinstance(cpp_sources, str) else list(cpp_sources)
            srcs = [cuda_sources] if isinstance(cuda_sources, str) else list(cuda_sources)
            # torch would prepend <torch/extension.h> to cpp and <torch/types.h>, <cuda.h>, <cuda_runtime.h> to
            # cuda; the fast variant keeps the cpp header (PCH) but drops the implicit <torch/types.h> from the .cu
            # and lightens an explicit <torch/extension.h> there. Any build failure -> the original call.
            fcpp = ["#include <torch/extension.h>"] + cpp
            fcu = ["#include <cuda.h>\n#include <cuda_runtime.h>"] + [x.replace("#include <torch/extension.h>", FAST_HDR)
                                                                     for x in srcs]
            marker = os.path.join(work, f".eval_nofast_{name}")
            if not os.path.exists(marker):
                try:
                    return orig(name + "_evalf", fcpp, fcu, no_implicit_headers=True, **kw)
                except Exception as e:  # noqa: BLE001
                    with open(marker, "w") as f:
                        f.write(str(e)[:2000])
        return orig(name, cpp_sources, cuda_sources, *args, **kw)

    ce.load_inline = load_inline


def main() -> None:
    mode, cfg_path = sys.argv[1], sys.argv[2]
    cfg = json.load(open(cfg_path))
    t0 = time.time()
    res: dict = {"mode": mode}
    try:
        import torch  # noqa: F401
        import common as C
        set_flags(torch, cfg)
        flags0 = torch_flags(torch)
        if mode in ("cand_compile", "cand_run"):
            build_accel()
        if mode == "cand_compile":
            mod = C.load_module(cfg["cand_path"], "eval_candidate")
            res.update(ok=hasattr(mod, "ModelNew"), error="" if hasattr(mod, "ModelNew") else "ModelNew not defined")
        elif mode == "ref_outputs":
            dev = torch.device("cuda:0")
            prob = C.load_module(cfg["problem_path"], "eval_problem")
            model = C.make_model(prob.Model, prob, cfg["init_seed"], dev)
            meta = {}
            with torch.no_grad():
                for tag, seed in cfg["seeds"].items():
                    inputs = C.make_inputs(prob, seed, dev)
                    out = model(*inputs)
                    torch.cuda.synchronize()
                    outs = C.to_cpu_outputs(out)
                    torch.save(outs, os.path.join(cfg["out_dir"], f"{tag}.pt"))
                    meta[tag] = [{"shape": list(o.shape), "dtype": str(o.dtype),
                                  "finite": bool(torch.isfinite(o).all().item())} for o in outs]
                    del inputs, out
            res.update(ok=True, outputs=meta, peak_mem_mb=torch.cuda.max_memory_allocated() / 2**20)
        elif mode in ("ref_time", "cand_run"):
            dev = torch.device("cuda:0")
            torch.cuda.set_device(0)
            timer = C.Timer(dev)  # captures timing primitives + calibrates BEFORE candidate import
            prob = C.load_module(cfg["problem_path"], "eval_problem")
            if mode == "ref_time":
                cls = prob.Model
            else:
                cand = C.load_module(cfg["cand_path"], "eval_candidate")
                cls = cand.ModelNew
            model = C.make_model(cls, prob, cfg["init_seed"], dev)
            if mode == "ref_time" and cfg.get("torch_compile"):
                model = torch.compile(model)  # baseline only (trusted): KernelBench's torch.compile default mode
            res["overhead_ms"] = timer.overhead_ms
            with torch.no_grad():
                if mode == "cand_run":
                    _correctness(torch, C, prob, model, dev, cfg, res)
                    if res.get("tamper") or res.get("runtime_error"):
                        raise _Stop()
                    if cfg.get("kernel_names") is not None:
                        _profile(torch, C, prob, model, dev, cfg, res)
                    emit("CHECKED", {k: res.get(k) for k in ("tamper", "runtime_error", "own_kernels_ran",
                                                              "own_kernel_time_frac", "correct_outputs")})
                if cfg.get("time_base") is not None:
                    _timing(torch, C, prob, model, dev, cfg, res, timer)
            tam = timer.tampered()
            flags1 = torch_flags(torch)
            if tam:
                res["tamper"] = (res.get("tamper") or []) + [f"timing/harness primitives replaced: {tam}"]
            if flags1 != flags0:
                res["tamper"] = (res.get("tamper") or []) + [
                    f"torch numerics flags changed: { {k: (flags0[k], flags1[k]) for k in flags0 if flags0[k] != flags1[k]} }"]
            if torch.cuda.current_stream() != torch.cuda.default_stream():
                res["tamper"] = (res.get("tamper") or []) + ["current CUDA stream was changed"]
            res["ok"] = not res.get("tamper") and not res.get("runtime_error")
            res["peak_mem_mb"] = torch.cuda.max_memory_allocated() / 2**20
        else:
            raise ValueError(f"unknown mode {mode}")
    except _Stop:
        res["ok"] = False
    except BaseException as e:  # noqa: BLE001 - report everything to the parent
        res.update(ok=False, runtime_error=f"{type(e).__name__}: {str(e)[:2000]}",
                   traceback=traceback.format_exc()[-4000:])
    res["wall_s"] = round(time.time() - t0, 3)
    emit("RESULT", res)
    sys.stdout.flush()
    os._exit(0)


class _Stop(Exception):
    pass


def _check_out(torch, out, dev):
    for o in (out if isinstance(out, (tuple, list)) else [out]):
        if isinstance(o, (tuple, list)):
            _check_out(torch, o, dev)
            continue
        if type(o) is not torch.Tensor:
            raise TypeError(f"output must be a plain torch.Tensor (got {type(o).__name__})")
        if o.device.type != "cuda":
            raise TypeError(f"output must be on the GPU (got {o.device})")


def _correctness(torch, C, prob, model, dev, cfg, res):
    res["correct_outputs"] = {}
    for tag, seed in cfg["seeds"].items():
        inputs = C.make_inputs(prob, seed, dev)
        fp0 = C.input_fingerprint(inputs)
        out = model(*inputs)
        torch.cuda.synchronize()
        _check_out(torch, out, dev)
        outs = C.to_cpu_outputs(out)
        fp1 = C.input_fingerprint(inputs)
        if fp0 != fp1:
            res["tamper"] = (res.get("tamper") or []) + ["forward modified its input tensors in place"]
        torch.save(outs, os.path.join(cfg["out_dir"], f"{tag}.pt"))
        res["correct_outputs"][tag] = [list(o.shape) for o in outs]
        del inputs, out, outs


def _profile(torch, C, prob, model, dev, cfg, res):
    """Which GPU kernels ran during one forward; at least one must be the candidate's own kernel."""
    seed = next(iter(cfg["seeds"].values()))
    inputs = C.make_inputs(prob, seed, dev)
    model(*inputs)
    torch.cuda.synchronize()
    from torch.profiler import ProfilerActivity, profile
    with profile(activities=[ProfilerActivity.CUDA]) as prof:
        model(*inputs)
        torch.cuda.synchronize()
    names, own, tot_us, own_us = {}, [], 0.0, 0.0
    wanted = [k for k in cfg["kernel_names"] if k]
    for ev in prof.events():
        dt = getattr(ev, "device_type", None)
        if dt is None or "CUDA" not in str(dt):
            continue
        us = float(getattr(ev, "device_time_total", 0.0) or getattr(ev, "device_time", 0.0) or 0.0)
        names[ev.name] = names.get(ev.name, 0.0) + us
        tot_us += us
        if any(w in ev.name for w in wanted):
            own.append(ev.name)
            own_us += us
    res["gpu_kernels"] = sorted(names.items(), key=lambda kv: -kv[1])[:25]
    res["own_kernels_ran"] = sorted(set(own))[:20]
    res["own_kernel_time_frac"] = round(own_us / tot_us, 4) if tot_us > 0 else 0.0
    del inputs


def _timing(torch, C, prob, model, dev, cfg, res, timer):
    """Every timed trial gets FRESH inputs (new seed, generated on the GPU outside the timed region), so no
    result can be cached across calls by pointer or by content. The output of the last trial of every block
    (block sizes chosen at random by the parent, unknown here until the GO arrives) is saved for verification
    against the trusted reference, which runs the identical trial sequence."""
    base = int(cfg["time_base"])
    for w in range(int(cfg.get("warmup", 3))):
        model(*C.make_inputs(prob, C.trial_seed(base, -100 - w), dev))
    torch.cuda.synchronize()
    est = []
    for w in range(3):
        inp = C.make_inputs(prob, C.trial_seed(base, -200 - w), dev)
        _, evt, wall = timer.trial(model, inp)
        est.append(max(evt, wall - timer.overhead_ms))
        del inp
    emit("READY", {"est_ms": sorted(est)[1], "overhead_ms": timer.overhead_ms})
    evts, walls = [], []
    k = 0
    blk = 0
    for line in sys.stdin:
        cmd = line.strip().split()
        if not cmd:
            continue
        if cmd[0] == "GO":
            n = int(cmd[1])
            for j in range(n):
                inp = C.make_inputs(prob, C.trial_seed(base, k), dev)
                out, evt, wall = timer.trial(model, inp)
                evts.append(evt)
                walls.append(wall)
                if j == n - 1:
                    _check_out(torch, out, dev)
                    torch.save(C.to_cpu_outputs(out), os.path.join(cfg["out_dir"], f"blk{blk}.pt"))
                del inp, out
                k += 1
            blk += 1
            emit("DONE", {"n": n})
        elif cmd[0] == "EXIT":
            break
    res["evt_ms"] = evts
    res["wall_ms"] = walls
    res["blocks"] = blk


if __name__ == "__main__":
    main()