Generate an image with Python

"""Python 3.10+. Standard library only. Server-side; local-fixture tested, not a live model test."""
import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

EXAMPLE_PAYLOAD = {
    "model": "google/gemini-3.1-flash-lite-image",
    "prompt": "A single blue ceramic cup on a plain cream background, soft studio light.",
    "resolution": "1K", "aspect_ratio": "1:1", "n": 1, "async": True,
}
ACTIVE = {"submitted", "queued", "pending", "running", "processing"}


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        return None


def safe_url(base, path):
    import re
    origin = urllib.parse.urlsplit(base)
    url = urllib.parse.urlsplit(urllib.parse.urljoin(base + "/", path))
    valid_scheme = origin.scheme == "https" or (origin.scheme == "http" and origin.hostname in {"127.0.0.1", "::1"})
    def address(value):
        return (value.scheme, value.hostname, value.port or (443 if value.scheme == "https" else 80))
    if (not valid_scheme or origin.path not in {"", "/"} or origin.query or origin.fragment
            or origin.username or origin.password or url.username or url.password or url.query or url.fragment
            or address(url) != address(origin) or not re.fullmatch(r"/v1/images/(generations|edits)(/[^/?]+)?", url.path)):
        raise RuntimeError("Unsafe API or polling URL; no credentials sent.")
    return urllib.parse.urlunsplit(url)


def run_image(key, payload=None, base="https://api.maxapi.dev", poll_url=None,
              interval=5, max_polls=120, timeout=30, on_task=lambda task: None):
    if not key:
        raise RuntimeError("Set MAXAPI_KEY in your server environment.")
    if not isinstance(max_polls, int) or not 1 <= max_polls <= 720 or interval < 0 or timeout <= 0:
        raise RuntimeError("Invalid polling limits.")
    opener = urllib.request.build_opener(NoRedirect())

    def request(path, body=None):
        import re
        headers = {"Authorization": "Bearer " + key}
        if body is not None:
            headers["Content-Type"] = "application/json"
        req = urllib.request.Request(safe_url(base, path), headers=headers,
                                     data=json.dumps(body).encode() if body is not None else None,
                                     method="POST" if body is not None else "GET")
        try:
            response = opener.open(req, timeout=timeout)
        except urllib.error.HTTPError as error:
            response = error
        except (urllib.error.URLError, TimeoutError, OSError) as error:
            raise RuntimeError("Network or timeout error; outcome may be unknown. No automatic resubmission.") from error
        with response:
            status = response.code
            try:
                result = json.load(response)
            except (ValueError, UnicodeError) as error:
                raise RuntimeError(f"HTTP {status}: non-JSON response; inspect history before resubmitting.") from error
        if not isinstance(result, dict):
            raise RuntimeError("Invalid JSON response object.")
        if not 200 <= status < 300 or result.get("error"):
            detail = result.get("error")
            code = detail.get("code", "api_error") if isinstance(detail, dict) else "api_error"
            code = code if isinstance(code, str) and re.fullmatch(r"[a-zA-Z0-9_-]{1,80}", code) else "api_error"
            raise RuntimeError(f"HTTP {status}: {code}; no automatic retry.")
        return result

    current = request(poll_url) if poll_url else request("/v1/images/generations", EXAMPLE_PAYLOAD if payload is None else payload)
    path = poll_url or current.get("poll_url")
    if path:
        path = safe_url(base, path)
        on_task({"id": current.get("id"), "poll_url": path})
    count = 0
    while True:
        if current.get("status") == "succeeded" or (not current.get("status") and isinstance(current.get("data"), list) and current["data"]):
            return current
        if current.get("status") not in ACTIVE:
            raise RuntimeError("Task failed, cancelled or returned an unknown state. Check history; no resubmission.")
        if not path:
            raise RuntimeError("Pending task has no polling URL. Check history; do not resubmit automatically.")
        if count >= max_polls:
            raise RuntimeError("Polling limit reached; task was NOT cancelled. Resume using the saved polling URL.")
        time.sleep(interval)
        current = request(path)
        count += 1


if __name__ == "__main__":
    try:
        result = run_image(os.environ.get("MAXAPI_KEY"), poll_url=os.environ.get("MAXAPI_POLL_URL"),
                           on_task=lambda task: print("Save to resume:", json.dumps(task), file=sys.stderr))
        print(json.dumps({"id": result.get("id"), "status": result.get("status", "succeeded"),
                          "urls": [item["url"] for item in result.get("data", []) if item.get("url")]}, indent=2))
    except (RuntimeError, ValueError) as error:
        print(str(error), file=sys.stderr)
        sys.exit(1)