"""Generate an image; optionally save its bytes using the same owner's app session."""
import argparse
import base64
import json
import os
from pathlib import Path
import re
from urllib.request import Request, urlopen
from kendr_http import API_ORIGIN, post_json


def generate(prompt, request_id, reference=None):
    payload = {"model": "kendr-image", "prompt": prompt,
               "size": "1536x1024", "quality": "medium"}
    if reference:
        source = Path(reference)
        if source.stat().st_size > 8 * 1024 * 1024:
            raise ValueError("Reference image exceeds 8 MiB")
        # The server infers and validates PNG/JPEG/WebP from the bytes.
        payload["input_images"] = [{"name": source.name,
            "data": base64.b64encode(source.read_bytes()).decode("ascii")}]
    return post_json("/api/v1/images/generations", payload, request_id)


def download(image, destination):
    # Construct a fixed-origin path; never send the session credential to an arbitrary URL.
    image_id = image["id"]
    if not re.fullmatch(r"img_[a-f0-9]{24}", image_id):
        raise ValueError("Unexpected image identifier")
    request = Request(API_ORIGIN + "/api/me/generated-images/" + image_id,
        headers={"X-Kendr-Session": os.environ["KENDR_SESSION_TOKEN"]})
    with urlopen(request, timeout=180) as response:
        content_type = response.headers.get_content_type()
        extension = {"image/png": ".png", "image/jpeg": ".jpg", "image/webp": ".webp"}.get(content_type)
        if not extension:
            raise ValueError("Download did not return an image")
        path = Path(destination).with_suffix(extension)
        with path.open("xb") as output:
            output.write(response.read())
    return path


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("prompt")
    parser.add_argument("--request-id", required=True)
    parser.add_argument("--reference", help="Local PNG, JPEG, or WebP to edit")
    parser.add_argument("--output", help="Output filename; requires KENDR_SESSION_TOKEN")
    args = parser.parse_args()
    if args.output and not os.environ.get("KENDR_SESSION_TOKEN"):
        parser.error("--output requires KENDR_SESSION_TOKEN for the same account as the API key")
    result = generate(args.prompt, args.request_id, args.reference)
    print(json.dumps({"request_id": result["request_id"], "image": result["image"],
        "credit_micros_charged": result.get("credit_micros_charged")}, indent=2))
    if args.output:
        print(download(result["image"], args.output))