import streamlit as st
from clarifai_grpc.channel.clarifai_channel import ClarifaiChannel
from clarifai_grpc.grpc.api import resources_pb2, service_pb2, service_pb2_grpc
from clarifai_grpc.grpc.api.status import status_code_pb2
from PIL import Image
from io import BytesIO
from base64 import b64decode

# Streamlit page configuration
st.set_page_config(page_title='DALL-E Image Generator', layout='wide')

# Streamlit sidebar elements
st.sidebar.title("DALL-E Image Generator")
prompt = st.sidebar.text_input("Enter your prompt:", value='A penguin watching the sunset.')
PAT = st.sidebar.text_input("Enter your Personal Access Token:", type="password")

# Constants
USER_ID = 'openai'
APP_ID = 'dall-e'
MODEL_ID = 'dall-e-3'
MODEL_VERSION_ID = 'dc9dcb6ee67543cebc0b9a025861b868'

# Streamlit main page
st.title('Generate Images with DALL-E and Clarifai')
if st.sidebar.button("Generate Image"):
    if PAT:  # Check if PAT is filled
        with st.spinner('Generating Image...'):
            # Clarifai gRPC setup
            channel = ClarifaiChannel.get_grpc_channel()
            stub = service_pb2_grpc.V2Stub(channel)
            metadata = (('authorization', 'Key ' + PAT),)
            userDataObject = resources_pb2.UserAppIDSet(user_id=USER_ID, app_id=APP_ID)

            post_model_outputs_response = stub.PostModelOutputs(
                service_pb2.PostModelOutputsRequest(
                    user_app_id=userDataObject,
                    model_id=MODEL_ID,
                    version_id=MODEL_VERSION_ID,
                    inputs=[
                        resources_pb2.Input(
                            data=resources_pb2.Data(
                                text=resources_pb2.Text(raw=prompt)
                            )
                        )
                    ]
                ),
                metadata=metadata
            )

            if post_model_outputs_response.status.code != status_code_pb2.SUCCESS:
                st.error("Error in generating image: " + post_model_outputs_response.status.description)
            else:
                # Decode the base64 image and display
                output = post_model_outputs_response.outputs[0].data.image.base64
                image = Image.open(BytesIO(output))
                st.image(image, caption='Generated Image', use_column_width=True)

    else:
        st.sidebar.warning("Please enter your Personal Access Token.")