import streamlit as st
from PIL import Image
from utils import get_gpt4V_response, get_str_to_json


st.set_page_config(page_title="GPT-4V Demo", page_icon="🧠", layout="wide")

with st.sidebar:
    st.title("Parameters")
    st.write("This is a demo of GPT-4V model. It takes a story, goal, entity and an image as input and generates a response.")

    st.subheader("Sampling Temperature")
    temperature = st.slider(label="", min_value=0.1, max_value=1.0, value=0.5, step=0.1)
    st.write("The sampling temperature, between 0 and 1. Higher values like 0.8 will make the output more random, while lower values like 0.2 will make it more focused and deterministic.")

    st.subheader("Entity?")
    entity_opt = st.radio(label="With or Without", options=[1, 0], format_func=lambda x: ["Without", "With"][x])


def main():
    global temperature, entity
    st.title('What can go wrong?')

    col1, col2 = st.columns(2)

    with col1:
        story = st.text_area("Story", placeholder="Enter the story here")

        entity = None
        if entity_opt:
            entity = st.text_input("Entity", placeholder="Enter the entity here")

        goal = st.text_area("Goal", placeholder="Enter the goal here")

        image = st.file_uploader("Upload Image", type=['jpg', 'png'])

        if image:
            cols = st.columns(3)
            with cols[1]:
                st.image(image, caption="Uploaded Image", use_column_width=True)
        

        if st.button("Submit"):
            if not story or not goal or (entity_opt and not entity) or not image:
                st.error("Please fill all the fields")
                return
            with col2:
                with st.status("Generating response...", expanded=True):
                    response = get_gpt4V_response(story, goal, entity, image, temperature=temperature)
                
                try:
                    response_json = get_str_to_json(response)
                    if "condition" not in response_json or "alternate_condition" not in response_json:
                        raise ValueError("Invalid JSON - 1")
                    if not entity_opt and "entity" not in response_json:
                        raise ValueError("Invalid JSON - 2")
                    
                    if not entity_opt:
                        with st.expander("Entity", expanded=True):
                            st.write(response_json["entity"])
                    
                    with st.expander("Condition", expanded=True):
                        st.write(response_json["condition"])
                    
                    with st.expander("Alternate Condition", expanded=True):
                        st.write(response_json["alternate_condition"])

                except Exception as e:
                    print(e)
                    st.warning(f"Failed to parse JSON. Going for full output")
                    st.write(response)
main()