#!/usr/bin/env python3
# Generate bestiary portraits with fal-ai/nano-banana-pro/edit; the reference image is the
# front and front three-quarter view of the creature's 64 px idle sheet, scaled up 6x.
import base64, io, json, os, sys, urllib.request, concurrent.futures as cf
from PIL import Image

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

OUT = "/home/zink/Projects/dungeoncraft/art/generated/portraits"
KEY = open(os.path.expanduser("~/.config/dungeoncraft/fal-key")).read().strip()
EXTRA = os.environ.get("EXTRA", "")


def ref_uri(cid):
    im = Image.open(sheet(cid)).convert("RGBA")
    h = im.height
    two = im.crop((0, 0, 2 * h, h)).resize((12 * h, 6 * h), Image.NEAREST)
    bg = Image.new("RGBA", two.size, (40, 36, 44, 255))
    bg.alpha_composite(two)
    buf = io.BytesIO()
    bg.convert("RGB").save(buf, "PNG")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()


def gen(cid, suffix=""):
    body = {"prompt": prompt_for(cid, EXTRA), "image_urls": [ref_uri(cid)], "aspect_ratio": "1:1",
            "resolution": "1K", "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)
        with urllib.request.urlopen(res["images"][0]["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)