"""Cold start: fresh process -> import -> load weights (already on disk) -> first answer.
Each model is measured 3 times in a new Python process; median reported.

  python bench\\cold_start.py
"""
import json
import os
import statistics
import subprocess
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
PY = sys.executable

SNIPPET = r'''
import os, time, json
t0 = time.perf_counter()
os.environ.setdefault("USE_TF", "0")
KIND, NAME = %r, %r
st = {"ticket": "My card was charged twice, I want a refund"}
q = {"team": {"type": "choice", "instructions": "Which team?", "criteria": {"billing": None, "technical": None}}}
if KIND == "laya":
    import torch; torch.set_num_threads(8)
    import laya
    m = laya.load(NAME, device="cpu")
    t1 = time.perf_counter()
    m.system_one(st, q)
else:
    import nirnay
    m = nirnay.load(NAME, threads=8)
    t1 = time.perf_counter()
    m.predict(st, q)
t2 = time.perf_counter()
print(json.dumps({"load_s": t1 - t0, "first_answer_s": t2 - t0}))
'''


def run(kind, name):
    env = dict(os.environ, HF_HUB_OFFLINE="1", PYTHONPATH=ROOT)
    out = subprocess.run([PY, "-c", SNIPPET % (kind, name)], capture_output=True, text=True, env=env, cwd=ROOT)
    line = [l for l in out.stdout.splitlines() if l.startswith("{")]
    if not line:
        raise RuntimeError(out.stderr[-2000:])
    return json.loads(line[-1])


def main():
    models = [("nirnay", os.path.join(ROOT, "models", n)) for n in os.environ.get("COLD_NIRNAY", "nirnay-v1").split(",")]
    models += [("laya", "laya"), ("laya", "laya-multilingual"), ("laya", "laya-typed-decisions")]
    only = os.environ.get("COLD_ONLY")
    out_path = os.path.join(HERE, "results", "cold_start.json")
    res = json.load(open(out_path)) if (only and os.path.exists(out_path)) else {}
    if only:
        models = [m for m in models if m[0] == only]
    for kind, name in models:
        runs = [run(kind, name) for _ in range(3)]
        key = "%s:%s" % (kind, os.path.basename(name))
        res[key] = {"first_answer_s": round(statistics.median(r["first_answer_s"] for r in runs), 2),
                    "load_s": round(statistics.median(r["load_s"] for r in runs), 2)}
        print(key, res[key], flush=True)
    json.dump(res, open(out_path, "w"), indent=2)


if __name__ == "__main__":
    main()