#!/usr/bin/env python3
"""OpenRouter Image Generation CLI — ported from Codex image_gen.py.

Supports generate (text-to-image) and edit (image-to-image) via OpenRouter
chat/completions endpoint with GPT Image models.

Features ported from Codex:
- Structured prompt augmentation (--style, --lighting, --palette, etc.)
- Auto-downscale for web assets
- Retry with exponential backoff on transient/rate-limit errors
- Async batch processing from JSONL
- Dry-run mode for payload inspection
"""

from __future__ import annotations

import argparse
import asyncio
import base64
import json
import os
import re
import sys
import time
from io import BytesIO
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

DEFAULT_MODEL = "openai/gpt-5.4-image-2"
DEFAULT_SIZE = "auto"
DEFAULT_QUALITY = "medium"
DEFAULT_OUTPUT_FORMAT = "webp"
DEFAULT_CONCURRENCY = 3
DEFAULT_DOWNSCALE_SUFFIX = "-web"
DEFAULT_MAX_TOKENS = 4096
OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"
TMPFILES_UPLOAD_URL = "https://tmpfiles.org/api/v1/upload"
MAX_IMAGE_BYTES = 50 * 1024 * 1024
MAX_BATCH_JOBS = 100

CRED_PATH = os.path.expanduser("~/.gemini/antigravity/.credentials.json")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _die(message: str, code: int = 1) -> None:
    print(f"Error: {message}", file=sys.stderr)
    raise SystemExit(code)


def _warn(message: str) -> None:
    print(f"Warning: {message}", file=sys.stderr)


def _load_api_key() -> str:
    key = os.getenv("OPENROUTER_API_KEY")
    if key:
        return key
    try:
        with open(CRED_PATH, "r") as f:
            creds = json.load(f)
            key = creds.get("openrouter_api_key")
    except Exception:
        pass
    if not key:
        _die(
            "OpenRouter API key not found. Set OPENROUTER_API_KEY env var "
            "or add openrouter_api_key to ~/.gemini/antigravity/.credentials.json"
        )
    return key


def _ensure_requests():
    try:
        import requests  # noqa: F401
        return requests
    except ImportError:
        _die("requests library is required. Install with: pip install requests")


# ---------------------------------------------------------------------------
# Image upload (local file → tmpfiles.org → public URL)
# ---------------------------------------------------------------------------

def _upload_to_tmpfiles(image_path: str) -> str:
    requests = _ensure_requests()
    print(f"📤 Uploading local image to tmpfiles.org...", file=sys.stderr)
    try:
        with open(image_path, "rb") as f:
            resp = requests.post(TMPFILES_UPLOAD_URL, files={"file": f}, timeout=60)
        resp.raise_for_status()
        data = resp.json()
        if "data" in data and "url" in data["data"]:
            url = data["data"]["url"]
            direct_url = url.replace("tmpfiles.org/", "tmpfiles.org/dl/")
            print(f"✅ Uploaded: {direct_url}", file=sys.stderr)
            return direct_url
    except Exception as e:
        _warn(f"tmpfiles.org upload failed ({e}), falling back to base64")

    # Fallback: inline base64
    import mimetypes
    mime_type, _ = mimetypes.guess_type(image_path)
    if not mime_type:
        mime_type = "image/png"
    with open(image_path, "rb") as f:
        encoded = base64.b64encode(f.read()).decode("utf-8")
    return f"data:{mime_type};base64,{encoded}"


def _resolve_image_url(image_input: str) -> str:
    if image_input.startswith(("http://", "https://", "data:")):
        return image_input
    path = Path(image_input)
    if not path.exists():
        _die(f"Image file not found: {path}")
    if path.stat().st_size > MAX_IMAGE_BYTES:
        _warn(f"Image exceeds 50MB: {path}")
    return _upload_to_tmpfiles(str(path))


# ---------------------------------------------------------------------------
# Prompt augmentation (ported from Codex)
# ---------------------------------------------------------------------------

def _fields_from_args(args: argparse.Namespace) -> Dict[str, Optional[str]]:
    return {
        "use_case": getattr(args, "use_case", None),
        "scene": getattr(args, "scene", None),
        "subject": getattr(args, "subject", None),
        "style": getattr(args, "style", None),
        "composition": getattr(args, "composition", None),
        "lighting": getattr(args, "lighting", None),
        "palette": getattr(args, "palette", None),
        "materials": getattr(args, "materials", None),
        "text": getattr(args, "text", None),
        "constraints": getattr(args, "constraints", None),
        "negative": getattr(args, "negative", None),
    }


def _augment_prompt(augment: bool, prompt: str, fields: Dict[str, Optional[str]]) -> str:
    if not augment:
        return prompt
    sections: List[str] = []
    mapping = [
        ("use_case", "Use case"),
        ("scene", "Scene/background"),
        ("subject", "Subject"),
        ("style", "Style/medium"),
        ("composition", "Composition/framing"),
        ("lighting", "Lighting/mood"),
        ("palette", "Color palette"),
        ("materials", "Materials/textures"),
        ("text", "Text (verbatim)"),
        ("constraints", "Constraints"),
        ("negative", "Avoid"),
    ]
    for key, label in mapping:
        val = fields.get(key)
        if val and key != "text":
            sections.append(f"{label}: {val}")
        elif val and key == "text":
            sections.append(f'{label}: "{val}"')
    if sections:
        return f"Primary request: {prompt}\n" + "\n".join(sections)
    return prompt


# ---------------------------------------------------------------------------
# Retry logic (ported from Codex)
# ---------------------------------------------------------------------------

def _is_transient(status_code: int) -> bool:
    return status_code in (429, 500, 502, 503, 504)


def _extract_retry_after(headers: dict) -> Optional[float]:
    val = headers.get("retry-after") or headers.get("Retry-After")
    if val:
        try:
            return float(val)
        except ValueError:
            pass
    return None


def _call_openrouter(
    api_key: str,
    payload: dict,
    *,
    max_attempts: int = 3,
    timeout: int = 180,
    label: str = "",
) -> dict:
    requests = _ensure_requests()
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://github.com/antigravity-awf",
        "X-Title": "Antigravity Image Gen",
    }
    last_exc = None
    for attempt in range(1, max_attempts + 1):
        try:
            resp = requests.post(
                OPENROUTER_API_URL,
                headers=headers,
                json=payload,
                timeout=timeout,
            )
            if resp.status_code == 200:
                return resp.json()
            if _is_transient(resp.status_code) and attempt < max_attempts:
                sleep_s = _extract_retry_after(dict(resp.headers)) or min(60.0, 2.0 ** attempt)
                print(
                    f"{label} attempt {attempt}/{max_attempts} got {resp.status_code}; "
                    f"retrying in {sleep_s:.1f}s",
                    file=sys.stderr,
                )
                time.sleep(sleep_s)
                continue
            # Non-transient error or last attempt — log response body
            try:
                err_body = resp.text[:500]
            except Exception:
                err_body = "(unreadable)"
            print(f"{label} API error {resp.status_code}: {err_body}", file=sys.stderr)
            resp.raise_for_status()
        except Exception as e:
            last_exc = e
            if attempt == max_attempts:
                raise
            sleep_s = min(60.0, 2.0 ** attempt)
            print(
                f"{label} attempt {attempt}/{max_attempts} failed ({e.__class__.__name__}); "
                f"retrying in {sleep_s:.1f}s",
                file=sys.stderr,
            )
            time.sleep(sleep_s)
    raise last_exc or RuntimeError("unknown error")


async def _call_openrouter_async(
    api_key: str,
    payload: dict,
    *,
    max_attempts: int = 3,
    timeout: int = 180,
    label: str = "",
) -> dict:
    """Async wrapper — runs sync call in executor to avoid blocking event loop."""
    loop = asyncio.get_event_loop()
    return await loop.run_in_executor(
        None,
        lambda: _call_openrouter(
            api_key, payload, max_attempts=max_attempts, timeout=timeout, label=label
        ),
    )


# ---------------------------------------------------------------------------
# Response parsing — extract image from OpenRouter chat/completions response
# ---------------------------------------------------------------------------

def _extract_images_from_response(data: dict) -> List[bytes]:
    """Extract base64 image bytes from OpenRouter response."""
    images: List[bytes] = []
    if "choices" not in data or not data["choices"]:
        return images

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

    # Path 1: message.images[] (OpenAI native format via OpenRouter)
    if "images" in message:
        for img in message["images"]:
            if isinstance(img, dict):
                url = img.get("image_url", {}).get("url", "")
            else:
                url = str(img)
            if url.startswith("data:image"):
                _, encoded = url.split(",", 1)
                images.append(base64.b64decode(encoded))
            elif url.startswith("http"):
                requests = _ensure_requests()
                resp = requests.get(url, timeout=60)
                resp.raise_for_status()
                images.append(resp.content)

    # Path 2: content contains markdown image or inline base64
    if not images and "content" in message:
        content = message["content"] or ""
        # Check for base64 data URIs
        b64_matches = re.findall(r'data:image/[^;]+;base64,([A-Za-z0-9+/=]+)', content)
        for b64 in b64_matches:
            try:
                images.append(base64.b64decode(b64))
            except Exception:
                pass
        # Check for HTTP URLs
        if not images:
            url_matches = re.findall(r'(https?://\S+\.(?:png|jpg|jpeg|webp|gif))', content)
            if url_matches:
                requests = _ensure_requests()
                for url in url_matches[:1]:
                    try:
                        resp = requests.get(url, timeout=60)
                        resp.raise_for_status()
                        images.append(resp.content)
                    except Exception:
                        pass

    return images


# ---------------------------------------------------------------------------
# Output helpers
# ---------------------------------------------------------------------------

def _build_output_paths(
    out: str, output_format: str, count: int, out_dir: Optional[str]
) -> List[Path]:
    ext = "." + output_format
    if out_dir:
        base = Path(out_dir)
        base.mkdir(parents=True, exist_ok=True)
        return [base / f"image_{i}{ext}" for i in range(1, count + 1)]

    out_path = Path(out)
    if out_path.suffix == "":
        out_path = out_path.with_suffix(ext)

    if count == 1:
        return [out_path]
    return [
        out_path.with_name(f"{out_path.stem}-{i}{out_path.suffix}")
        for i in range(1, count + 1)
    ]


def _derive_downscale_path(path: Path, suffix: str) -> Path:
    if suffix and not suffix.startswith(("-", "_")):
        suffix = "-" + suffix
    return path.with_name(f"{path.stem}{suffix}{path.suffix}")


def _downscale_image_bytes(image_bytes: bytes, *, max_dim: int, output_format: str) -> bytes:
    try:
        from PIL import Image
    except ImportError:
        _die("Downscaling requires Pillow. Install with: pip install Pillow")

    with Image.open(BytesIO(image_bytes)) as img:
        img.load()
        w, h = img.size
        scale = min(1.0, float(max_dim) / float(max(w, h)))
        target = (max(1, int(round(w * scale))), max(1, int(round(h * scale))))
        resized = img if target == (w, h) else img.resize(target, Image.Resampling.LANCZOS)

        fmt = output_format.lower()
        if fmt in ("jpg", "jpeg"):
            fmt = "jpeg"
            if resized.mode in ("RGBA", "LA"):
                bg = Image.new("RGB", resized.size, (255, 255, 255))
                bg.paste(resized, mask=resized.split()[-1])
                resized = bg
            else:
                resized = resized.convert("RGB")

        out = BytesIO()
        resized.save(out, format=fmt.upper())
        return out.getvalue()


def _convert_image_format(image_bytes: bytes, output_format: str) -> bytes:
    try:
        from PIL import Image
    except ImportError:
        _die("Format conversion requires Pillow. Install with: pip install Pillow")

    with Image.open(BytesIO(image_bytes)) as img:
        fmt = output_format.lower()
        if fmt in ("jpg", "jpeg"):
            fmt = "jpeg"
            if img.mode in ("RGBA", "LA"):
                bg = Image.new("RGB", img.size, (255, 255, 255))
                bg.paste(img, mask=img.split()[-1])
                img = bg
            else:
                img = img.convert("RGB")

        out = BytesIO()
        img.save(out, format=fmt.upper())
        return out.getvalue()


def _write_images(
    images: List[bytes],
    outputs: List[Path],
    *,
    force: bool = False,
    downscale_max_dim: Optional[int] = None,
    downscale_suffix: str = DEFAULT_DOWNSCALE_SUFFIX,
    output_format: str = DEFAULT_OUTPUT_FORMAT,
) -> None:
    for idx, img_bytes in enumerate(images):
        if idx >= len(outputs):
            break
            
        fmt = output_format.lower()
        is_webp = img_bytes.startswith(b"RIFF") and b"WEBP" in img_bytes[:16]
        is_png = img_bytes.startswith(b"\x89PNG\r\n\x1a\n")
        is_jpeg = img_bytes.startswith(b"\xff\xd8")
        
        needs_convert = False
        if fmt == "webp" and not is_webp:
            needs_convert = True
        elif fmt in ("jpg", "jpeg") and not is_jpeg:
            needs_convert = True
        elif fmt == "png" and not is_png:
            needs_convert = True
            
        if needs_convert:
            try:
                img_bytes = _convert_image_format(img_bytes, output_format)
            except Exception as e:
                _warn(f"Failed to convert image format to {output_format}: {e}")

        out_path = outputs[idx]
        if out_path.exists() and not force:
            _die(f"Output already exists: {out_path} (use --force to overwrite)")
        out_path.parent.mkdir(parents=True, exist_ok=True)
        out_path.write_bytes(img_bytes)
        print(f"Wrote {out_path}")

        if downscale_max_dim is not None:
            derived = _derive_downscale_path(out_path, downscale_suffix)
            if derived.exists() and not force:
                _die(f"Output already exists: {derived} (use --force)")
            resized = _downscale_image_bytes(
                img_bytes, max_dim=downscale_max_dim, output_format=output_format
            )
            derived.write_bytes(resized)
            print(f"Wrote {derived} (downscaled to max {downscale_max_dim}px)")


# ---------------------------------------------------------------------------
# Build OpenRouter payload
# ---------------------------------------------------------------------------

def _build_messages(prompt: str, image_urls: Optional[List[str]] = None) -> list:
    if image_urls:
        content: list = [{"type": "text", "text": prompt}]
        for url in image_urls:
            content.append({"type": "image_url", "image_url": {"url": url}})
        return [{"role": "user", "content": content}]
    return [{"role": "user", "content": prompt}]


def _build_payload(
    model: str,
    messages: list,
    *,
    max_tokens: int = DEFAULT_MAX_TOKENS,
) -> dict:
    return {
        "model": model,
        "messages": messages,
        "max_tokens": max_tokens,
    }


# ---------------------------------------------------------------------------
# Commands: generate, edit, generate-batch
# ---------------------------------------------------------------------------

def _cmd_generate(args: argparse.Namespace) -> None:
    api_key = _load_api_key()
    prompt = _read_prompt(args.prompt, getattr(args, "prompt_file", None))
    fields = _fields_from_args(args)
    prompt = _augment_prompt(args.augment, prompt, fields)

    messages = _build_messages(prompt)
    payload = _build_payload(args.model, messages, max_tokens=args.max_tokens)

    output_format = args.output_format or DEFAULT_OUTPUT_FORMAT
    outputs = _build_output_paths(args.out, output_format, args.n, args.out_dir)

    if args.dry_run:
        print(json.dumps({"outputs": [str(p) for p in outputs], **payload}, indent=2))
        return

    print(f"🎨 [OpenRouter] Generating {args.n} image(s) with {args.model}...", file=sys.stderr)
    print(f"📝 Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}", file=sys.stderr)

    started = time.time()
    images = []
    for i in range(args.n):
        if args.n > 1:
            print(f"⏳ Generating image {i+1}/{args.n}...", file=sys.stderr)
            
        data = _call_openrouter(api_key, payload, max_attempts=args.max_attempts, label=f"[generate {i+1}]" if args.n > 1 else "[generate]")
        batch_images = _extract_images_from_response(data)
        if batch_images:
            images.extend(batch_images)
        else:
            _die(f"No images found in response for image {i+1}.\nResponse: {json.dumps(data, indent=2)}")

    elapsed = time.time() - started
    print(f"✅ Generation completed in {elapsed:.1f}s.", file=sys.stderr)

    _write_images(
        images,
        outputs,
        force=args.force,
        downscale_max_dim=args.downscale_max_dim,
        downscale_suffix=args.downscale_suffix,
        output_format=output_format,
    )

    for p in outputs:
        if p.exists():
            print(f"✨ file://{p}")


def _cmd_edit(args: argparse.Namespace) -> None:
    api_key = _load_api_key()
    prompt = _read_prompt(args.prompt, getattr(args, "prompt_file", None))
    fields = _fields_from_args(args)
    prompt = _augment_prompt(args.augment, prompt, fields)

    # Resolve all input images to URLs
    image_urls = [_resolve_image_url(img) for img in args.image]

    messages = _build_messages(prompt, image_urls)
    payload = _build_payload(args.model, messages, max_tokens=args.max_tokens)

    output_format = args.output_format or DEFAULT_OUTPUT_FORMAT
    outputs = _build_output_paths(args.out, output_format, args.n, args.out_dir)

    if args.dry_run:
        print(json.dumps({"outputs": [str(p) for p in outputs], **payload}, indent=2))
        return

    print(f"🎨 [OpenRouter] Editing with {args.model} ({len(image_urls)} image(s)) x {args.n}...", file=sys.stderr)
    print(f"📝 Prompt: {prompt[:120]}{'...' if len(prompt) > 120 else ''}", file=sys.stderr)

    started = time.time()
    images = []
    for i in range(args.n):
        if args.n > 1:
            print(f"⏳ Editing image {i+1}/{args.n}...", file=sys.stderr)
            
        data = _call_openrouter(api_key, payload, max_attempts=args.max_attempts, label=f"[edit {i+1}]" if args.n > 1 else "[edit]")
        batch_images = _extract_images_from_response(data)
        if batch_images:
            images.extend(batch_images)
        else:
            _die(f"No images found in response for edit {i+1}.\nResponse: {json.dumps(data, indent=2)}")

    elapsed = time.time() - started
    print(f"✅ Edit completed in {elapsed:.1f}s.", file=sys.stderr)

    _write_images(
        images,
        outputs,
        force=args.force,
        downscale_max_dim=args.downscale_max_dim,
        downscale_suffix=args.downscale_suffix,
        output_format=output_format,
    )

    for p in outputs:
        if p.exists():
            print(f"✨ file://{p}")


def _cmd_generate_batch(args: argparse.Namespace) -> None:
    api_key = _load_api_key()
    jobs = _read_jobs_jsonl(args.input)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    base_fields = _fields_from_args(args)
    output_format = args.output_format or DEFAULT_OUTPUT_FORMAT

    print(f"🎨 [OpenRouter Batch] {len(jobs)} jobs, concurrency={args.concurrency}", file=sys.stderr)

    async def run_all():
        sem = asyncio.Semaphore(args.concurrency)
        failed = 0

        async def run_one(i: int, job: dict):
            nonlocal failed
            prompt = str(job["prompt"]).strip()
            job_fields = {k: job.get(k, base_fields.get(k)) for k in base_fields}
            augmented = _augment_prompt(args.augment, prompt, job_fields)

            image_urls = None
            if "image" in job:
                raw_images = job["image"] if isinstance(job["image"], list) else [job["image"]]
                image_urls = [_resolve_image_url(img) for img in raw_images]

            messages = _build_messages(augmented, image_urls)
            payload = _build_payload(args.model, messages, max_tokens=args.max_tokens)

            ext = "." + output_format
            out_path = out_dir / f"{i:03d}-{_slugify(prompt[:60])}{ext}"
            label = f"[job {i}/{len(jobs)}]"

            try:
                async with sem:
                    print(f"{label} starting", file=sys.stderr)
                    started = time.time()
                    data = await _call_openrouter_async(
                        api_key, payload, max_attempts=args.max_attempts, label=label
                    )
                    elapsed = time.time() - started
                    print(f"{label} completed in {elapsed:.1f}s", file=sys.stderr)

                images = _extract_images_from_response(data)
                if images:
                    _write_images(images, [out_path], force=args.force, output_format=output_format)
                else:
                    print(f"{label} no images in response", file=sys.stderr)
                    failed += 1
            except Exception as e:
                print(f"{label} failed: {e}", file=sys.stderr)
                failed += 1
                if args.fail_fast:
                    raise

        tasks = [asyncio.create_task(run_one(i, job)) for i, job in enumerate(jobs, 1)]
        await asyncio.gather(*tasks, return_exceptions=not args.fail_fast)
        return failed

    failed = asyncio.run(run_all())
    if failed:
        raise SystemExit(1)


# ---------------------------------------------------------------------------
# Utility
# ---------------------------------------------------------------------------

def _read_prompt(prompt: Optional[str], prompt_file: Optional[str]) -> str:
    if prompt and prompt_file:
        _die("Use --prompt or --prompt-file, not both.")
    if prompt_file:
        p = Path(prompt_file)
        if not p.exists():
            _die(f"Prompt file not found: {p}")
        return p.read_text(encoding="utf-8").strip()
    if prompt:
        return prompt.strip()
    _die("Missing prompt. Use --prompt or --prompt-file.")
    return ""


def _slugify(value: str) -> str:
    value = value.strip().lower()
    value = re.sub(r"[^a-z0-9]+", "-", value)
    value = re.sub(r"-{2,}", "-", value).strip("-")
    return value[:60] if value else "output"


def _read_jobs_jsonl(path: str) -> List[Dict[str, Any]]:
    p = Path(path)
    if not p.exists():
        _die(f"Input file not found: {p}")
    jobs: List[Dict[str, Any]] = []
    for line_no, raw in enumerate(p.read_text(encoding="utf-8").splitlines(), start=1):
        line = raw.strip()
        if not line or line.startswith("#"):
            continue
        try:
            if line.startswith("{"):
                item = json.loads(line)
            else:
                item = {"prompt": line}
            if "prompt" not in item or not str(item["prompt"]).strip():
                _die(f"Missing prompt at line {line_no}")
            jobs.append(item)
        except json.JSONDecodeError as e:
            _die(f"Invalid JSON on line {line_no}: {e}")
    if not jobs:
        _die("No jobs found in input file.")
    if len(jobs) > MAX_BATCH_JOBS:
        _die(f"Too many jobs ({len(jobs)}). Max is {MAX_BATCH_JOBS}.")
    return jobs


# ---------------------------------------------------------------------------
# CLI parser
# ---------------------------------------------------------------------------

def _add_shared_args(parser: argparse.ArgumentParser) -> None:
    parser.add_argument("--model", default=DEFAULT_MODEL, help="OpenRouter model ID")
    parser.add_argument("--prompt", help="Text prompt for generation")
    parser.add_argument("--prompt-file", help="Read prompt from file")
    parser.add_argument("--n", type=int, default=1, help="Number of images (1-10)")
    parser.add_argument("--max-tokens", type=int, default=DEFAULT_MAX_TOKENS)
    parser.add_argument("--out", default="output.webp", help="Output path")
    parser.add_argument("--out-dir", help="Output directory (for batch)")
    parser.add_argument("--output-format", default=DEFAULT_OUTPUT_FORMAT, choices=["png", "webp", "jpeg"])
    parser.add_argument("--force", action="store_true", help="Overwrite existing files")
    parser.add_argument("--dry-run", action="store_true", help="Print payload without calling API")
    parser.add_argument("--max-attempts", type=int, default=3, help="Retry attempts on transient errors")

    # Prompt augmentation
    parser.add_argument("--augment", dest="augment", action="store_true")
    parser.add_argument("--no-augment", dest="augment", action="store_false")
    parser.set_defaults(augment=True)
    parser.add_argument("--use-case", help="Augment: use case context")
    parser.add_argument("--scene", help="Augment: scene/background description")
    parser.add_argument("--subject", help="Augment: subject description")
    parser.add_argument("--style", help="Augment: art style/medium")
    parser.add_argument("--composition", help="Augment: composition/framing")
    parser.add_argument("--lighting", help="Augment: lighting/mood")
    parser.add_argument("--palette", help="Augment: color palette")
    parser.add_argument("--materials", help="Augment: materials/textures")
    parser.add_argument("--text", help="Augment: text to render (verbatim)")
    parser.add_argument("--constraints", help="Augment: constraints")
    parser.add_argument("--negative", help="Augment: things to avoid")

    # Post-processing: downscale
    parser.add_argument("--downscale-max-dim", type=int, help="Generate additional downscaled copy")
    parser.add_argument("--downscale-suffix", default=DEFAULT_DOWNSCALE_SUFFIX)


def main() -> int:
    parser = argparse.ArgumentParser(
        description="OpenRouter Image Generation CLI (Codex-grade)"
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    # generate
    gen_parser = subparsers.add_parser("generate", help="Create a new image from text")
    _add_shared_args(gen_parser)
    gen_parser.set_defaults(func=_cmd_generate)

    # edit (image-to-image)
    edit_parser = subparsers.add_parser("edit", help="Edit/transform an existing image")
    _add_shared_args(edit_parser)
    edit_parser.add_argument("--image", action="append", required=True, help="Input image path/URL")
    edit_parser.set_defaults(func=_cmd_edit)

    # generate-batch
    batch_parser = subparsers.add_parser("generate-batch", help="Batch generate from JSONL")
    _add_shared_args(batch_parser)
    batch_parser.add_argument("--input", required=True, help="Path to JSONL file")
    batch_parser.add_argument("--concurrency", type=int, default=DEFAULT_CONCURRENCY)
    batch_parser.add_argument("--fail-fast", action="store_true")
    batch_parser.set_defaults(func=_cmd_generate_batch)

    args = parser.parse_args()

    # Validation
    if args.n < 1 or args.n > 10:
        _die("--n must be between 1 and 10")
    if args.max_attempts < 1 or args.max_attempts > 10:
        _die("--max-attempts must be between 1 and 10")
    if args.command == "generate-batch" and not args.out_dir:
        _die("generate-batch requires --out-dir")

    args.func(args)
    return 0


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