"""
title: Black Forest Labs API - Flux Pro Image Generator
author: Henning Kosmalla
author: https://github.com/hkosm
version: 0.0.1
license: MIT
description: BFL API Docs: https://api.bfl.ml/docs | images are not stored inside openwebui
"""

import time
import requests


# Replace  with your own API key
BFL_API_KEY = ""


class Tools:
    def __init__(self):
        pass

    async def create_flux_image(
            self,
            prompt: str,
            image_format: str,
            __event_emitter__=None,
    ) -> str:
        """
        This Tool creates visually stunning images with text prompts using the black forest labs API with the Flux.1-pro model.

        :param prompt: the prompt to generate the image
        :param image_format: either 'default' for a square image, 'landscape' for a landscape format or 'portrait' for a portrait of mobile format
        """

        await __event_emitter__(
            {
                "type": "status",
                "data": {"description": "Creating FLUX Image...", "done": False},
            }
        )

        try:
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {"description": "Creating FLUX Image...", "done": False},
                }
            )

            request_id = send_image_generation_request(
                prompt=prompt, image_format=image_format, steps=25
            )
            print(f"flux image request: {request_id}")
            image_url = poll_result(request_id)
            print(f"flux image url: {image_url}")

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

            await __event_emitter__(
                {
                    "type": "message",
                    "data": {"content": f"![Image]({image_url}) \n"},
                }
            )
            await __event_emitter__(
                {
                    "type": "message",
                    "data": {"content": f"Prompt: {prompt} \n"},
                }
            )

            return f"Only tell the user that the image was successfully generated with the format '{image_format}'. Do not show any links."

        except Exception as e:
            await __event_emitter__(
                {
                    "type": "status",
                    "data": {"description": f"An error occured: {e}", "done": True},
                }
            )

            return f"Tell the user that an error occurred and the image generation was not successful. Reason: {e}"


FORMATS = {
    "default": (1024, 1024),
    "square": (1024, 1024),
    "landscape": (1024, 768),
    "landscape_large": (1440, 1024),
    "portrait": (768, 1024),
    "portrait_large": (1024, 1440),
}


def send_image_generation_request(
        prompt: str, image_format: str = "default", steps: int = 30
) -> str:
    """
    Send a request to the BFL API to generate an image based on the given prompt.
    :param prompt: The prompt to generate the image from.
    :param image_format: The format of the image to generate.
    Must be one of "default", "square", "landscape", "landscape_large", "portrait", "portrait_large".
    :param steps: The number of steps to generate the image.
    :return:
    """

    width, heigth = FORMATS[image_format]
    request = requests.post(
        "https://api.bfl.ml/v1/image",
        headers={
            "accept": "application/json",
            "x-key": BFL_API_KEY,
            "Content-Type": "application/json",
        },
        json={"prompt": prompt, "width": width, "height": heigth, "steps": steps},
    ).json()
    print(request)
    return request["id"]


def poll_result(request_id: str) -> str:
    """
    Poll the BFL API for the result of the image generation request. Returns the resulting url.
    :param request_id: The ID of the image generation request
    :return: The URL of the generated image
    :raises RuntimeError: If the status is not "Ready"
    """
    while get_result(request_id)["status"] not in [
        "Ready",
        "Error",
        "Content Moderated",
        "Request Moderated",
        "Task not found",
    ]:
        time.sleep(1)

    result = get_result(request_id)
    status = result["status"]

    if status == "Ready":
        return result["result"]["sample"]
    elif status in [
        "Error",
        "Content Moderated",
        "Request Moderated",
        "Task not found",
    ]:
        raise RuntimeError(f"Image generation failed. Status: {status}")


def get_result(request_id: str):
    result = requests.get(
        "https://api.bfl.ml/v1/get_result",
        headers={
            "accept": "application/json",
            "x-key": BFL_API_KEY,
        },
        params={
            "id": request_id,
        },
    ).json()
    return result