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

# GCP credentials stuff

import json
import pybase64
from google.oauth2 import service_account 
import google.auth

import json
import os
from dotenv import load_dotenv

load_dotenv( )

# Load the decrypted JSON file or its content as environment variable
decrypted_service_account_path = "decrypted.json"

# Get the secret from environment variables
decrypter_json = os.getenv("DECRYPTER_JSON")

# with open(decrypted_service_account_path, 'r') as f:
#     service_account_info = json.load(f)
service_account_info = json.loads(decrypter_json)

# 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="elemental-shine-437217-t6"
aiplatform.init(project=project, credentials=credentials)

def generate_image(prompt, model_name):
    try:
        print(f"Generating image with prompt: {prompt}, model: {model_name}")
        model = ImageGenerationModel.from_pretrained(model_name)
        
        response = model.generate_images(
            prompt=prompt,
            number_of_images=1
        )
        
        print("Received response:", response)
        
        image_bytes = response[0]._image_bytes
        image_url = Image.open(BytesIO(image_bytes))
        
    except Exception as e:
        print(f"Error: {e}")
        raise gr.Error(f"An error occurred while generating the image: {str(e)}")
    
    return image_url

with gr.Blocks() as demo:

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

    #instructtons
    with gr.Accordion("Instructions & Tips",label="instructtons", open=False):
        with gr.Row():
            gr.Markdown("***Tips**; Use adjectives (size, color,mood), spectty the visual style (realistic, cartoon, 8-btt), explain the point of vi" )


    #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=["imagen-3.0-fast-generate-001", "imagen-3.0-generate-001"], label="Model", value="imagen-3.0-generate-001")

    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")

demo.launch(share=False)