#!/usr/bin/env python3
"""Print, or explicitly run, a recorded single-context benchmark recipe."""
import argparse
import json
import os
from pathlib import Path
import re
import shlex
import subprocess

assets = Path(__file__).resolve().parent
cases = json.loads((assets / "launch-arguments.json").read_text())
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("case", choices=[*cases, "original-strata"])
parser.add_argument("--execute", action="store_true")
args = parser.parse_args()
values = dict(os.environ, ASSETS=str(assets))

def expand(value):
    def substitute(match):
        key = match.group(1)
        if key not in values:
            parser.error("set " + key + " before using this recipe")
        return values[key]
    return re.sub(r"\$\{([A-Z_]+)\}", substitute, value)

if args.case == "original-strata":
    case = json.loads((assets / "original-strata-command.json").read_text())
    command = [expand(v) for v in case["command"]]
    configured = case["env"]
else:
    case = cases[args.case]
    command = ["nice", "-n", "10", "taskset", "-c", "0-15"] + [expand(v) for v in case["argv"]]
    configured = case["environment"]
    configured = dict(configured, LD_LIBRARY_PATH="${ENGINE}:/opt/cuda/lib64")
env = {k: v for k, v in os.environ.items()
       if not k.startswith(("GGML_MOE_", "GGML_CPU_PREPARED_", "MOE_", "STRATA_"))
       and k not in ("LD_PRELOAD", "GGML_CUDA_DISABLE_GRAPHS", "GGML_CUDA_DISABLE_FUSION")}
env.update({k: expand(v) for k, v in configured.items()})
print(shlex.join(["env"] + [k + "=" + env[k] for k in sorted(configured)] + command), flush=True)
if not args.execute:
    raise SystemExit(0)
# The benchmark helper is for isolated pool creation, not concurrent serving.
if case.get("benchmark_affinity") and not Path(env["LD_PRELOAD"]).is_file():
    parser.error("build the optional benchmark affinity helper first")
if args.case != "original-strata":
    os.execvpe(command[0], command, env)
request = (assets / "strata-request.txt").read_text()
with subprocess.Popen(command, env=env, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                      text=True, bufsize=1) as process:
    sent = finished = False
    for line in process.stdout:
        print(line, end="", flush=True)
        if line.startswith("READY ") and not sent:
            process.stdin.write(request)
            process.stdin.flush()
            sent = True
        if line.startswith("DONE "):
            process.stdin.write("QUIT\n")
            process.stdin.flush()
            finished = True
    rc = process.wait()
    if rc or not finished:
        raise SystemExit(rc or "original engine exited without DONE")