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

Takes a user-uploaded image (base64) + edit instruction and sends
them directly to the model. No S3, no prompt rewriting, no description.
Returns the model's output URL directly.
"""

import base64
import replicate
from loguru import logger


async def edit_image(
    image_base64: str,
    instruction: str,
    mime_type: str = "image/jpeg",
    aspect_ratio: str = "1:1",
) -> dict:
    """Edit an image using the user's raw instruction.

    Args:
        image_base64: Base64-encoded image data from the user.
        instruction: The user's edit instruction as-is (e.g. "change background to beach").
        mime_type: MIME type of the image.
        aspect_ratio: Output aspect ratio.

    Returns:
        dict with "url" on success, or "error" on failure.
    """
    try:
        logger.info(f"[EditService] Editing image: {instruction!r}")

        # Build a data URI from the base64 image
        data_uri = f"data:{mime_type};base64,{image_base64}"

        output = await replicate.async_run(
            "black-forest-labs/flux-2-klein-4b",
            input={
                "prompt": instruction,
                "images": [data_uri],
                "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"[EditService] Edit complete: {url}")
            return {"url": url, "success": True}

        logger.warning("[EditService] No output returned from Replicate")
        return {"error": "Image edit returned no result. Try again.", "success": False}

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