"""
title: DALL·E Azure Manifold
description: A manifold function to integrate Azure's DALL-E 3 models into Open WebUI. Get your API Key & Endpoint in the Azure AI Studio.
author: Raskoll
version: 0.1.0
license: MIT
requirements: pydantic
environment_variables: AZURE_OPENAI_API_KEY, AZURE_ENDPOINT
"""

import os
from typing import Iterator, List, Union
from open_webui.utils.misc import get_last_user_message
from pydantic import BaseModel, Field
import json
from openai import AzureOpenAI

#Thanks bro https://openwebui.com/f/bgeneto/dall_e/

class Pipe:
    """Azure DALL-E 3 pipeline"""

    class Valves(BaseModel):
        AZURE_ENDPOINT: str = Field(
            default="https://yourAzureEndpoint.openai.azure.com/",
            description="Azure OpenAI API Endpoint",
        )
        AZURE_OPENAI_API_KEY: str = Field(
            default="", description="Your Azure OpenAI API key"
        )
        IMAGE_SIZE: str = Field(default="1024x1024", description="Generated image size")
        NUM_IMAGES: int = Field(default=1, description="Number of images to generate")

    def __init__(self):
        self.type = "manifold"
        self.id = "DALL_E_Azure"
        self.name = "DALL-E Azure"
        self.valves = self.Valves(
            AZURE_OPENAI_API_KEY=os.getenv("AZURE_OPENAI_API_KEY", ""),
            AZURE_ENDPOINT=os.getenv(
                "AZURE_ENDPOINT", "https://yourAzureEndpoint.openai.azure.com/"
            ),
        )

        self.client = AzureOpenAI(
            api_version="2024-05-01-preview",
            azure_endpoint=self.valves.AZURE_ENDPOINT,
            api_key=self.valves.AZURE_OPENAI_API_KEY,
        )

    def get_openai_assistants(self) -> List[dict]:
        """Get the available ImageGen models from Azure OpenAI"""

        if self.valves.AZURE_OPENAI_API_KEY:
            self.client = AzureOpenAI(
                api_version="2024-05-01-preview",
                azure_endpoint=self.valves.AZURE_ENDPOINT,
                api_key=self.valves.AZURE_OPENAI_API_KEY,
            )

            return [{"id": "dall-e-3", "name": ""}]

        return []

    def pipes(self) -> List[dict]:
        return self.get_openai_assistants()

    def pipe(self, body: dict) -> Union[str, Iterator[str]]:
        if not self.valves.AZURE_OPENAI_API_KEY:
            return "Error: AZURE_OPENAI_API_KEY is not set"

        self.client = AzureOpenAI(
            api_version="2024-05-01-preview",
            azure_endpoint=self.valves.AZURE_ENDPOINT,
            api_key=self.valves.AZURE_OPENAI_API_KEY,
        )

        # Extract the user message (prompt) from the body
        user_message = get_last_user_message(body["messages"])

        try:
            # Send the request to Azure DALL-E 3 API
            response = self.client.images.generate(
                model="dall-e-3",  # The model ID for DALL-E 3 in Azure
                prompt=user_message,
                n=self.valves.NUM_IMAGES,
            )

            # Debugging: print raw response for inspection
            print("Raw API Response:", response)

            # Attempt to parse the response as JSON
            image_urls = json.loads(response.model_dump_json())["data"]

            message = ""
            for image in image_urls:
                if image["url"]:
                    message += "![image](" + image["url"] + ")\n"

            yield message

        except json.JSONDecodeError as e:
            return f"Error parsing JSON response: {str(e)}. Raw response: {response}"

        except Exception as e:
            return f"An error occurred: {str(e)}. Raw response: {response}"