"""
title: Nano GPT Image Gen
author: Elliott Groves
version: 1.0.1
date: 2024-09-27
description: Nano GPT Image Gen action button for openwebui.
author_url: https://github.com/Orciotrox/Nano-GPT.com_OpenWebUI
funding_url: https://github.com/Orciotrox/Nano-GPT.com_OpenWebUI
nano_address: nano_1pkmodta8fg8ti39pr1doe1mjbwo8cu3c9mt5u38d73d5t57d9nmgmnheifk
icon_url: data:image/svg+xml;base64,IDxzdmcgIHZlcnNpb249IjEuMCIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIiAgd2lkdGg9IjMwMC4wMDAwMDBwdCIgaGVpZ2h0PSIzMDAuMDAwMDAwcHQiIHZpZXdCb3g9IjAgMCAzMDAuMDAwMDAwIDMwMC4wMDAwMDAiICBwcmVzZXJ2ZUFzcGVjdFJhdGlvPSJ4TWlkWU1pZCBtZWV0Ij4gIDxnIHRyYW5zZm9ybT0idHJhbnNsYXRlKDAuMDAwMDAwLDMwMC4wMDAwMDApIHNjYWxlKDAuMTAwMDAwLC0wLjEwMDAwMCkiIGZpbGw9IiM5YjliOWIiIHN0cm9rZT0ibm9uZSI+IDxwYXRoIGQ9Ik0zOTggMjg4MSBjLTE0NSAtNDcgLTI2OSAtMTc0IC0zMTIgLTMyMCAtMTQgLTQ4IC0xNiAtMTc0IC0xNiAtMTA4MiAwIC05OTIgMSAtMTAyOSAyMCAtMTA5MCA1MyAtMTcyIDIwOCAtMzA0IDM4NSAtMzI5IDk3IC0xMyAxOTQ5IC0xMyAyMDM1IDAgMjA3IDMzIDM3MCAxOTggNDAwIDQwNSAxMyA5MSAxMyAxOTIzIDAgMjAyNCAtMjQgMTgyIC0xNTQgMzM4IC0zMjkgMzkxIC02MSAxOSAtOTcgMjAgLTEwOTUgMTkgLTk3MiAwIC0xMDM0IC0xIC0xMDg4IC0xOHogbTIxMTIgLTMxMyBjMzMgLTE3IDUxIC0zNSA2OCAtNjggbDIyIC00NSAtMiAtOTkxIC0zIC05OTEgLTI5IC0zNyBjLTUzIC03MCA5IC02NiAtMTA3OCAtNjYgbC05ODMgMCAtMzkgMjMgYy01NCAzMCAtNzggNzQgLTgzIDE1NSBsLTQgNjcgMjc2IDI0NyBjMzA0IDI3MiAzMDEgMjcwIDQwNSAyNDggNjggLTE0IDEwMiAtOCAxNDkgMjcgNjQgNDkgODAgMTU1IDMyIDIyMSAtNDUgNjMgLTIwMCA5NiAtMzI5IDY4IC0xMTEgLTIzIC0xNzIgLTYzIC0zNTQgLTIyNiAtOTIgLTgzIC0xNjkgLTE1MCAtMTczIC0xNTAgLTMgMCAtNCAzMjEgLTMgNzEzIGwzIDcxMiAyMyAzNyBjMTMgMjAgNDEgNDUgNjUgNTcgNDIgMjEgNDggMjEgMTAxNyAyMSBsOTc2IDAgNDQgLTIyeiIvPiA8cGF0aCBkPSJNMjIwNSAyNDE0IGMtNjkgLTM2IC05MiAtNzMgLTE4MSAtMzA0IC0zMSAtODAgLTU3IC0xNDYgLTU4IC0xNDcgLTEgLTIgLTI0IDUyIC01MSAxMjAgLTg1IDIxNSAtOTYgMjM2IC0xNDMgMjc3IC01NSA0OSAtOTggNTkgLTE0OSAzNCAtNTcgLTI4IC03NSAtNzIgLTcwIC0xNjQgNCAtNjYgMTcgLTEwOCAxMTEgLTM0OSA1OCAtMTUwIDEwNiAtMjc3IDEwNiAtMjgyIDAgLTUgLTgwIC05IC0xNzkgLTkgLTIwMCAwIC0yMzAgLTQgLTI2MSAtMzIgLTI5IC0yNyAtMjIgLTU4IDIwIC04MiAzMSAtMTggNTcgLTIxIDIyNiAtMjYgbDE5MSAtNSAtMjUgLTY3IC0yNSAtNjggLTE0OSAwIGMtMTUyIDAgLTE5OCAtOSAtMjI2IC00MyAtMTkgLTIzIC02IC01NCAzMSAtNzQgMjQgLTEyIDY2IC0xOCAxNjAgLTIxIDExNCAtNCAxMjcgLTYgMTIyIC0yMSAtODMgLTIwNiAtMTU4IC00MTcgLTE2MiAtNDU4IC0xMSAtMTAwIDQ0IC0xODMgMTIwIC0xODMgNDUgMCAxMDUgMzcgMTQxIDg3IDI4IDM5IDg4IDE4MSAxOTcgNDY4IGwxMSAzMCAxNCAtMzAgYzcgLTE2IDQ1IC0xMTEgODMgLTIxMCA5NyAtMjQ5IDk3IC0yNTEgMTM4IC0yOTYgNDYgLTUxIDgxIC02OSAxMzMgLTY5IDc2IDAgMTMwIDg4IDExNSAxODggLTQgMjUgLTQ1IDE0NCAtOTEgMjY1IC00NiAxMjEgLTg0IDIyMSAtODQgMjIzIDAgMiA1MCA0IDExMCA0IDEyMSAwIDE3MCAxMyAxOTAgNTAgMTAgMTggOCAyNiAtMTAgNDUgLTMzIDM1IC01NSAzOSAtMjcxIDQzIGwtNzYgMiAtMjEgNTMgYy0xMSAyOCAtMjMgNjAgLTI3IDcwIC03IDE2IDQgMTcgMTQ2IDE3IDEzMCAwIDE2MSAzIDE5NyAyMCAzMiAxNSA0MiAyNCA0MiA0MyAwIDI1IC0yMCA0OCAtNTUgNjIgLTExIDUgLTkyIDExIC0xODEgMTQgbC0xNjAgNiAxMTQgMjk1IGMxNDIgMzY4IDE1MSA0MzAgNzYgNTA1IC00MCA0MCAtODcgNDYgLTEzOSAxOXoiLz4gPHBhdGggZD0iTTEwNzIgMjI1MiBjLTY2IC0yNCAtMTE4IC02NSAtMTU5IC0xMjYgLTQzIC02MiAtNTYgLTExNiAtNTEgLTE5NiAxMyAtMTc0IDE2MCAtMzAwIDMzNyAtMjg4IDI1OCAxOSAzODIgMzI2IDIxMCA1MjEgLTgxIDkyIC0yMjIgMTMwIC0zMzcgODl6Ii8+IDwvZz4gPC9zdmc+IA==
"""

#Nano GPT Pipeline https://openwebui.com/f/elliott/nano_gpt/
#Integrates nano-gpt.com into OpenWebUI for seamless text-based AI tasks, like document retrieval and generation.

#Nano GPT Image Generation Button https://openwebui.com/f/elliott/nanoimages/
#Allows users to generate images from text prompts with nano-gpt.com.

#Nano GPT Receive Nano Button https://openwebui.com/f/elliott/nanogpt_receive_nano/
#Automates receiving Nano cryptocurrency through OpenWebUI for use with nano-gpt.com.


from pydantic import BaseModel, Field
from typing import Optional, List
import requests
import asyncio


class ImageModel(BaseModel):
    id: str
    name: str
    description: str
    cost: dict
    maxImages: int
    resolutions: list
    engine: str


class Action:
    class Valves(BaseModel):
        NANO_GPT_API_KEY: str = Field(
            default="Your Nano GPT API Key",
            description="API key for authenticating requests to the Nano GPT API.",
        )
        NAME_PREFIX: str = Field(
            default="NanoGPT-Image/",
            description="Prefix to be added before model names.",
        )
        NANO_GPT_API_BASE_URL: str = Field(
            default="https://nano-gpt.com/api",
            description="Base URL for accessing Nano GPT API endpoints.",
        )

    def __init__(self):
        self.valves = self.Valves()
        self.models = []

    async def fetch_models(self) -> List[ImageModel]:
        try:
            headers = {
                "Authorization": f"Bearer {self.valves.NANO_GPT_API_KEY}",
                "Content-Type": "application/json",
            }

            response = requests.get(
                f"{self.valves.NANO_GPT_API_BASE_URL}/models", headers=headers
            )
            response.raise_for_status()

            models = response.json().get("models", {}).get("image", {})
            if not models:
                return []

            return [
                ImageModel(
                    id=model_info.get("model"),
                    name=model_info.get("name", model_key),
                    description=model_info.get("description", ""),
                    cost=model_info.get("cost", {}),
                    maxImages=model_info.get("maxImages", 1),
                    resolutions=[res["value"] for res in model_info.get("resolutions", [])],
                    engine=model_info.get("engine", ""),
                )
                for model_key, model_info in models.items()
            ]

        except requests.RequestException as e:
            print(f"Request failed: {e}")
            return []

    async def action(self, body: dict, __user__=None, __event_emitter__=None, __event_call__=None) -> Optional[dict]:
        print(f"action: {__name__}")

        # Fetch available models
        models = await self.fetch_models()
        if not models:
            print("No models available for selection.")
            return

        # Display model selection options
        selection_text = "\n\n".join([f"{i + 1}. {model.name} - {model.description}" for i, model in enumerate(models)])

        selection_data = {
            "type": "input",
            "data": {
                "title": "Nano GPT Image Models",
                "message": f"Input the number of the model you would like to use:\n\n{selection_text}",
                "placeholder": selection_text,
            },
        }

        try:
            # Get user-selected model
            selected_index = await __event_call__(selection_data)
            if not selected_index or not (1 <= int(selected_index) <= len(models)):
                print("Invalid model selection.")
                return

            selected_model = models[int(selected_index) - 1]
            print(f"Selected Model: {selected_model.name}")

            # Get image width
            width_data = {
                "type": "input",
                "data": {
                    "title": "Image Width",
                    "message": "Enter the width of the image (default is 1024):",
                    "placeholder": "1024",
                },
            }
            width = await __event_call__(width_data)
            width = int(width) if width else 1024

            # Get image height
            height_data = {
                "type": "input",
                "data": {
                    "title": "Image Height",
                    "message": "Enter the height of the image (default is 1024):",
                    "placeholder": "1024",
                },
            }
            height = await __event_call__(height_data)
            height = int(height) if height else 1024

            # Get prompt for image generation
            prompt_data = {
                "type": "input",
                "data": {
                    "title": "Image Prompt",
                    "message": "Enter the prompt for image generation:",
                    "placeholder": "A beautiful landscape",
                },
            }
            image_prompt = await __event_call__(prompt_data)

            # Generate the image
            response_data = self.generate_image(image_prompt, selected_model.id, width, height)
            if not response_data or "data" not in response_data or not response_data["data"]:
                raise Exception("No image data in the response")

            image_data = response_data["data"][0].get("b64_json")
            if not isinstance(image_data, str):
                raise Exception("Invalid image data format")

            content = f"![Generated Image](data:image/png;base64,{image_data})\n\n"

            await __event_emitter__({"type": "message", "data": {"content": content}})
            await __event_emitter__({"type": "image", "data": {"content": image_data}})

            if __event_emitter__:
                await __event_emitter__({"type": "status", "data": {"description": "Image Generated", "done": True}})

        except Exception as e:
            error_message = f"Error processing image: {e}"
            print(error_message)
            if __event_emitter__:
                await __event_emitter__({"type": "status", "data": {"description": "Error Generating Image", "done": True}})
                await __event_emitter__({"type": "message", "data": {"content": error_message}})

    def generate_image(self, prompt, model_id, width, height):
        data = {
            "prompt": prompt,
            "model": model_id,
            "width": width,
            "height": height,
        }
        headers = {
            "x-api-key": self.valves.NANO_GPT_API_KEY,
            "Content-Type": "application/json",
        }

        response = requests.post(
            f"{self.valves.NANO_GPT_API_BASE_URL}/generate-image", headers=headers, json=data
        )

        if response.status_code == 200:
            return response.json()
        else:
            print(f"Error Content: {response.text}")
            return None