"""
title: DALL·E Manifold
description: A manifold function to integrate OpenAI's DALL-E models into Open WebUI.
author: bgeneto (based on Marc Lopez pipeline)
funding_url: https://github.com/open-webui
version: 0.1.4
license: MIT
requirements: pydantic
environment_variables: OPENAI_API_BASE_URL, OPENAI_IMG_API_KEY
"""

import os
from typing import Iterator, List, Union

from open_webui.utils.misc import get_last_user_message
from openai import OpenAI
from pydantic import BaseModel, Field


class Pipe:
    """OpenAI ImageGen pipeline"""

    class Valves(BaseModel):

        OPENAI_API_BASE_URL: str = Field(
            default="https://api.openai.com/v1", description="OpenAI API Base URL"
        )
        OPENAI_IMG_API_KEY: str = Field(default="", description="your 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"
        self.name = "DALL-E"
        self.valves = self.Valves(
            **{
                "OPENAI_IMG_API_KEY": os.getenv("OPENAI_IMG_API_KEY", ""),
                "OPENAI_API_BASE_URL": os.getenv(
                    "OPENAI_API_BASE_URL", "https://api.openai.com/v1"
                ),
            }
        )

        self.client = OpenAI(
            base_url=self.valves.OPENAI_API_BASE_URL,
            api_key=self.valves.OPENAI_IMG_API_KEY,
        )

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

        Returns:
            List[dict]: The list of ImageGen models
        """

        if self.valves.OPENAI_IMG_API_KEY:
            self.client = OpenAI(
                base_url=self.valves.OPENAI_API_BASE_URL,
                api_key=self.valves.OPENAI_IMG_API_KEY,
            )

            models = self.client.models.list()
            return [
                {
                    "id": model.id,
                    "name": model.id,
                }
                for model in models
                if "dall-e" in model.id
            ]

        return []

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

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

        self.client = OpenAI(
            base_url=self.valves.OPENAI_API_BASE_URL,
            api_key=self.valves.OPENAI_IMG_API_KEY,
        )

        model_id = body["model"]
        model_id = model_id.split(".")[1]
        user_message = get_last_user_message(body["messages"])

        response = self.client.images.generate(
            model=model_id,
            prompt=user_message,
            size=self.valves.IMAGE_SIZE,
            n=self.valves.NUM_IMAGES,
        )

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

        yield message