# file stuff
import os
from io import BytesIO

#image generation stuff
from PIL import Image

# gradio / hf / image gen stuff
import gradio as gr
from dotenv import load_dotenv


from google.cloud import aiplatform
import vertexai
from vertexai.preview.vision_models import ImageGenerationModel
from vertexai import preview
import uuid #for generating unique filenames

# GCP credentials stuff
import json
import tempfile
import pybase64
from google.oauth2 import service_account
import google.auth

load_dotenv()

service_account_json = pybase64.b64decode(os.getenv("IMAGEN"))
service_account_info = json.loads(service_account_json)
credentials = service_account.Credentials.from_service_account_info(service_account_info)
project="pdr-imagen"
aiplatform.init(project=project, credentials=credentials)

def generate_image(prompt,model_name):
    try:
        model = ImageGenerationModel.from_pretrained(model_name)
        response = model.generate_images(
            prompt=prompt,
            number_of_images=1,
        )

        image_bytes = response[0]._image_bytes
        image_url = Image.open(BytesIO(image_bytes))

    except Exception as e:
        print(e)
        raise gr.Error(f"An error occurred while generating the image for: {entry}")
    return image_url

with gr.Blocks() as demo:

    gr.Markdown("# <center>Google Vertex Imagen Generator</center>")

    #instructions
    with gr.Accordion("Instructions & Tips",label="instructions",open=False):
        with gr.Row():
            gr.Markdown("**Tips**: Use adjectives (size,color,mood), specify the visual style (realistic,cartoon,8-bit), explain the point of view (from above,first person,wide angle) ")

    #prompts
    with gr.Accordion("Prompt",label="Prompt",open=True):
        text = gr.Textbox(label="What do you want to create?", placeholder="Enter your text and then click on the \"Image Generate\" button")

    model = gr.Dropdown(choices=["imagegeneration@002", "imagegeneration@005"], label="Model", value="imagegeneration@005")

    with gr.Row():
        btn = gr.Button("Generate Images")

    #output
    with gr.Accordion("Image Output",label="Image Output",open=True):
        output_image = gr.Image(label="Image")

    btn.click(fn=generate_image, inputs=[text, model ], outputs=output_image, api_name=False)
    text.submit(fn=generate_image, inputs=[text, model ], outputs=output_image, api_name="generate_image") # Generate an api endpoint in Gradio / HF

demo.launch(share=False)