# /// script
# requires-python = ">=3.11"
# dependencies = ["requests>=2.32,<3"]
# ///
"""Estimate for free. Paid creation needs --generate and confirmation."""
import argparse
import json
import os
import re
import time
from pathlib import Path
from urllib.parse import urljoin, urlparse

import requests

API_BASE = "https://api.entirefeed.com/v1"
DEFAULT_PROMPT = "Studio product photograph of a cobalt-blue ceramic mug on a warm ivory background. A small white sun symbol on the front, no lettering. Soft window light from the left, realistic glaze, gentle shadow, square composition with generous empty space."


def api(method, path, *, key=None, **kwargs):
    url = urljoin(API_BASE + "/", path)
    if urlparse(url).netloc != urlparse(API_BASE).netloc or urlparse(url).scheme != "https":
        raise ValueError("Unexpected API URL")
    headers = kwargs.pop("headers", {})
    if key:
        headers["Authorization"] = f"Bearer {key}"
    response = requests.request(method, url, headers=headers, timeout=60, **kwargs)
    response.raise_for_status()
    return response.json()


def main():
    parser = argparse.ArgumentParser(description="Estimate, create or edit, and download a Nano Banana image.")
    action = parser.add_mutually_exclusive_group()
    action.add_argument("--generate", action="store_true", help="Enable paid creation after confirmation.")
    action.add_argument("--resume", metavar="GENERATION_ID", help="Retrieve existing work without creating again.")
    parser.add_argument("--idempotency-key", help="Keep this key and the exact inputs to recover a request.")
    parser.add_argument("--model", choices=["nano-banana-2", "nano-banana-2-lite", "nano-banana-pro"], default="nano-banana-2")
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--resolution", default="1K", choices=["1K", "2K", "4K"])
    parser.add_argument("--aspect-ratio", default="1:1")
    parser.add_argument("--edit", action="append", metavar="HTTPS_IMAGE_URL", help="Edit this image; repeat for additional references.")
    parser.add_argument("--output", type=Path, default=Path("nano-banana.png"))
    args = parser.parse_args()
    key = os.environ.get("ENTIREFEED_API_KEY")

    if (args.generate or args.resume) and args.output.exists():
        raise SystemExit(f"{args.output} already exists. Choose another --output path.")

    if args.resume:
        if not key:
            raise SystemExit("Set ENTIREFEED_API_KEY to resume.")
        if not re.fullmatch(r"gen_[a-zA-Z0-9]+", args.resume):
            raise SystemExit("Expected a gen_... generation ID.")
        current = api("GET", f"generations/{args.resume}", key=key)
        status_url = f"{API_BASE}/generations/{args.resume}"
    else:
        image_limit = {"nano-banana-2": 14, "nano-banana-2-lite": 10, "nano-banana-pro": 8}[args.model]
        prompt_limit = 10000 if args.model == "nano-banana-pro" else 20000
        if not args.prompt.strip() or len(args.prompt) > prompt_limit:
            raise SystemExit(f"Use a nonblank prompt of at most {prompt_limit} characters.")
        if args.model == "nano-banana-2-lite" and args.resolution != "1K":
            raise SystemExit("Nano Banana 2 Lite supports 1K output only.")
        payload = {
            "task": "image.edit" if args.edit else "image.create",
            "model": args.model,
            "input": {"prompt": args.prompt, "resolution": args.resolution, "aspect_ratio": args.aspect_ratio},
        }
        if args.edit:
            if len(args.edit) > image_limit or any(urlparse(url).scheme != "https" for url in args.edit):
                raise SystemExit(f"Use up to {image_limit} publicly reachable HTTPS image URLs for {args.model}.")
            payload["input"]["images"] = [{"url": url} for url in args.edit]
        estimate = api("POST", "estimate", json=payload)
        print("Current estimate:", json.dumps(estimate["estimated_cost"], indent=2))
        if not args.generate:
            print("Estimate only: no image generated and no balance spent.")
            print("To generate, add --generate --idempotency-key YOUR_UNIQUE_REQUEST_KEY.")
            return
        if not key or not args.idempotency_key:
            raise SystemExit("Set ENTIREFEED_API_KEY and pass --idempotency-key before generating.")
        print("Creation reserves balance using the current price, which can change after this estimate.")
        if input("Type GENERATE to create this paid request: ").strip() != "GENERATE":
            print("No generation created.")
            return
        print(f"Request key: {args.idempotency_key}. Keep it with the unchanged inputs after a network error.")
        current = api("POST", "generations", key=key, json=payload, headers={"Idempotency-Key": args.idempotency_key})
        status_url = current["links"]["self"]

    generation_id = current["id"]
    print(f"Generation: {generation_id}")
    print(f"Resume with: uv run nano-banana.py --resume {generation_id}")
    deadline = time.monotonic() + 900
    while True:
        # Fetch the complete record even if create already returned success.
        current = api("GET", status_url, key=key)
        print("Status:", current["status"])
        if current["status"] not in {"queued", "running"}:
            break
        if time.monotonic() >= deadline:
            raise SystemExit(f"Still active. Resume {generation_id}; do not create another request.")
        time.sleep(5)

    if current["status"] != "succeeded":
        print("Error:", json.dumps(current.get("error"), indent=2))
        print("Usage:", json.dumps(current.get("usage"), indent=2))
        raise SystemExit(f"Stopped at {current['status']}. See Console for details.")

    assets = (current.get("output") or {}).get("assets", [])
    image = next((asset for asset in assets if asset.get("type") == "image"), None)
    if not image or urlparse(image["url"]).scheme != "https":
        raise SystemExit("Succeeded, but no HTTPS image asset was returned.")
    print("Image URL:", image["url"])
    print("MIME type:", image.get("mime_type"))
    print("Expiry:", image.get("expires_at"))
    partial = args.output.with_name(args.output.name + ".part")
    output = partial.open("xb")
    try:
        with output:
            # Separate request: never send the API key to the image host.
            with requests.get(image["url"], stream=True, timeout=120) as response:
                response.raise_for_status()
                for chunk in response.iter_content(chunk_size=1024 * 1024):
                    output.write(chunk)
    except (requests.RequestException, OSError):
        partial.unlink(missing_ok=True)
        raise
    partial.rename(args.output)
    print("Saved:", args.output)
    print("Usage:", json.dumps(current.get("usage"), indent=2))


if __name__ == "__main__":
    try:
        main()
    except requests.RequestException as error:
        raise SystemExit(
            f"Request failed: {error}. Resume the generation ID if available; "
            "otherwise reuse the same idempotency key and unchanged inputs."
        ) from error