#!/usr/bin/env python3
"""Sequential, resumable LLaVA-1.5-7B (GGUF Q8_0, llama.cpp CPU) evaluation.
Resource discipline: one server, -t 3 (auto-drop to -t 2 if 5-min load > 6), nice 15 + ionice idle,
watchdog before every image (1-min load > 4.5 -> sleep 60, repeat).
Speed-up that does not change answers: loop order arm -> image -> all questions of that image, so the image
(system prompt + 576 image tokens) is prefilled once and every later question reuses the KV prompt cache.
Each image's rows are appended to results jsonl in one write, then progress.json is updated -> kill/reboot safe."""
import base64, json, os, signal, subprocess, sys, time, urllib.request
from pathlib import Path
PKG = Path(__file__).resolve().parents[2]; W = PKG / "work"
cfg = json.load(open(sys.argv[1] if len(sys.argv) > 1 else W / "code/run_config.json"))
PNG = Path(cfg["png_dir"]); TAG = cfg.get("tag", "main"); OUT = W / f"results/answers_{TAG}.jsonl"; PROG = W / f"results/progress_{TAG}.json"; LOG = W / f"logs/03_run_vlm_{TAG}.log"
SYS = ("A chat between a curious human and an artificial intelligence assistant. "
       "The assistant gives helpful, detailed, and polite answers to the human's questions.")
SUFFIX = "\nAnswer the question using a single word or phrase."
PORT = cfg.get("port", 18131)
def log(*x):
    with open(LOG, "a") as f: print(time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), *x, file=f, flush=True)
def loadavg(): return [float(v) for v in open("/proc/loadavg").read().split()[:3]]
threads = cfg.get("threads", 3); srv = None; MARK = None
def start_server():
    global srv, MARK
    cmd = ["nice", "-n", "15", "ionice", "-c", "3", f"{cfg['llama_dir']}/llama-server", "-m", cfg["model"], "--mmproj", cfg["mmproj"],
           "-t", str(threads), "-tb", str(threads), "-c", "4096", "-np", "1", "--port", str(PORT), "--host", "127.0.0.1"]
    srv = subprocess.Popen(cmd, stdout=open(W / "logs/llama_server.log", "a"), stderr=subprocess.STDOUT,
                           env=dict(os.environ, LD_LIBRARY_PATH=cfg["llama_dir"]), start_new_session=True)
    t = time.time()
    while True:
        try:
            if json.load(urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health", timeout=3)).get("status") == "ok": break
        except Exception: pass
        if srv.poll() is not None: raise RuntimeError("llama-server died on start")
        time.sleep(2)
    MARK = json.load(urllib.request.urlopen(f"http://127.0.0.1:{PORT}/props"))["media_marker"]
    log(f"server up pid={srv.pid} threads={threads} load_s={time.time()-t:.1f} cmd={' '.join(cmd)}")
def stop_server():
    if srv and srv.poll() is None:
        os.killpg(srv.pid, signal.SIGINT)
        try: srv.wait(60)
        except Exception: os.killpg(srv.pid, signal.SIGKILL)
def ask(png_b64, q, n_predict):
    body = {"prompt": {"prompt_string": f"{SYS} USER: {MARK}\n{q}{SUFFIX} ASSISTANT:", "multimodal_data": [png_b64]},
            "n_predict": n_predict, "temperature": 0, "top_k": 1, "cache_prompt": True, "id_slot": 0}
    t = time.time()
    r = json.load(urllib.request.urlopen(urllib.request.Request(f"http://127.0.0.1:{PORT}/completion", json.dumps(body).encode(),
                                                                {"Content-Type": "application/json"}), timeout=3600))
    tm = r.get("timings", {})
    return r["content"], dict(wall_s=round(time.time() - t, 3), prompt_n=tm.get("prompt_n"), cache_n=tm.get("cache_n"),
                              prompt_ms=round(tm.get("prompt_ms", 0)), predicted_n=tm.get("predicted_n"))
def main():
    global threads
    qs = json.load(open(W / "data/questions.json")); by_img = {}
    for x in qs: by_img.setdefault(x["image"], []).append(x)
    sel = json.load(open(W / "data/selection.json")); order = [n for s in sel["subsets"].values() for n in s]
    done = set(map(tuple, json.load(open(PROG))["done"])) if PROG.exists() else set()
    todo = [(arm, n) for arm in cfg["arms"] for n in (order[:cfg["limit_images"]] if cfg.get("limit_images") else order) if (arm, n) not in done]
    log(f"start: {len(done)} done, {len(todo)} to do; arms={cfg['arms']}")
    start_server()
    try:
        for k, (arm, n) in enumerate(todo):
            if (W / "results/STOP").exists(): log("STOP file found, exiting cleanly"); break
            waited = 0
            while loadavg()[0] > cfg.get("load_max", 4.5):
                log(f"watchdog: load1={loadavg()[0]} > {cfg.get('load_max', 4.5)}, sleep 60 (waited {waited}s)"); time.sleep(60); waited += 60
            if loadavg()[1] > 6 and threads > 2:
                log(f"load5={loadavg()[1]} > 6 -> restarting server with -t 2"); stop_server(); threads = 2; start_server()
            b64 = base64.b64encode((PNG / arm / (n[:-4] + ".png")).read_bytes()).decode()
            rows, t_img = [], time.time()
            for i, x in enumerate(by_img[n]):
                for attempt in range(3):
                    try:
                        ans, tm = ask(b64, x["question"], 3 if x["task"] == "pope" else 10); break
                    except Exception as e:
                        log(f"request error {arm} {n} {x['qid']}: {e!r}; restarting server"); stop_server(); start_server()
                else: raise RuntimeError("3 failed attempts")
                if i > 0 and (tm["prompt_n"] or 0) > 200: log(f"WARN cache miss {arm} {n} {x['qid']} prompt_n={tm['prompt_n']}")
                rows.append(dict(arm=arm, image=n, qid=x["qid"], task=x["task"], subset=x["subset"], answer_raw=ans, threads=threads, **tm))
            with open(OUT, "a") as f:
                f.write("".join(json.dumps(r, ensure_ascii=False) + "\n" for r in rows)); f.flush(); os.fsync(f.fileno())
            done.add((arm, n)); tmp = PROG.with_suffix(".tmp")
            json.dump(dict(done=sorted(done), updated=time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), threads=threads), open(tmp, "w")); os.replace(tmp, PROG)
            log(f"[{len(done)}] {arm} {n} q={len(rows)} img_s={time.time()-t_img:.1f} first_prompt_n={rows[0]['prompt_n']} load={loadavg()} waited={waited}")
    finally:
        stop_server()
    log("exit")
if __name__ == "__main__": main()