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
import openai

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


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


def get_image(text: str, translate: bool):
    print(text)
    if translate:
        text = translate_text(text)
    response = openai.Image.create(
        prompt=text, n=1, size="512x512", model="dall-e-2"
    )  # ,response_format="b64_json"
    print(response)

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

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

    return img

    # 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 de OpenAI.
    </h2>
    <h2>
     Obtendrás mejores resultados cuanto más larga sea la descripción.
    </h2>
    
    </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():
            output = gr.Image(
                label="Resultado", height=512, width=512
            )  # 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],
    )  # steps,guidance,width,height]
demo.queue(concurrency_count=4)
demo.launch(debug=True)