import argparse
import base64
import json
import os
import sys
import time
import urllib.request
from pathlib import Path

ap = argparse.ArgumentParser()
ap.add_argument("--input", action="append", default=[])
ap.add_argument("--outdir", required=True)
ap.add_argument("--params", default="{}")
a = ap.parse_args()
p = json.loads(a.params)

if not p.get("prompt"):
    print("ERROR: prompt is required")
    sys.exit(1)
key = os.environ.get("OPENROUTER_API_KEY")
if not key:
    print("ERROR: OPENROUTER_API_KEY not set — add it in Settings")
    sys.exit(1)

model = p.get("model", "google/gemini-2.5-flash-image")
body = {
    "model": model,
    "messages": [{"role": "user", "content": p["prompt"]}],
    "modalities": ["image", "text"],
}
req = urllib.request.Request(
    "https://openrouter.ai/api/v1/chat/completions",
    data=json.dumps(body).encode(),
    headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json",
             "X-Title": "MODELBEAST"},
    method="POST",
)
print(f"calling {model} via OpenRouter ...", flush=True)
t0 = time.time()
try:
    with urllib.request.urlopen(req, timeout=300) as resp:
        result = json.loads(resp.read())
except urllib.error.HTTPError as e:
    print(f"ERROR {e.code}: {e.read().decode()[:1500]}")
    sys.exit(1)
elapsed = round(time.time() - t0, 1)

outdir = Path(a.outdir)
(outdir / "openrouter_result.json").write_text(json.dumps(
    {k: v for k, v in result.items() if k != "choices"} |
    {"note": "choices omitted here; images extracted below"}, indent=2))


def data_urls(obj, acc):
    """Collect data:image/... URLs from any nesting (message.images, content parts)."""
    if isinstance(obj, dict):
        for v in obj.values():
            data_urls(v, acc)
    elif isinstance(obj, list):
        for v in obj:
            data_urls(v, acc)
    elif isinstance(obj, str) and obj.startswith("data:image/"):
        acc.append(obj)


urls: list[str] = []
data_urls(result.get("choices", []), urls)
if not urls:
    text = ""
    try:
        text = result["choices"][0]["message"].get("content") or ""
    except (KeyError, IndexError):
        pass
    print("ERROR: no image in response.", ("Model said: " + str(text)[:800]) if text else "")
    print(json.dumps(result, indent=2)[:1500])
    sys.exit(1)

outputs = []
for i, u in enumerate(urls):
    header, b64 = u.split(",", 1)
    ext = ".png" if "png" in header else ".webp" if "webp" in header else ".jpg"
    name = f"nano_banana_{i}{ext}" if len(urls) > 1 else f"nano_banana{ext}"
    (outdir / name).write_bytes(base64.b64decode(b64))
    outputs.append({"path": name, "meta": {"tool": "openrouter", "model": model,
                                           "seconds": elapsed}})

usage = result.get("usage", {})
print(f"done in {elapsed}s — {len(outputs)} image(s); usage: {json.dumps(usage)}", flush=True)
(outdir / "result.json").write_text(json.dumps({"outputs": outputs}))