"""
Image generation service using Replicate (black-forest-labs/flux-2-klein-4b).
"""

import replicate
from loguru import logger


async def generate_image(prompt: str, aspect_ratio: str = "1:1", resolution: str = "1K") -> dict:
    """Generate an image using Replicate's flux-2-klein-4b model.

    Args:
        prompt: Text description of the image to generate.
        aspect_ratio: Aspect ratio (e.g. "1:1", "16:9", "4:3").
        resolution: Output resolution ("1K", "2K", "4K").

    Returns:
        dict with "url" key on success, or "error" key on failure.
    """
    try:
        logger.info(f"[ImageService] Generating image: {prompt!r} ({aspect_ratio}, {resolution})")

        # output = await replicate.async_run(
        #     "google/nano-banana-pro",
        #     input={
        #         "prompt": prompt,
        #         "resolution": resolution,
        #         "image_input": [],
        #         "aspect_ratio": aspect_ratio,
        #         "output_format": "png",
        #         "safety_filter_level": "block_only_high",
        #     },
        # )

        output = await replicate.async_run(
            "black-forest-labs/flux-2-klein-4b",
            input={
                "prompt": prompt,
                "images": [],
                "go_fast": False,
                "aspect_ratio": aspect_ratio,
                "output_format": "jpg",
                "output_quality": 95,
                "output_megapixels": "1",
            },
        )

        if output:
            url = str(output[0].url)
            logger.info(f"[ImageService] Image generated: {url}")
            return {"url": url, "success": True}

        logger.warning("[ImageService] No image returned from Replicate")
        return {"error": "No image was generated. Please try again.", "success": False}

    except Exception as e:
        logger.error(f"[ImageService] Error generating image: {e}")
        return {"error": str(e), "success": False}