#!/usr/bin/env python3
"""lfit-quick - fast, free cloud image generation via Pollinations (FLUX).

No API key, no account, no cost. Use this for FAST draft/concept images.
For final high-definition assets, prefer local ComfyUI (openclaw infer image generate).

Usage:
  lfit-quick --prompt "neon cyberpunk alley, game key art" \
                 [--out PATH] [--width 1024] [--height 1024] \
                 [--model flux] [--seed 12345] [--negative "blurry, watermark"]

On success prints the absolute path of the saved image (and nothing else on stdout),
so a calling agent can capture/attach it. Saves to
~/.openclaw/workspace/generated-images/ by default.
"""
import argparse
import datetime
import os
import sys
import time
import urllib.parse
import urllib.request

DEFAULT_DIR = os.path.expanduser("~/.openclaw/workspace/generated-images")
PNG_MAGIC = b"\x89PNG\r\n\x1a\n"
JPG_MAGIC = b"\xff\xd8\xff"


def build_url(args):
    params = {
        "width": args.width,
        "height": args.height,
        "model": args.model,
        "nologo": "true",
    }
    if args.seed is not None:
        params["seed"] = args.seed
    if args.negative:
        params["negative"] = args.negative
    return (
        "https://image.pollinations.ai/prompt/"
        + urllib.parse.quote(args.prompt, safe="")
        + "?"
        + urllib.parse.urlencode(params)
    )


def main():
    ap = argparse.ArgumentParser(description="Fast free cloud image gen (Pollinations/FLUX).")
    ap.add_argument("--prompt", required=True)
    ap.add_argument("--out")
    ap.add_argument("--width", type=int, default=1024)
    ap.add_argument("--height", type=int, default=1024)
    ap.add_argument("--model", default="flux", help="flux (default), flux-realism, turbo, etc.")
    ap.add_argument("--seed", type=int)
    ap.add_argument("--negative")
    args = ap.parse_args()

    os.makedirs(DEFAULT_DIR, exist_ok=True)
    out = args.out or os.path.join(
        DEFAULT_DIR, "lfit-quick-{:%Y%m%d-%H%M%S}.jpg".format(datetime.datetime.now())
    )
    out = os.path.abspath(os.path.expanduser(out))
    os.makedirs(os.path.dirname(out), exist_ok=True)

    url = build_url(args)
    last_err = None
    for attempt in range(1, 4):
        try:
            req = urllib.request.Request(url, headers={"User-Agent": "openclaw-lfit-quick/1.0"})
            with urllib.request.urlopen(req, timeout=180) as resp:
                data = resp.read()
            if len(data) < 1024 or not (data.startswith(JPG_MAGIC) or data.startswith(PNG_MAGIC)):
                raise ValueError("response was not a valid image ({} bytes)".format(len(data)))
            with open(out, "wb") as fh:
                fh.write(data)
            print(out)
            return 0
        except Exception as exc:  # noqa: BLE001 - report and retry
            last_err = exc
            sys.stderr.write("[attempt {}/3] {}\n".format(attempt, exc))
            if attempt < 3:
                time.sleep(3 * attempt)
    sys.stderr.write("FAILED after 3 attempts: {}\n".format(last_err))
    return 1


if __name__ == "__main__":
    sys.exit(main())
