import gradio as gr
import boto3
import banana_dev as client
import os
from typing import Iterable
import gradio as gr
from gradio.themes.base import Base
from gradio.themes.utils import colors, fonts, sizes
from PIL import Image
import io
import base64
from io import BytesIO
import traceback
import uuid
from dotenv import load_dotenv
import gzip

load_dotenv()

BANANA_API_KEY = os.getenv("BANANA_API_KEY")
BANANA_MODEL_KEY = os.getenv("BANANA_MODEL_KEY")
BANANA_URL = os.getenv("BANANA_URL")

def predict(image, prompt,qr_code, negative_prompt, seed):
    try:
        my_model = client.Client(
            api_key=BANANA_API_KEY,
            model_key=BANANA_MODEL_KEY,
            url=BANANA_URL,
        )
        #target image
        with open(image.name, "rb") as f:
            image_bytes = f.read()
        image_encoded = base64.b64encode(image_bytes)
        image = image_encoded.decode("utf-8")

        print(image)
        print(negative_prompt)
        print(seed)
        print(qr_code)
        print(prompt)

        inputs = {
            "prompt" : prompt,
            "qr_code_content" : qr_code,
            "negative_prompt": negative_prompt,
            "image": image,
            "seed": seed
        }
        result, meta = my_model.call("/", inputs)
        output_bytes = result["outputs"]

        image_encoded = output_bytes.encode('utf-8')
        image_data = base64.b64decode(image_encoded)
        image_data = gzip.decompress(image_data)
        image_io = BytesIO(image_data)
        output_image = Image.open(image_io)
        return output_image, None
    except Exception as e:
        tb_str = traceback.format_exception(type(e), e, e.__traceback__)
        tb_str = "".join(tb_str)
        return None, tb_str

with gr.Blocks() as demo:                         
    title_with_logo = gr.Markdown("# Powered by Banana <img src='https://files.umso.co/lib_udQIfMXlGlflDxFd/rnqmq7j8oqug1kae.png' align='left' width='50' height='50'>")
    with gr.Column():                        
        image = gr.File(label="Image", file_count="single", scale=1)
        prompt=gr.Textbox(label="Prompt")
        qr_code=gr.Textbox(label="QR code link")
        negative_prompt= gr.Textbox(label="Negative Prompt")
        seed= gr.Textbox(label="Seed")
        btn = gr.Button("Run illusion-diffusion-hq")                       
        output_image = gr.Image(label="Output Image")  
        output_error = gr.Textbox(label="Error Message")  
    btn.click(predict, [image, prompt,qr_code, negative_prompt, seed], [output_image, output_error])

demo.launch()