generator.py

#!/usr/bin/env python3
"""Image generation/editing pipeline with pluggable API backends."""

import asyncio
import base64
from abc import ABC, abstractmethod
from typing import Optional

import aiohttp
import boto3

import config
import storage

# Edit prompts by category
EDIT_PROMPTS = {
    "car_add_damage": "Add realistic moderate body damage to this car - major scratches or moderate dents on the body panels. Make the damage look natural and real, as if from a minor collision. Preserve existing photo style and lighting as much as possible.",
    "delivery_proof": "In this existing photo, add a food delivery paper bag on the ground right in front of the door. The bag should be a small brown paper bag with a solid plain red sticker on the top keeping it closed. Preserve existing photo style and lighting as much as possible.",
    "receipts": "Edit this receipt image to exactly double the total amount and all intermediate monetary values, including line item prices, subtotal, tax amount, total, and any other amount necessary. Keep everything else in the photo exactly the same.",
    "product": "Edit this photo so this product looks like it arrived from the seller with a realistic issue that a customer would complain about. The issue can be either minor shipping damage (scuff/tear/rip/dent, paint chip, or a small crack) OR a manufacturing defect (uneven paint finish, misaligned seam, or a surface blemish). It still has to look like a real photo.",
}


class ImageGenerator(ABC):
    """Base class for image generation/editing APIs."""

    name: str = "base"
    generator_name: str = "base"  # Used for DB storage
    requires_presigned_url: bool = (
        False  # If True, needs S3 presigned URL instead of bytes
    )

    @abstractmethod
    def edit(self, image_bytes: bytes, prompt: str) -> bytes:
        """Edit an image given raw bytes and a prompt. Returns edited image bytes."""
        raise NotImplementedError

    async def edit_async(
        self,
        image_bytes: bytes,
        prompt: str,
        session: aiohttp.ClientSession = None,
        presigned_url: str = None,
    ) -> bytes:
        """Async version of edit. Override for async implementations."""
        return self.edit(image_bytes, prompt)


class GrokGenerator(ImageGenerator):
    """Grok (xAI) image editing API."""

    name = "grok-imagine-image-beta"
    generator_name = "grok"

    def __init__(self, api_key: str = None, model: str = "grok-imagine-image-beta"):
        self.api_key = api_key or config.GROK_API_KEY
        self.model = model
        self.api_base = "https://api.x.ai/v1"

    def edit(self, image_bytes: bytes, prompt: str) -> bytes:
        """Edit image using Grok API (sync version)."""
        return asyncio.get_event_loop().run_until_complete(
            self.edit_async(image_bytes, prompt)
        )

    async def edit_async(
        self,
        image_bytes: bytes,
        prompt: str,
        session: aiohttp.ClientSession = None,
        presigned_url: str = None,
    ) -> bytes:
        """Edit image using Grok API (async version)."""
        # Convert to base64 data URI
        b64_data = base64.b64encode(image_bytes).decode("utf-8")
        image_uri = f"data:image/jpeg;base64,{b64_data}"

        url = f"{self.api_base}/images/edits"
        payload = {
            "prompt": prompt,
            "image": {"url": image_uri},
            "model": self.model,
        }
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }

        close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            close_session = True

        try:
            async with session.post(
                url,
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=120),
            ) as response:
                if response.status != 200:
                    text = await response.text()
                    raise RuntimeError(f"Grok API error {response.status}: {text}")

                result = await response.json()

            if "data" not in result or len(result["data"]) == 0:
                raise RuntimeError(f"Unexpected Grok response: {result}")

            img_data = result["data"][0]

            if "b64_json" in img_data:
                return base64.b64decode(img_data["b64_json"])
            elif "url" in img_data:
                async with session.get(
                    img_data["url"], timeout=aiohttp.ClientTimeout(total=60)
                ) as img_response:
                    if img_response.status != 200:
                        raise RuntimeError(f"Failed to download from {img_data['url']}")
                    return await img_response.read()

            raise RuntimeError(f"No image data in response: {img_data}")
        finally:
            if close_session:
                await session.close()


class QwenGenerator(ImageGenerator):
    """Qwen Image Edit via Fal AI API."""

    name = "qwen-image-edit-2511"
    generator_name = "qwen"
    requires_presigned_url = True  # Qwen needs image URLs, not raw bytes

    def __init__(self, api_key: str = None):
        self.api_key = api_key or config.FAL_API_KEY
        self.api_url = "https://fal.run/fal-ai/qwen-image-edit-2511"

    def edit(self, image_bytes: bytes, prompt: str) -> bytes:
        """Edit image using Qwen via Fal API (sync version)."""
        return asyncio.get_event_loop().run_until_complete(
            self.edit_async(image_bytes, prompt)
        )

    async def edit_async(
        self,
        image_bytes: bytes,
        prompt: str,
        session: aiohttp.ClientSession = None,
        presigned_url: str = None,
    ) -> bytes:
        """Edit image using Qwen via Fal AI API (async version)."""
        if not presigned_url:
            raise RuntimeError("Qwen generator requires a presigned URL")

        headers = {
            "Authorization": f"Key {self.api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "prompt": prompt,
            "image_urls": [presigned_url],
            "num_inference_steps": 28,
            "guidance_scale": 4.5,
            "output_format": "png",
        }

        close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            close_session = True

        try:
            async with session.post(
                self.api_url,
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=180),
            ) as response:
                if response.status != 200:
                    text = await response.text()
                    raise RuntimeError(f"Qwen/Fal API error {response.status}: {text}")

                result = await response.json()

            # Extract image URL from response
            if "images" not in result or len(result["images"]) == 0:
                raise RuntimeError(f"Unexpected Qwen response: {result}")

            image_url = result["images"][0]["url"]

            # Download the generated image
            async with session.get(
                image_url, timeout=aiohttp.ClientTimeout(total=60)
            ) as img_response:
                if img_response.status != 200:
                    raise RuntimeError(f"Failed to download from {image_url}")
                return await img_response.read()
        finally:
            if close_session:
                await session.close()


class SeedreamGenerator(ImageGenerator):
    """Seedream v4.5 Image Edit via Fal AI API (ByteDance)."""

    name = "seedream-v4.5-1k"
    generator_name = "seedream"
    requires_presigned_url = True  # Seedream needs image URLs, not raw bytes

    def __init__(self, api_key: str = None):
        self.api_key = api_key or config.FAL_API_KEY
        self.api_url = "https://fal.run/fal-ai/bytedance/seedream/v4.5/edit"

    def edit(self, image_bytes: bytes, prompt: str) -> bytes:
        """Edit image using Seedream via Fal API (sync version)."""
        return asyncio.get_event_loop().run_until_complete(
            self.edit_async(image_bytes, prompt)
        )

    async def edit_async(
        self,
        image_bytes: bytes,
        prompt: str,
        session: aiohttp.ClientSession = None,
        presigned_url: str = None,
    ) -> bytes:
        """Edit image using Seedream via Fal AI API (async version)."""
        if not presigned_url:
            raise RuntimeError("Seedream generator requires a presigned URL")

        headers = {
            "Authorization": f"Key {self.api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "prompt": prompt,
            "image_urls": [presigned_url],
            "image_size": {
                "width": 1024,
                "height": 1024,
            },
            "seed": config.SEED,
        }

        close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            close_session = True

        try:
            async with session.post(
                self.api_url,
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=180),
            ) as response:
                if response.status != 200:
                    text = await response.text()
                    raise RuntimeError(
                        f"Seedream/Fal API error {response.status}: {text}"
                    )

                result = await response.json()

            # Extract image URL from response
            if "images" not in result or len(result["images"]) == 0:
                raise RuntimeError(f"Unexpected Seedream response: {result}")

            image_url = result["images"][0]["url"]

            # Download the generated image
            async with session.get(
                image_url, timeout=aiohttp.ClientTimeout(total=60)
            ) as img_response:
                if img_response.status != 200:
                    raise RuntimeError(f"Failed to download from {image_url}")
                return await img_response.read()
        finally:
            if close_session:
                await session.close()


class OpenAIGenerator(ImageGenerator):
    """GPT Image 1.5 via Fal AI API."""

    name = "gpt-image-1.5"
    generator_name = "openai"
    requires_presigned_url = True  # Fal needs image URLs, not raw bytes

    def __init__(self, api_key: str = None):
        self.api_key = api_key or config.FAL_API_KEY
        self.api_url = "https://fal.run/fal-ai/gpt-image-1.5"

    def edit(self, image_bytes: bytes, prompt: str) -> bytes:
        """Edit image using GPT Image 1.5 via Fal API (sync version)."""
        return asyncio.get_event_loop().run_until_complete(
            self.edit_async(image_bytes, prompt)
        )

    async def edit_async(
        self,
        image_bytes: bytes,
        prompt: str,
        session: aiohttp.ClientSession = None,
        presigned_url: str = None,
    ) -> bytes:
        """Edit image using GPT Image 1.5 via Fal API (async version)."""
        if not presigned_url:
            raise RuntimeError("OpenAI/GPT Image generator requires a presigned URL")

        headers = {
            "Authorization": f"Key {self.api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "prompt": prompt,
            "image_urls": [presigned_url],
            "image_size": "1024x1024",
            "quality": "high",
            "input_fidelity": "high",
            "output_format": "png",
        }

        close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            close_session = True

        try:
            async with session.post(
                f"{self.api_url}/edit",
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=180),
            ) as response:
                if response.status != 200:
                    text = await response.text()
                    raise RuntimeError(
                        f"GPT Image/Fal API error {response.status}: {text}"
                    )

                result = await response.json()

            # Extract image URL from response
            if "images" not in result or len(result["images"]) == 0:
                raise RuntimeError(f"Unexpected GPT Image response: {result}")

            image_url = result["images"][0]["url"]

            # Download the generated image
            async with session.get(
                image_url, timeout=aiohttp.ClientTimeout(total=60)
            ) as img_response:
                if img_response.status != 200:
                    raise RuntimeError(f"Failed to download from {image_url}")
                return await img_response.read()
        finally:
            if close_session:
                await session.close()

    async def generate_async(
        self, prompt: str, session: aiohttp.ClientSession = None
    ) -> bytes:
        """Generate image from text prompt only (text-to-image)."""
        headers = {
            "Authorization": f"Key {self.api_key}",
            "Content-Type": "application/json",
        }
        payload = {
            "prompt": prompt,
            "image_size": "1024x1024",
            "quality": "high",
            "output_format": "png",
        }

        close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            close_session = True

        try:
            async with session.post(
                self.api_url,
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=180),
            ) as response:
                if response.status != 200:
                    text = await response.text()
                    raise RuntimeError(
                        f"GPT Image/Fal API error {response.status}: {text}"
                    )

                result = await response.json()

            # Extract image URL from response
            if "images" not in result or len(result["images"]) == 0:
                raise RuntimeError(f"Unexpected GPT Image response: {result}")

            image_url = result["images"][0]["url"]

            # Download the generated image
            async with session.get(
                image_url, timeout=aiohttp.ClientTimeout(total=60)
            ) as img_response:
                if img_response.status != 200:
                    raise RuntimeError(f"Failed to download from {image_url}")
                return await img_response.read()
        finally:
            if close_session:
                await session.close()


class GeminiGenerator(ImageGenerator):
    """Google Gemini image generation/editing API (Nano Banana Pro / Imagen 3)."""

    name = "gemini-3-pro-image-preview"
    generator_name = "gemini"

    def __init__(self, api_key: str = None, model: str = "gemini-3-pro-image-preview"):
        self.api_key = api_key or config.GEMINI_API_KEY
        self.model = model
        self.name = model  # Update name to match model
        self.api_base = "https://generativelanguage.googleapis.com/v1beta"

    def edit(self, image_bytes: bytes, prompt: str) -> bytes:
        """Edit image using Gemini API (sync version)."""
        return asyncio.get_event_loop().run_until_complete(
            self.edit_async(image_bytes, prompt)
        )

    async def edit_async(
        self,
        image_bytes: bytes,
        prompt: str,
        session: aiohttp.ClientSession = None,
        presigned_url: str = None,
    ) -> bytes:
        """Edit image using Gemini API (async version).

        Sends the image + prompt and requests an edited image output.
        """
        # Convert image to base64
        b64_data = base64.b64encode(image_bytes).decode("utf-8")

        # Detect mime type from bytes
        mime_type = "image/jpeg"
        if image_bytes[:8] == b"\x89PNG\r\n\x1a\n":
            mime_type = "image/png"
        elif image_bytes[:2] == b"\xff\xd8":
            mime_type = "image/jpeg"

        url = f"{self.api_base}/models/{self.model}:generateContent"

        payload = {
            "contents": [
                {
                    "parts": [
                        {"text": prompt},
                        {"inline_data": {"mime_type": mime_type, "data": b64_data}},
                    ]
                }
            ],
            "generationConfig": {"responseModalities": ["Image"]},
        }

        headers = {
            "x-goog-api-key": self.api_key,
            "Content-Type": "application/json",
        }

        close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            close_session = True

        try:
            async with session.post(
                url,
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=180),
            ) as response:
                if response.status != 200:
                    text = await response.text()
                    raise RuntimeError(f"Gemini API error {response.status}: {text}")

                result = await response.json()

            # Extract image from response
            if "candidates" not in result or len(result["candidates"]) == 0:
                raise RuntimeError(f"No candidates in Gemini response: {result}")

            candidate = result["candidates"][0]
            if "content" not in candidate or "parts" not in candidate["content"]:
                raise RuntimeError(f"No content parts in Gemini response: {candidate}")

            # Find the image part in the response (API returns camelCase)
            for part in candidate["content"]["parts"]:
                # Check both camelCase (API response) and snake_case (docs)
                if "inlineData" in part:
                    img_data = part["inlineData"]
                    return base64.b64decode(img_data["data"])
                elif "inline_data" in part:
                    img_data = part["inline_data"]
                    return base64.b64decode(img_data["data"])

            raise RuntimeError(
                f"No image found in Gemini response parts: {candidate['content']['parts']}"
            )
        finally:
            if close_session:
                await session.close()

    async def generate_async(
        self, prompt: str, session: aiohttp.ClientSession = None
    ) -> bytes:
        """Generate image from text prompt only (text-to-image)."""
        url = f"{self.api_base}/models/{self.model}:generateContent"

        payload = {
            "contents": [{"parts": [{"text": prompt}]}],
            "generationConfig": {"responseModalities": ["Image"]},
        }

        headers = {
            "x-goog-api-key": self.api_key,
            "Content-Type": "application/json",
        }

        close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            close_session = True

        try:
            async with session.post(
                url,
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=180),
            ) as response:
                if response.status != 200:
                    text = await response.text()
                    raise RuntimeError(f"Gemini API error {response.status}: {text}")

                result = await response.json()

            # Extract image from response
            if "candidates" not in result or len(result["candidates"]) == 0:
                raise RuntimeError(f"No candidates in Gemini response: {result}")

            candidate = result["candidates"][0]
            if "content" not in candidate or "parts" not in candidate["content"]:
                raise RuntimeError(f"No content parts in Gemini response: {candidate}")

            # Find the image part in the response (API returns camelCase)
            for part in candidate["content"]["parts"]:
                # Check both camelCase (API response) and snake_case (docs)
                if "inlineData" in part:
                    img_data = part["inlineData"]
                    return base64.b64decode(img_data["data"])
                elif "inline_data" in part:
                    img_data = part["inline_data"]
                    return base64.b64decode(img_data["data"])

            raise RuntimeError(
                f"No image found in Gemini response parts: {candidate['content']['parts']}"
            )
        finally:
            if close_session:
                await session.close()


class GeminiFalGenerator(ImageGenerator):
    """Google Gemini via Fal.ai (nano-banana-pro model). Same output as GeminiGenerator."""

    name = "gemini-3-pro-image-preview"  # Same as direct Gemini
    generator_name = "gemini"  # Same as direct Gemini
    requires_presigned_url = True  # Fal needs image URLs

    def __init__(self, api_key: str = None):
        self.api_key = api_key or config.FAL_API_KEY
        self.api_url = "https://fal.run/fal-ai/nano-banana-pro/edit"

    def edit(self, image_bytes: bytes, prompt: str) -> bytes:
        """Edit image using Fal API (sync version)."""
        return asyncio.get_event_loop().run_until_complete(
            self.edit_async(image_bytes, prompt)
        )

    async def edit_async(
        self,
        image_bytes: bytes,
        prompt: str,
        session: aiohttp.ClientSession = None,
        presigned_url: str = None,
    ) -> bytes:
        """Edit image using Fal API (async version)."""
        if not presigned_url:
            raise RuntimeError("GeminiFal generator requires a presigned URL")

        headers = {
            "Authorization": f"Key {self.api_key}",
            "Content-Type": "application/json",
        }

        payload = {
            "prompt": prompt,
            "image_urls": [presigned_url],
        }

        close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            close_session = True

        try:
            async with session.post(
                self.api_url,
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=300),
            ) as response:
                if response.status != 200:
                    text = await response.text()
                    raise RuntimeError(f"Fal API error {response.status}: {text}")
                result = await response.json()

            if "images" not in result or len(result["images"]) == 0:
                raise RuntimeError(f"Unexpected Fal response: {result}")

            image_url = result["images"][0]["url"]

            async with session.get(
                image_url, timeout=aiohttp.ClientTimeout(total=60)
            ) as img_resp:
                if img_resp.status != 200:
                    raise RuntimeError(f"Failed to download from {image_url}")
                return await img_resp.read()
        finally:
            if close_session:
                await session.close()

    async def generate_async(
        self, prompt: str, session: aiohttp.ClientSession = None
    ) -> bytes:
        """Generate image from text prompt only (text-to-image)."""
        headers = {
            "Authorization": f"Key {self.api_key}",
            "Content-Type": "application/json",
        }

        payload = {
            "prompt": prompt,
        }

        close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            close_session = True

        try:
            # Use the non-edit endpoint for text-to-image
            t2i_url = "https://fal.run/fal-ai/nano-banana-pro"
            async with session.post(
                t2i_url,
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=300),
            ) as response:
                if response.status != 200:
                    text = await response.text()
                    raise RuntimeError(f"Fal API error {response.status}: {text}")
                result = await response.json()

            if "images" not in result or len(result["images"]) == 0:
                raise RuntimeError(f"Unexpected Fal response: {result}")

            image_url = result["images"][0]["url"]

            async with session.get(
                image_url, timeout=aiohttp.ClientTimeout(total=60)
            ) as img_resp:
                if img_resp.status != 200:
                    raise RuntimeError(f"Failed to download from {image_url}")
                return await img_resp.read()
        finally:
            if close_session:
                await session.close()


class GeminiReplicateGenerator(ImageGenerator):
    """Google Gemini via Replicate (nano-banana-pro model). Same output as GeminiGenerator."""

    name = "gemini-3-pro-image-preview"  # Same as direct Gemini
    generator_name = "gemini"  # Same as direct Gemini

    def __init__(self, api_token: str = None, model: str = "google/nano-banana-pro"):
        self.api_token = api_token or config.REPLICATE_API_TOKEN
        self.model = model
        self.api_base = "https://api.replicate.com/v1"

    def edit(self, image_bytes: bytes, prompt: str) -> bytes:
        """Edit image using Replicate API (sync version)."""
        return asyncio.get_event_loop().run_until_complete(
            self.edit_async(image_bytes, prompt)
        )

    async def edit_async(
        self,
        image_bytes: bytes,
        prompt: str,
        session: aiohttp.ClientSession = None,
        presigned_url: str = None,
    ) -> bytes:
        """Edit image using Replicate API (async version)."""
        b64_data = base64.b64encode(image_bytes).decode("utf-8")
        mime_type = (
            "image/png" if image_bytes[:8] == b"\x89PNG\r\n\x1a\n" else "image/jpeg"
        )
        data_uri = f"data:{mime_type};base64,{b64_data}"

        headers = {
            "Authorization": f"Bearer {self.api_token}",
            "Content-Type": "application/json",
        }

        payload = {
            "input": {
                "prompt": prompt,
                "image_input": [data_uri],
                "resolution": "2K",
                "output_format": "jpg",
            }
        }

        close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            close_session = True

        try:
            # Create prediction and wait for result
            async with session.post(
                f"{self.api_base}/models/{self.model}/predictions",
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=300),
            ) as response:
                if response.status not in (200, 201):
                    text = await response.text()
                    raise RuntimeError(f"Replicate API error {response.status}: {text}")
                result = await response.json()

            # If not completed, poll for result
            if result.get("status") not in ("succeeded", "failed", "canceled"):
                poll_url = (
                    result.get("urls", {}).get("get")
                    or f"{self.api_base}/predictions/{result['id']}"
                )
                for _ in range(120):  # Max 2 minutes polling
                    await asyncio.sleep(1)
                    async with session.get(poll_url, headers=headers) as poll_resp:
                        result = await poll_resp.json()
                        if result.get("status") in ("succeeded", "failed", "canceled"):
                            break

            if result.get("status") != "succeeded":
                raise RuntimeError(f"Replicate prediction failed: {result}")

            # Download the output image
            output_url = result.get("output")
            if isinstance(output_url, list):
                output_url = output_url[0] if output_url else None
            if not output_url:
                raise RuntimeError(f"No output URL in Replicate response: {result}")

            async with session.get(output_url) as img_resp:
                return await img_resp.read()
        finally:
            if close_session:
                await session.close()

    async def generate_async(
        self, prompt: str, session: aiohttp.ClientSession = None
    ) -> bytes:
        """Generate image from text prompt only (text-to-image)."""
        headers = {
            "Authorization": f"Bearer {self.api_token}",
            "Content-Type": "application/json",
        }

        payload = {
            "input": {
                "prompt": prompt,
                "resolution": "2K",
                "output_format": "jpg",
            }
        }

        close_session = False
        if session is None:
            session = aiohttp.ClientSession()
            close_session = True

        try:
            async with session.post(
                f"{self.api_base}/models/{self.model}/predictions",
                json=payload,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=300),
            ) as response:
                if response.status not in (200, 201):
                    text = await response.text()
                    raise RuntimeError(f"Replicate API error {response.status}: {text}")
                result = await response.json()

            if result.get("status") not in ("succeeded", "failed", "canceled"):
                poll_url = (
                    result.get("urls", {}).get("get")
                    or f"{self.api_base}/predictions/{result['id']}"
                )
                for _ in range(120):
                    await asyncio.sleep(1)
                    async with session.get(poll_url, headers=headers) as poll_resp:
                        result = await poll_resp.json()
                        if result.get("status") in ("succeeded", "failed", "canceled"):
                            break

            if result.get("status") != "succeeded":
                raise RuntimeError(f"Replicate prediction failed: {result}")

            output_url = result.get("output")
            if isinstance(output_url, list):
                output_url = output_url[0] if output_url else None
            if not output_url:
                raise RuntimeError(f"No output URL in Replicate response: {result}")

            async with session.get(output_url) as img_resp:
                return await img_resp.read()
        finally:
            if close_session:
                await session.close()


# Registry of available generators
GENERATORS = {
    "grok": GrokGenerator,
    "qwen": QwenGenerator,
    "seedream": SeedreamGenerator,
    "openai": OpenAIGenerator,
    "gemini": GeminiGenerator,
    "gemini-fal": GeminiFalGenerator,
    "gemini-replicate": GeminiReplicateGenerator,
}


def get_generator(name: str, **kwargs) -> ImageGenerator:
    """Get a generator instance by name."""
    if name not in GENERATORS:
        raise ValueError(
            f"Unknown generator: {name}. Available: {list(GENERATORS.keys())}"
        )
    return GENERATORS[name](**kwargs)


def process_image(
    image_id: str,
    generator: ImageGenerator,
    s3_client=None,
    dry_run: bool = False,
) -> Optional[str]:
    """Process a single image: download from S3, edit, upload result, update DB.

    Returns the S3 path of the generated image, or None on failure.
    """
    # Get image metadata from DB
    img = storage.get_image_by_id(image_id)
    if not img:
        print(f"  Image not found in DB: {image_id}")
        return None

    category = img["category"]
    source = img["source"]
    s3_processed_path = img.get("s3_processed_path")

    if not s3_processed_path:
        print(f"  No S3 path for image: {image_id}")
        return None

    # Get prompt for category
    prompt = EDIT_PROMPTS.get(category)
    if not prompt:
        print(f"  No prompt for category: {category}")
        return None

    # Initialize S3 client if needed
    if s3_client is None:
        s3_client = boto3.client("s3", region_name=config.S3_REGION)

    # Parse S3 path
    if s3_processed_path.startswith("s3://"):
        # Full S3 URI
        parts = s3_processed_path.replace("s3://", "").split("/", 1)
        bucket = parts[0]
        key = parts[1]
    else:
        # Relative path - assume it's under our bucket
        bucket = config.S3_BUCKET
        key = s3_processed_path.lstrip("/")
        # Handle local-style paths like "data/processed/..."
        if key.startswith("data/"):
            key = key[5:]  # Remove "data/" prefix

    print(f"  Downloading from s3://{bucket}/{key}")

    if dry_run:
        print(f"  [DRY RUN] Would process {image_id}")
        return "DRY_RUN"

    image_bytes = None
    presigned_url = None

    if generator.requires_presigned_url:
        # Generate presigned URL for generators that need it (e.g., Qwen)
        try:
            presigned_url = s3_client.generate_presigned_url(
                "get_object",
                Params={"Bucket": bucket, "Key": key},
                ExpiresIn=300,
            )
        except Exception as e:
            print(f"  Presigned URL failed: {e}")
            return None
    else:
        # Download image from S3
        try:
            response = s3_client.get_object(Bucket=bucket, Key=key)
            image_bytes = response["Body"].read()
        except Exception as e:
            print(f"  Failed to download: {e}")
            return None

    # Edit image (use async version via asyncio.run for presigned_url support)
    print(f"  Sending to {generator.name}...")
    try:
        edited_bytes = asyncio.get_event_loop().run_until_complete(
            generator.edit_async(image_bytes, prompt, presigned_url=presigned_url)
        )
    except Exception as e:
        print(f"  Generation failed: {e}")
        return None

    # Upload to S3
    output_key = f"generated/{category}/{generator.name}/{source}/{image_id}_edited.png"
    output_s3_path = f"s3://{config.S3_BUCKET}/{output_key}"

    print(f"  Uploading to {output_s3_path}")
    try:
        s3_client.put_object(
            Bucket=config.S3_BUCKET,
            Key=output_key,
            Body=edited_bytes,
            ContentType="image/png",
        )
    except Exception as e:
        print(f"  Failed to upload: {e}")
        return None

    # Update DB - insert into generated_images table
    storage.insert_generated_image(
        base_image_id=image_id,
        generator=generator.generator_name,
        model=generator.name,
        s3_path=output_s3_path,
        prompt=prompt,
    )

    return output_s3_path


async def process_image_async(
    image_id: str,
    generator: ImageGenerator,
    s3_client,
    session: aiohttp.ClientSession,
    dry_run: bool = False,
) -> Optional[str]:
    """Process a single image asynchronously."""
    # Get image metadata from DB
    img = storage.get_image_by_id(image_id)
    if not img:
        return None

    category = img["category"]
    source = img["source"]
    s3_processed_path = img.get("s3_processed_path")

    if not s3_processed_path:
        return None

    prompt = EDIT_PROMPTS.get(category)
    if not prompt:
        return None

    # Parse S3 path
    if s3_processed_path.startswith("s3://"):
        parts = s3_processed_path.replace("s3://", "").split("/", 1)
        bucket = parts[0]
        key = parts[1]
    else:
        bucket = config.S3_BUCKET
        key = s3_processed_path.lstrip("/")
        if key.startswith("data/"):
            key = key[5:]

    if dry_run:
        return "DRY_RUN"

    loop = asyncio.get_event_loop()
    image_bytes = None
    presigned_url = None

    if generator.requires_presigned_url:
        # Generate presigned URL for generators that need it (e.g., Qwen)
        try:
            presigned_url = await loop.run_in_executor(
                None,
                lambda: s3_client.generate_presigned_url(
                    "get_object",
                    Params={"Bucket": bucket, "Key": key},
                    ExpiresIn=300,
                ),
            )
        except Exception as e:
            print(f"  [{image_id}] Presigned URL failed: {e}")
            return None
    else:
        # Download from S3 (run in thread pool to not block)
        try:
            response = await loop.run_in_executor(
                None, lambda: s3_client.get_object(Bucket=bucket, Key=key)
            )
            image_bytes = response["Body"].read()
        except Exception as e:
            print(f"  [{image_id}] Download failed: {e}")
            return None

    # Edit image (async)
    try:
        edited_bytes = await generator.edit_async(
            image_bytes, prompt, session, presigned_url=presigned_url
        )
    except Exception as e:
        print(f"  [{image_id}] Generation failed: {e}")
        return None

    # Upload to S3 (run in thread pool)
    output_key = f"generated/{category}/{generator.name}/{source}/{image_id}_edited.png"
    output_s3_path = f"s3://{config.S3_BUCKET}/{output_key}"

    try:
        await loop.run_in_executor(
            None,
            lambda: s3_client.put_object(
                Bucket=config.S3_BUCKET,
                Key=output_key,
                Body=edited_bytes,
                ContentType="image/png",
            ),
        )
    except Exception as e:
        print(f"  [{image_id}] Upload failed: {e}")
        return None

    # Update DB - insert into generated_images table
    storage.insert_generated_image(
        base_image_id=image_id,
        generator=generator.generator_name,
        model=generator.name,
        s3_path=output_s3_path,
        prompt=prompt,
    )

    return output_s3_path


async def process_batch_async(
    images: list,
    generator: ImageGenerator,
    s3_client,
    session: aiohttp.ClientSession,
    dry_run: bool = False,
) -> tuple:
    """Process a batch of images concurrently."""
    tasks = [
        process_image_async(img["id"], generator, s3_client, session, dry_run)
        for img in images
    ]
    results = await asyncio.gather(*tasks, return_exceptions=True)

    success = 0
    failed = 0
    for img, result in zip(images, results):
        if isinstance(result, Exception):
            print(f"  [{img['id']}] Error: {result}")
            failed += 1
        elif result:
            print(f"  [{img['id']}] SUCCESS")
            success += 1
        else:
            print(f"  [{img['id']}] FAILED")
            failed += 1

    return success, failed


def process_all_validated(
    generator_name: str,
    category: str = None,
    limit: int = None,
    dry_run: bool = False,
    concurrency: int = 5,
    group: int = None,
):
    """Process all validated images with parallel execution."""
    asyncio.run(
        process_all_validated_async(
            generator_name=generator_name,
            category=category,
            limit=limit,
            dry_run=dry_run,
            concurrency=concurrency,
            group=group,
        )
    )


async def process_all_validated_async(
    generator_name: str,
    category: str = None,
    limit: int = None,
    dry_run: bool = False,
    concurrency: int = 5,
    group: int = None,
):
    """Process all validated images with parallel execution (async version)."""
    generator = get_generator(generator_name)
    s3_client = boto3.client("s3", region_name=config.S3_REGION)

    # Get validated images that don't have a generated version for this generator/model
    images = storage.get_images_without_generation(
        generator=generator.generator_name,
        model=generator.name,
        category=category,
        group=group,
    )

    if limit:
        images = images[:limit]

    print(
        f"Processing {len(images)} images with {generator.name} (concurrency={concurrency})"
    )
    print(f"{'=' * 60}")

    total_success = 0
    total_failed = 0

    async with aiohttp.ClientSession() as session:
        # Process in batches
        for i in range(0, len(images), concurrency):
            batch = images[i : i + concurrency]
            batch_num = i // concurrency + 1
            total_batches = (len(images) + concurrency - 1) // concurrency

            print(f"\nBatch {batch_num}/{total_batches} ({len(batch)} images)")

            success, failed = await process_batch_async(
                batch, generator, s3_client, session, dry_run
            )
            total_success += success
            total_failed += failed

    print(f"\n{'=' * 60}")
    print(f"Done! Success: {total_success}, Failed: {total_failed}")


def main():
    import argparse

    parser = argparse.ArgumentParser(
        description="Generate edited images from validated base images"
    )
    parser.add_argument("--generator", required=True, choices=list(GENERATORS.keys()))
    parser.add_argument(
        "--category",
        "-c",
        nargs="+",
        help="Process only these categories (space-separated)",
    )
    parser.add_argument(
        "--limit", "-l", type=int, help="Limit number of images to process"
    )
    parser.add_argument(
        "--dry-run",
        "-n",
        action="store_true",
        help="Don't actually process, just show what would be done",
    )
    parser.add_argument("--image-id", "-i", help="Process a single image by ID")
    parser.add_argument(
        "--concurrency",
        "-p",
        type=int,
        default=5,
        help="Number of parallel requests (default: 5)",
    )
    parser.add_argument(
        "--group",
        "-g",
        type=int,
        default=None,
        help="Only process images in this group (0, 1, 2, etc.)",
    )

    args = parser.parse_args()

    if args.image_id:
        generator = get_generator(args.generator)
        result = process_image(args.image_id, generator, dry_run=args.dry_run)
        if result:
            print(f"Generated: {result}")
        else:
            print("Failed")
    else:
        process_all_validated(
            generator_name=args.generator,
            category=args.category,
            limit=args.limit,
            dry_run=args.dry_run,
            concurrency=args.concurrency,
            group=args.group,
        )


if __name__ == "__main__":
    main()