import streamlit as st
from freeGPT import Client
from PIL import Image
from io import BytesIO

# Initialize the client with your API key
client = Client()

st.title("ChatBot")
st.write("Welcome to our ChatBot! Ask me anything.")

# Create a text area for user input
user_input = st.text_area("Enter your message here: ", height=100)

# Create a dropdown menu to select the mode
mode_dropdown = st.selectbox("Select Mode", ["Text", "Image"], index=0)

# Create a submit button
submit_button = st.button("Submit")

# Define a function to handle form submission
def handle_form():
    global user_input
    if user_input != "":
        # Check the selected mode
        if mode_dropdown == "Text":
            # Generate response using the freeGPT model
            response = client.create_completion("gpt3", user_input)
            
            # Print the response
            st.write(response)
        elif mode_dropdown == "Image":
            # Generate image using the freeGPT model
            img_bytes = client.create_image("prodia", user_input)
            
            # Convert bytes to image object
            img = Image.open(BytesIO(img_bytes))
            
            # Show the image
            st.image(img, width=None, height=None)
    else:
        st.write("Please enter a valid message.")

# Set up event handler for submit button click
submit_button.onclick(handle_form)

# Display the form
st.form(user_input, mode_dropdown, submit_button)