#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "httpx>=0.27.0",
#     "python-dotenv>=1.0.0",
# ]
# ///
"""
Generate plan images via OpenRouter (Gemini 3 image-capable models).

Drop-in alternative to generate_gpt_image.py — same CLI signature so the
planf3 workflows need no changes. Uses your existing OPENROUTER_API_KEY
instead of a paid OpenAI key.

Usage:
    python generate_or_image.py "prompt" output.png [options]

Examples:
    python generate_or_image.py "A sunset over mountains" sunset.png
    python generate_or_image.py "Wide architecture diagram" wide.png --size 1536x1024 --quality high

Environment:
    OPENROUTER_API_KEY - Required (get from https://openrouter.ai/settings/keys)
"""

import argparse
import base64
import json
import os
import re
import sys
from pathlib import Path

import httpx
from dotenv import load_dotenv

load_dotenv(Path.cwd() / ".env")

# OpenRouter image-capable models. Gemini 3 Flash Image is cheap + fast;
# Pro Image is higher quality. Pick via --model or OPENROUTER_IMAGE_MODEL env.
DEFAULT_MODEL = os.environ.get("OPENROUTER_IMAGE_MODEL", "google/gemini-3.1-flash-image")
API_BASE = "https://openrouter.ai/api/v1/chat/completions"

# OpenRouter uses this to attribute usage in their dashboard.
HTTP_REFERER = os.environ.get("OPENROUTER_REFERER", "https://github.com/disler/planf3")
APP_TITLE = "planf3"


def parse_size(size: str) -> str:
    """Validate and normalize the size argument. We pass it through to the model
    via the prompt; OpenRouter image models derive dimensions from the prompt
    context, so we keep the explicit size in the request text as a hint."""
    if size == "auto":
        return "auto"
    if not re.match(r"^\d+x\d+$", size):
        raise ValueError(f"invalid size '{size}' — expected WxH (e.g. 1536x1024) or 'auto'")
    return size


def generate(prompt: str, output_path: str, size: str, quality: str, model: str) -> str:
    api_key = os.environ.get("OPENROUTER_API_KEY")
    if not api_key:
        raise EnvironmentError(
            "OPENROUTER_API_KEY environment variable not set. "
            "Get one from https://openrouter.ai/settings/keys and add to .env"
        )

    # Compose the image request. Image-capable Gemini models on OpenRouter
    # return an inline base64 image in the message content when asked.
    size_hint = f" Image dimensions: {size}." if size != "auto" else ""
    quality_hint = f" Quality: {quality}." if quality != "auto" else ""
    user_content = f"Generate a single professional, minimal image:{size_hint}{quality_hint} {prompt}"

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": HTTP_REFERER,
        "X-Title": APP_TITLE,
    }
    payload = {
        "model": model,
        "messages": [{"role": "user", "content": user_content}],
        # request an image back
        "modalities": ["image", "text"],
        # cap output tokens to control cost — image bytes count against this.
        # Gemini image output is ~4 tokens/px so keep this modest.
        "max_tokens": 4096,
    }

    print(f"[generate_or_image] model={model} size={size} quality={quality}", file=sys.stderr)
    print(f"[generate_or_image] prompt: {prompt[:120]}{'...' if len(prompt)>120 else ''}", file=sys.stderr)

    with httpx.Client(timeout=180.0) as client:
        resp = client.post(API_BASE, headers=headers, json=payload)

    if resp.status_code != 200:
        raise RuntimeError(f"OpenRouter API error {resp.status_code}: {resp.text[:400]}")

    data = resp.json()
    message = data.get("choices", [{}])[0].get("message", {})
    content = message.get("content", "")

    # OpenRouter returns image-capable model output in two possible shapes:
    #  1. A list of content parts with type "image_url" (data URI)
    #  2. A markdown string like ![...](data:image/png;base64,...)
    b64_data = None
    mime = "image/png"

    if isinstance(content, list):
        for part in content:
            if isinstance(part, dict):
                if part.get("type") == "image_url":
                    url = part.get("image_url", {}).get("url", "")
                    b64_data, mime = _extract_data_uri(url)
                    break
    elif isinstance(content, str):
        b64_data, mime = _extract_data_uri(content)

    if not b64_data:
        raise RuntimeError(
            "No image returned by model. Response message content:\n"
            + (json.dumps(content)[:500] if content else "(empty)")
        )

    # Write the decoded bytes
    out = Path(output_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_bytes(base64.b64decode(b64_data))
    print(f"[generate_or_image] wrote {out} ({out.stat().st_size} bytes, {mime})", file=sys.stderr)
    return str(out)


def _extract_data_uri(text: str):
    """Pull base64 image data out of a data: URI or a markdown image with a data URI."""
    if not text:
        return None, "image/png"
    m = re.search(r"data:(image/[a-zA-Z+]+);base64,([A-Za-z0-9+/=\s]+)", text)
    if m:
        mime = m.group(1)
        # strip any whitespace the API may have injected
        b64 = re.sub(r"\s+", "", m.group(2))
        return b64, mime
    return None, "image/png"


def main():
    ap = argparse.ArgumentParser(description="Generate a plan image via OpenRouter.")
    ap.add_argument("prompt", help="Image prompt")
    ap.add_argument("output_path", help="Where to save the PNG")
    ap.add_argument("--size", default="1536x1024", help="WxH or 'auto' (default 1536x1024)")
    ap.add_argument("--quality", default="high", choices=["auto", "low", "medium", "high"], help="Quality hint")
    ap.add_argument("--model", default=DEFAULT_MODEL, help=f"OpenRouter model id (default {DEFAULT_MODEL})")
    args = ap.parse_args()

    try:
        size = parse_size(args.size)
        generate(args.prompt, args.output_path, size, args.quality, args.model)
    except Exception as e:
        print(f"ERROR: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()