import io
import logging
from enum import Enum

from prisma.models import AgentGraph
from replicate.client import Client as ReplicateClient
from replicate.exceptions import ReplicateError
from replicate.helpers import FileOutput

from backend.data.graph import GraphBaseMeta
from backend.data.model import CredentialsMetaInput, ProviderName
from backend.integrations.credentials_store import ideogram_credentials
from backend.util.request import Requests
from backend.util.settings import Settings

logger = logging.getLogger(__name__)
settings = Settings()


class ImageSize(str, Enum):
    LANDSCAPE = "1024x768"


class ImageStyle(str, Enum):
    DIGITAL_ART = "digital art"


async def generate_agent_image(agent: GraphBaseMeta | AgentGraph) -> io.BytesIO:
    if settings.config.use_agent_image_generation_v2:
        return await generate_agent_image_v2(graph=agent)
    else:
        return await generate_agent_image_v1(agent=agent)


async def generate_agent_image_v2(graph: GraphBaseMeta | AgentGraph) -> io.BytesIO:
    """
    Generate an image for an agent using Ideogram model.
    Returns:
        str: The URL of the generated image
    """
    if not ideogram_credentials.api_key:
        raise ValueError("Missing Ideogram API key")

    from backend.blocks.ideogram import (
        AspectRatio,
        ColorPalettePreset,
        IdeogramModelBlock,
        IdeogramModelName,
        MagicPromptOption,
        StyleType,
        UpscaleOption,
    )

    name = graph.name
    description = f"{name} ({graph.description})" if graph.description else name

    prompt = (
        "Create a visually striking retro-futuristic vector pop art illustration "
        f'prominently featuring "{name}" in bold typography. The image clearly and '
        f"literally depicts a {description}, along with recognizable objects directly "
        f"associated with the primary function of a {name}. "
        f"Ensure the imagery is concrete, intuitive, and immediately understandable, "
        f"clearly conveying the purpose of a {name}. "
        "Maintain vibrant, limited-palette colors, sharp vector lines, "
        "geometric shapes, flat illustration techniques, and solid colors "
        "without gradients or shading. Preserve a retro-futuristic aesthetic "
        "influenced by mid-century futurism and 1960s psychedelia, "
        "prioritizing clear visual storytelling and thematic clarity above all else."
    )

    custom_colors = [
        "#000030",
        "#1C0C47",
        "#9900FF",
        "#4285F4",
        "#FFFFFF",
    ]

    # Run the Ideogram model block with the specified parameters
    url = await IdeogramModelBlock().run_once(
        IdeogramModelBlock.Input(
            credentials=CredentialsMetaInput(
                id=ideogram_credentials.id,
                provider=ProviderName.IDEOGRAM,
                title=ideogram_credentials.title,
                type=ideogram_credentials.type,
            ),
            prompt=prompt,
            ideogram_model_name=IdeogramModelName.V3,
            aspect_ratio=AspectRatio.ASPECT_16_9,
            magic_prompt_option=MagicPromptOption.OFF,
            style_type=StyleType.AUTO,
            upscale=UpscaleOption.NO_UPSCALE,
            color_palette_name=ColorPalettePreset.NONE,
            custom_color_palette=custom_colors,
            seed=None,
            negative_prompt=None,
        ),
        "result",
        credentials=ideogram_credentials,
    )
    response = await Requests().get(url)
    return io.BytesIO(response.content)


async def generate_agent_image_v1(agent: GraphBaseMeta | AgentGraph) -> io.BytesIO:
    """
    Generate an image for an agent using Flux model via Replicate API.

    Args:
        agent (GraphBaseMeta | AgentGraph): The agent to generate an image for

    Returns:
        io.BytesIO: The generated image as bytes
    """
    try:
        if not settings.secrets.replicate_api_key:
            raise ValueError("Missing Replicate API key in settings")

        # Construct prompt from agent details
        prompt = (
            "Create a visually engaging app store thumbnail for the AI agent "
            "that highlights what it does in a clear and captivating way:\n"
            f"- **Name**: {agent.name}\n"
            f"- **Description**: {agent.description}\n"
            f"Focus on showcasing its core functionality with an appealing design."
        )

        # Set up Replicate client
        client = ReplicateClient(api_token=settings.secrets.replicate_api_key)

        # Model parameters
        input_data = {
            "prompt": prompt,
            "width": 1024,
            "height": 768,
            "aspect_ratio": "4:3",
            "output_format": "jpg",
            "output_quality": 90,
            "num_inference_steps": 30,
            "guidance": 3.5,
            "negative_prompt": "blurry, low quality, distorted, deformed",
            "disable_safety_checker": True,
        }

        try:
            # Run model
            output = client.run("black-forest-labs/flux-1.1-pro", input=input_data)

            # Depending on the model output, extract the image URL or bytes
            # If the output is a list of FileOutput or URLs
            if isinstance(output, list) and output:
                if isinstance(output[0], FileOutput):
                    image_bytes = output[0].read()
                else:
                    # If it's a URL string, fetch the image bytes
                    result_url = output[0]
                    response = await Requests().get(result_url)
                    image_bytes = response.content
            elif isinstance(output, FileOutput):
                image_bytes = output.read()
            elif isinstance(output, str):
                # Output is a URL
                response = await Requests().get(output)
                image_bytes = response.content
            else:
                raise RuntimeError("Unexpected output format from the model.")

            return io.BytesIO(image_bytes)

        except ReplicateError as e:
            if e.status == 401:
                raise RuntimeError("Invalid Replicate API token") from e
            raise RuntimeError(f"Replicate API error: {str(e)}") from e

    except Exception as e:
        logger.exception("Failed to generate agent image")
        raise RuntimeError(f"Image generation failed: {str(e)}")
