#!/usr/bin/env python3
# Generate creature sprite sheets with fal-ai/nano-banana-pro/edit, the imp's lower row as style ref.
import base64, json, os, sys, urllib.request, concurrent.futures as cf

OUT = "/home/zink/Projects/dungeoncraft/art/generated/sprites"
REF = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "style-ref.png")
KEY = open(os.path.expanduser("~/.config/dungeoncraft/fal-key")).read().strip()

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from creatures import CREATURES, prompt_for

ref_uri = "data:image/png;base64," + base64.b64encode(open(REF, "rb").read()).decode()


def gen(cid, suffix=""):
    body = {"prompt": prompt_for(cid), "image_urls": [ref_uri], "aspect_ratio": "21:9",
            "resolution": "2K", "num_images": 1, "output_format": "png"}
    req = urllib.request.Request("https://fal.run/fal-ai/nano-banana-pro/edit",
                                 data=json.dumps(body).encode(),
                                 headers={"Authorization": "Key " + KEY, "Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=600) as r:
            res = json.load(r)
        url = res["images"][0]["url"]
        with urllib.request.urlopen(url, timeout=300) as r:
            data = r.read()
        path = f"{OUT}/{cid}-raw{suffix}.png"
        open(path, "wb").write(data)
        return cid, "ok " + path
    except Exception as e:
        msg = str(e)
        if hasattr(e, "read"):
            try:
                msg += " " + e.read().decode()[:300]
            except Exception:
                pass
        return cid, "ERROR " + msg


if __name__ == "__main__":
    ids = sys.argv[1:] or [c["id"] for c in CREATURES]
    suffix = os.environ.get("SUFFIX", "")
    with cf.ThreadPoolExecutor(4) as ex:
        for cid, st in ex.map(lambda i: gen(i, suffix), ids):
            print(cid, st, flush=True)