import gradio as gr
from dotenv import load_dotenv
from transformers import pipeline
import urllib.request
from PIL import Image

# Load environment variables from the .env file de forma local
load_dotenv()
import base64

with open("Iso_Logotipo_Ceibal.png", "rb") as image_file:
    encoded_image = base64.b64encode(image_file.read()).decode()

import os
from openai import OpenAI


es_en_translator = pipeline("translation", model="Helsinki-NLP/opus-mt-es-en")


def translate_text(text):
    text = es_en_translator(text)[0].get("translation_text")  # type: ignore
    return text


client = OpenAI()

client.api_key = os.environ["OPENAI_API_KEY"]


def get_image(text: str, translate: bool):
    print(text)
    if translate:
        text = translate_text(text)
    response = client.images.generate(
        model="dall-e-3",
        prompt=text,
        size="1024x1024",
        quality="standard",
        n=1,
    )  # ,response_format="b64_json"
    print(response)

    urllib.request.urlretrieve(response.data[0].url, "img.png")  # type: ignore

    img = Image.open("img.png")

    return img, response.data[0].revised_prompt

    # return response["data"][0]["url"]  # type: ignore


with gr.Blocks() as demo:
    gr.Markdown(
        """
    <center>
    <h1>
    Uso de AI para la generación de imagenes a partir de texto.
    </h1>
    <img src='data:image/jpg;base64,{}' width=200px>
    <h2>
    Con este espacio podrás generar imagenes a partir de texto. Utiliza el modelo DALL-E 3 de OpenAI.
    </h2>
    <h3>
     Tenés que tener en cuenta que el modelo ahora toma el mensaje predeterminado proporcionado y lo reescribe automáticamente por razones de seguridad y para agregar más detalles.
    </h3>
    
    </center>
    """.format(
            encoded_image
        )
    )
    with gr.Row():
        with gr.Column():
            with gr.Row():
                gr.Markdown(
                    "El modelo funciona mejor con entradas en inglés, podés elegir traducir al inglés con este check:"
                )
                check = gr.Checkbox(label="Traducir al inglés")
            with gr.Row():
                gr.Markdown("Primero debes ingresar el texto para generar la imagen:")

            with gr.Row():
                with gr.Column(scale=4):
                    prompt = gr.Textbox(
                        label="Texo base para generar la imagen"
                    )  # Give prompt some real estate
                with gr.Column(scale=1, min_width=50):
                    btn = gr.Button("Generar")  # Submit button side by side!

        with gr.Column():
            text_output = gr.TextArea(label="Prompt revisado por el modelo:")
            output = gr.Image(
                label="Resultado", height=1024, width=1024
            )  # Move the output up too
            # examples = gr.Examples(
            #     inputs=[prompt]
            #     examples=[["Un perro en el parque", "low quality"]],
            # )

    btn.click(
        fn=get_image,
        inputs=[prompt, check],
        outputs=[output, text_output],
    )  # steps,guidance,width,height]
demo.queue()
demo.launch(debug=True)