#!/usr/bin/env python3
import argparse
import base64
import json
import os
import sys
import time
import urllib.error
import urllib.request


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate an image with gpt-image-2 through an OpenAI-compatible /v1/images/generations endpoint."
    )
    parser.add_argument("--base-url", default=os.environ.get("BASE_URL", "https://claude.omniclaw.store/v1"))
    parser.add_argument("--api-key-env", default="API_KEY")
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--output", default="image.png")
    parser.add_argument("--model", default="gpt-image-2")
    parser.add_argument("--size", default="1024x1024")
    parser.add_argument("--quality", default="medium", choices=["low", "medium", "high", "auto"])
    parser.add_argument("--format", default="png", choices=["png", "jpeg", "webp"])
    parser.add_argument("--compression", type=int, default=None)
    parser.add_argument("--background", default=None, choices=["auto", "opaque"])
    parser.add_argument("--moderation", default=None, choices=["auto", "low"])
    parser.add_argument("--n", type=int, default=1)
    parser.add_argument("--timeout", type=float, default=300.0)
    parser.add_argument("--dry-run", action="store_true")
    return parser.parse_args()


def build_payload(args: argparse.Namespace) -> dict:
    payload = {
        "model": args.model,
        "prompt": args.prompt,
        "size": args.size,
        "quality": args.quality,
        "output_format": args.format,
        "n": args.n,
    }
    if args.compression is not None:
        payload["output_compression"] = args.compression
    if args.background:
        payload["background"] = args.background
    if args.moderation:
        payload["moderation"] = args.moderation
    return payload


def main() -> int:
    args = parse_args()
    payload = build_payload(args)
    url = args.base_url.rstrip("/") + "/images/generations"

    if args.dry_run:
        print(json.dumps({"url": url, "payload": payload}, ensure_ascii=False, indent=2))
        return 0

    api_key = os.environ.get(args.api_key_env)
    if not api_key:
        print(f"Missing API key env: {args.api_key_env}", file=sys.stderr)
        return 2

    body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    req = urllib.request.Request(
        url,
        data=body,
        method="POST",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
    )

    started = time.monotonic()
    try:
        with urllib.request.urlopen(req, timeout=args.timeout) as resp:
            status = resp.status
            raw = resp.read()
    except urllib.error.HTTPError as exc:
        raw = exc.read()
        print(f"HTTP {exc.code}", file=sys.stderr)
        print(raw.decode("utf-8", "replace")[:4000], file=sys.stderr)
        return 1
    except Exception as exc:
        print(f"request failed: {type(exc).__name__}: {exc}", file=sys.stderr)
        return 1

    elapsed = time.monotonic() - started
    data = json.loads(raw.decode("utf-8"))
    if "error" in data:
        print(json.dumps(data["error"], ensure_ascii=False), file=sys.stderr)
        return 1

    items = data.get("data") or []
    if not items or "b64_json" not in items[0]:
        print("No b64_json image returned", file=sys.stderr)
        print(json.dumps(data, ensure_ascii=False)[:4000], file=sys.stderr)
        return 1

    image = base64.b64decode(items[0]["b64_json"])
    with open(args.output, "wb") as f:
        f.write(image)

    summary = {
        "status": status,
        "elapsed_seconds": round(elapsed, 3),
        "output": args.output,
        "bytes": len(image),
        "model": data.get("model"),
        "size": data.get("size"),
        "quality": data.get("quality"),
        "output_format": data.get("output_format"),
        "usage": data.get("usage"),
    }
    print(json.dumps(summary, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())