#!/usr/bin/env python3
"""Resident MTP forward probe; temporarily lock clocks and restore on exit."""
import argparse
import datetime
import glob
import json
import os
from pathlib import Path
import signal
import subprocess
import time


def read(path):
    try:
        return Path(path).read_text().strip()
    except OSError:
        return None


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("binary")
    parser.add_argument("model")
    parser.add_argument("--output", required=True)
    parser.add_argument("--runtime", required=True)
    parser.add_argument("--seqlen", type=int, default=128)
    parser.add_argument("--tokens", type=int, default=64)
    parser.add_argument("--rounds", type=int, default=3)
    parser.add_argument("--context", type=int, default=32768)
    parser.add_argument("--cases")
    args = parser.parse_args()
    if os.geteuid() != 0:
        parser.error("run with sudo for temporary clock locking and NPU debugfs sampling")
    help_env = os.environ.copy()
    help_env["LD_LIBRARY_PATH"] = args.runtime + ":" + help_env.get("LD_LIBRARY_PATH", "")
    probe = subprocess.run([args.binary], stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                           text=True, env=help_env, timeout=15)
    if "Usage:" not in probe.stdout:
        parser.error("cannot query engine CLI: " + probe.stdout.strip())
    if "mtp-benchmark" not in probe.stdout:
        parser.error("MTP is not compiled; rebuild with -DLING3_EXPERIMENTAL_MTP=ON")
    output = Path(args.output)
    output.mkdir(parents=True, exist_ok=False)
    original = {}
    applied = {}
    child = None

    def write_clock(path, value):
        original[path] = read(path)
        if original[path] is None:
            raise RuntimeError("cannot snapshot " + path)
        Path(path).write_text(str(value))
        applied[path] = read(path)

    def stop(signum, frame):
        raise KeyboardInterrupt

    signal.signal(signal.SIGTERM, stop)
    signal.signal(signal.SIGINT, stop)
    try:
        for policy in sorted(glob.glob("/sys/devices/system/cpu/cpufreq/policy*")):
            write_clock(policy + "/scaling_governor", "performance")
            write_clock(policy + "/scaling_min_freq", read(policy + "/scaling_max_freq"))
        for device in ("fdab0000.npu", "dmc"):
            base = "/sys/class/devfreq/" + device
            frequencies = read(base + "/available_frequencies")
            maximum = max(map(int, frequencies.split())) if frequencies else int(read(base + "/max_freq"))
            write_clock(base + "/governor", "performance")
            write_clock(base + "/min_freq", maximum)
        env = os.environ.copy()
        for key in ("LING3_GDN_PREFILL_DIR", "LING3_PREFILL_W4A4", "LING3_EXPERT_ALL_CORES",
                    "LING3_MLA_DECODE_NPU", "LING3_DISABLE_BATCH_INPUT_REUSE",
                    "LING3_DISABLE_PARALLEL_GATHER"):
            env.pop(key, None)
        for key in ("LING3_PREWARM_EXPERTS", "LING3_EXPERT_BALANCED", "LING3_EXPERT_ZERO_COPY",
                    "LING3_GDN_CPU_PREFILL", "LING3_GDN_CPU_FP32_STATE", "LING3_GDN_CPU_DECODE",
                    "LING3_GDN_FULL_FP32", "LING3_MLA_SIMD", "LING3_VECTOR_MATH"):
            env[key] = "1"
        env["LING3_MLA_BACKEND"] = "cpu"
        env["LD_LIBRARY_PATH"] = args.runtime + ":" + env.get("LD_LIBRARY_PATH", "")
        command = [args.binary, "mtp-benchmark", args.model, str(args.seqlen), str(args.tokens), str(args.rounds), str(args.context)]
        if args.cases:
            command.append(args.cases)
        metadata = {"started_utc": datetime.datetime.now(datetime.timezone.utc).isoformat(),
                    "command": command, "original_clocks": original, "meminfo_before": read("/proc/meminfo"), "applied_clocks": applied,
                    "profile": {k: v for k, v in env.items() if k.startswith("LING3_")}}
        (output / "metadata.json").write_text(json.dumps(metadata, indent=2))
        with (output / "benchmark.log").open("w") as log, (output / "telemetry.jsonl").open("w") as telemetry:
            child = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT, env=env)
            print("pid=" + str(child.pid) + " output=" + str(output), flush=True)
            start = time.monotonic()
            while child.poll() is None:
                row = {"elapsed_s": time.monotonic() - start,
                       "process_stat": read(f"/proc/{child.pid}/stat"),
                       "process_status": read(f"/proc/{child.pid}/status"),
                       "npu_load": read("/sys/kernel/debug/rknpu/load"),
                       "meminfo": read("/proc/meminfo")}
                row["frequencies"] = {p: read(p) for p in glob.glob("/sys/devices/system/cpu/cpufreq/policy*/scaling_cur_freq")
                                      + ["/sys/class/devfreq/fdab0000.npu/cur_freq", "/sys/class/devfreq/dmc/cur_freq"]}
                row["temperatures"] = {p: read(p) for p in glob.glob("/sys/class/thermal/thermal_zone*/temp")}
                telemetry.write(json.dumps(row) + "\n")
                telemetry.flush()
                time.sleep(0.5)
            metadata["exit_code"] = child.returncode
            metadata["elapsed_s"] = time.monotonic() - start
            print("exit_code=" + str(child.returncode), flush=True)
    finally:
        if child is not None and child.poll() is None:
            child.terminate()
            try:
                child.wait(timeout=15)
            except subprocess.TimeoutExpired:
                child.kill()
                child.wait()
        restored = {}
        for path, value in reversed(list(original.items())):
            try:
                Path(path).write_text(value)
                restored[path] = read(path)
            except OSError as error:
                restored[path] = str(error)
        final = locals().get("metadata", {"original_clocks": original, "applied_clocks": applied})
        final["restored_clocks"] = restored
        (output / "metadata.json").write_text(json.dumps(final, indent=2))
        if "SUDO_UID" in os.environ:
            for path in [output] + list(output.iterdir()):
                os.chown(path, int(os.environ["SUDO_UID"]), int(os.environ["SUDO_GID"]))
        print("clocks_restored", flush=True)
    return child.returncode if child else 1


if __name__ == "__main__":
    raise SystemExit(main())