import argparse
import base64
import json
import os
import sys
import time
import urllib.error
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,
    "prompt": p["prompt"],
    "resolution": p.get("resolution", "1K"),
    "aspect_ratio": p.get("aspect_ratio", "1:1"),
    "n": int(p.get("n", 1)),
}
if p.get("seed") not in (None, ""):
    body["seed"] = int(p["seed"])

req = urllib.request.Request(
    "https://openrouter.ai/api/v1/images",
    data=json.dumps(body).encode(),
    headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json",
             "X-OpenRouter-Title": "MODELBEAST"},
    method="POST",
)
print(f"calling {model} via OpenRouter Image API ...", 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)
outputs = []
for i, item in enumerate(result.get("data", [])):
    b64 = item.get("b64_json")
    if not b64:
        continue
    media = item.get("media_type", "image/png")
    ext = ".png" if "png" in media else ".webp" if "webp" in media else ".jpg"
    short = model.split("/")[-1]
    name = f"{short}_{i}{ext}" if len(result.get("data", [])) > 1 else f"{short}{ext}"
    (outdir / name).write_bytes(base64.b64decode(b64))
    outputs.append({"path": name, "meta": {"tool": "openrouter", "model": model,
                                           "seconds": elapsed,
                                           "cost_usd": result.get("usage", {}).get("cost")}})

if not outputs:
    print("ERROR: no image in response:")
    print(json.dumps(result, indent=2)[:1500])
    sys.exit(1)

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