#!/usr/bin/env python3
"""Tiny fal.ai queue client. Reads FAL_KEY from env or infography-mvp/.env."""
import json, os, re, sys, time, urllib.request, urllib.error

ROOT = "/root/notchlight"
ENVF = "/root/infography-mvp/.env"

def key():
    k = os.environ.get("FAL_KEY")
    if k:
        return k
    for line in open(ENVF):
        if line.startswith("FAL_KEY="):
            return line.split("=", 1)[1].strip().strip('"').strip("'")
    raise SystemExit("no FAL_KEY")

H = {"Authorization": "Key " + key(), "Content-Type": "application/json"}

def req(url, data=None):
    r = urllib.request.Request(url, data=None if data is None else json.dumps(data).encode(), headers=H,
                               method="GET" if data is None else "POST")
    with urllib.request.urlopen(r, timeout=120) as f:
        return json.loads(f.read().decode())

def run(endpoint, payload, label="gen", tries=240):
    app = "/".join(endpoint.split("/")[:2]) if endpoint.startswith("blackforestlabs") else "/".join(endpoint.split("/")[:2])
    # fal queue app id is the first two path segments for blackforestlabs/flux-3/... -> app is full endpoint minus method?
    sub = req(f"https://queue.fal.run/{endpoint}", payload)
    rid = sub["request_id"]
    status_url = sub["status_url"]
    print(f"[{label}] queued {rid}", flush=True)
    t0 = time.time()
    while time.time() - t0 < 600:
        s = req(status_url)
        st = s.get("status")
        if st == "COMPLETED":
            out = req(s["response_url"])
            print(f"[{label}] done in {time.time()-t0:.1f}s", flush=True)
            return out
        if st in ("FAILED", "ERROR"):
            raise RuntimeError(json.dumps(s)[:800])
        time.sleep(1.2)
    raise TimeoutError(label)

def download(url, path):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    urllib.request.urlretrieve(url, path)
    return path

if __name__ == "__main__":
    print(key()[:8] + "...")