"""Python 3 example: Gemini image generation with images, documents and video."""
import argparse
import base64
import json
import os
from pathlib import Path
import urllib.error
import urllib.parse
import urllib.request


def payload(prompt, files, size, ratio, pdfs=(), videos=(), attachments=(), text_files=()):
    parts = [{"text": prompt}]
    for name in text_files:
        parts.append({"text": Path(name).read_text(encoding="utf-8-sig")})
    for kind, name in [("image", f) for f in files] + [("pdf", f) for f in pdfs] + [("video", f) for f in videos] + list(attachments):
        data = Path(name).read_bytes()
        suffix = Path(name).suffix.lower()
        if kind in ("text/plain", "text/csv", "application/json"):
            text = data.decode("utf-8-sig")
            if kind == "application/json":
                json.loads(text)
            mime = kind
        elif kind in ("pdf", "application/pdf") and data.startswith(b"%PDF-"):
            mime = "application/pdf"
        elif kind in ("video", "video/mp4") and len(data) >= 16 and data[4:8] == b"ftyp" and suffix == ".mp4":
            # ponytail: container signatures only; upstream validates tracks and codecs.
            mime = "video/mp4"
        elif kind in ("video", "video/quicktime") and len(data) >= 16 and data[4:8] in (b"ftyp", b"moov", b"mdat", b"wide") and suffix == ".mov":
            mime = "video/quicktime"
        elif kind in ("video", "video/webm") and data.startswith(b"\x1aE\xdf\xa3") and b"webm" in data[:4096] and suffix == ".webm":
            mime = "video/webm"
        elif kind in ("image", "image/png") and data.startswith(b"\x89PNG\r\n\x1a\n"):
            mime = "image/png"
        elif kind in ("image", "image/jpeg") and data.startswith(b"\xff\xd8\xff"):
            mime = "image/jpeg"
        elif kind in ("image", "image/webp") and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
            mime = "image/webp"
        else:
            raise ValueError("Unsupported or mismatched file type " + kind + ": " + name)
        parts.append({"inlineData": {"mimeType": mime, "data": base64.b64encode(data).decode("ascii")}})
    return {"contents": [{"role": "user", "parts": parts}], "generationConfig": {
        "responseModalities": ["TEXT", "IMAGE"], "imageConfig": {"imageSize": size, "aspectRatio": ratio}}}


def save_images(data, output):
    images = []
    for candidate in data.get("candidates", []):
        for part in candidate.get("content", {}).get("parts", []):
            if part.get("thought"):
                continue
            if part.get("text"):
                print(part["text"])
            inline = part.get("inlineData", {})
            linked = part.get("fileData", {})
            if inline.get("data") or linked.get("fileUri"):
                if inline.get("data"):
                    raw = base64.b64decode(inline["data"], validate=True)
                    mime = inline.get("mimeType")
                else:
                    # Group 139 returns a one-hour image link by default; download it promptly.
                    req = urllib.request.Request(linked["fileUri"], headers={"User-Agent": "YM2API-Example/1.0"})
                    with urllib.request.urlopen(req, timeout=120) as response:
                        raw = response.read()
                    mime = linked.get("mimeType")
                suffix = {"image/png": ".png", "image/jpeg": ".jpg", "image/webp": ".webp"}.get(mime)
                if not suffix:
                    raise ValueError("Unexpected image MIME type")
                path = Path(str(output) + "-" + str(len(images) + 1) + suffix)
                with path.open("xb") as f:
                    f.write(raw)
                images.append(str(path))
    if not images:
        status = {"promptFeedback": data.get("promptFeedback"), "finishReasons": [c.get("finishReason") for c in data.get("candidates", [])]}
        raise RuntimeError("No output image: " + json.dumps(status, ensure_ascii=False))
    return images


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--list-models", action="store_true")
    parser.add_argument("--model", default="gemini-nano-banana-2.1")
    parser.add_argument("--prompt", default="Create a product photo of a yellow ceramic mug on a wooden desk.")
    parser.add_argument("--image", action="append", default=[], help="Repeat once per reference image")
    parser.add_argument("--pdf", action="append", default=[], help="Local PDF file; sent as application/pdf")
    parser.add_argument("--video", action="append", default=[], help="Local MP4, MOV or WebM file")
    parser.add_argument("--attachment", nargs=2, action="append", default=[], metavar=("MIME_TYPE", "FILE"), help="Explicit MIME type and local file; use a documented supported type")
    parser.add_argument("--text-file", action="append", default=[], help="Read a UTF-8 text file as a text part, without inlineData")
    parser.add_argument("--size", choices=["1K", "2K", "4K"], default="1K")
    parser.add_argument("--ratio", default="3:2")
    parser.add_argument("--output", default="result")
    args = parser.parse_args()
    base = os.environ["BASE_URL"].rstrip("/")
    if urllib.parse.urlsplit(base).scheme != "https":
        raise ValueError("BASE_URL must use HTTPS")
    body = None if args.list_models else json.dumps(payload(args.prompt, args.image, args.size, args.ratio, args.pdf, args.video, args.attachment, args.text_file)).encode()
    path = "/v1beta/models" if args.list_models else "/v1beta/models/" + urllib.parse.quote(args.model, safe="") + ":generateContent"
    req = urllib.request.Request(base + path, data=body, headers={"x-goog-api-key": os.environ["API_KEY"], "Content-Type": "application/json"})
    try:
        with urllib.request.urlopen(req, timeout=180) as response:
            data = json.load(response)
    except urllib.error.HTTPError as error:
        raise SystemExit("HTTP " + str(error.code) + ": " + error.read(1500).decode(errors="replace")) from None
    if args.list_models:
        print(json.dumps(data, ensure_ascii=False, indent=2))
    else:
        print("Saved:", ", ".join(save_images(data, args.output)))


if __name__ == "__main__":
    main()