#!/usr/bin/env python3
"""Self-healing single-arm runner: submits a graph JSON, watches queue+history,
resubmits if the item vanishes from queue without completing (queue-clear footgun).
Usage: arm_runner.py <graph.json> <tag> [max_minutes=60]
Success => print RESULT json line and exit 0.
"""
import json, sys, time, urllib.request

def api(path, obj=None, timeout=20):
    url = "http://127.0.0.1:8188" + path
    data = json.dumps(obj).encode() if obj is not None else None
    req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
    return json.load(urllib.request.urlopen(req, timeout=timeout))

def queue_has(pid):
    q = api("/queue")
    return any(i[1] == pid for i in q.get("queue_running", []) + q.get("queue_pending", []))

def hist_status(pid):
    h = api(f"/history/{pid}")
    if pid not in h:
        return None, []
    e = h[pid]
    outs = [im.get("filename") for o in e.get("outputs", {}).values() for im in o.get("images", [])]
    return e.get("status", {}).get("status_str"), outs

def hist_timing(pid):
    try:
        h = api(f"/history/{pid}")
        msgs = h[pid].get("status", {}).get("messages", [])
        ts = {m[0]: m[1] for m in msgs if m}
        if "execution_start" in ts and "execution_success" in ts:
            return round((ts["execution_success"] - ts["execution_start"]) / 1000)
    except Exception:
        pass
    return None

def main():
    gf, tag = sys.argv[1], sys.argv[2]
    deadline = time.time() + (int(sys.argv[3]) * 60 if len(sys.argv) > 3 else 60 * 60)
    payload = json.load(open(gf))["graph"]
    pid = None
    t_submit0 = time.time()
    while time.time() < deadline:
        if pid is None:
            r = api("/prompt", {"prompt": payload})
            if r.get("node_errors"):
                print(json.dumps({"tag": tag, "error": "node_errors",
                                  "detail": str(r["node_errors"])[:800]})); sys.exit(2)
            pid = r["prompt_id"]
            print(f"[{tag}] submitted {pid}", flush=True)
        # watch
        st, outs = hist_status(pid)
        if st == "success":
            dur = hist_timing(pid)
            print("RESULT " + json.dumps({"tag": tag, "prompt_id": pid, "status": "success",
                                          "render_sec": dur, "outputs": outs}), flush=True)
            sys.exit(0)
        if st not in (None, "success") :  # error etc
            print("RESULT " + json.dumps({"tag": tag, "prompt_id": pid, "status": st}), flush=True)
            sys.exit(4)
        # not in history yet
        if not queue_has(pid):
            print(f"[{tag}] {pid} vanished from queue w/o result -> resubmit", flush=True)
            pid = None
            continue
        time.sleep(25)
    print("RESULT " + json.dumps({"tag": tag, "status": "global_timeout"})); sys.exit(5)

if __name__ == "__main__":
    main()