import streamlit as st
from generator2 import response
from retrieve import retrieve_molecule_index
from PIL import Image

st.title("Chem 210 Autograder")

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Use text_area for text input with a smaller height
text_input = st.text_input("Indicate the molecule you want graded")

# Use file_uploader for image input
image_input = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])

if st.button("Submit"):
    if text_input and image_input:
        if text_input:
            with st.chat_message("user"):
                st.markdown(f"Does the following image indicate correct chemical structure of {text_input}?")
            st.session_state.messages.append({"role": "user", "content": text_input})
            index = int(retrieve_molecule_index(text_input)[0])
            image_path = f"test/CID_{index}.png"
            image2 = Image.open(image_path).convert('RGB')
            
        if image_input:
            with st.chat_message("user"):
                st.image(image_input, caption="User image", use_column_width=True)
            st.session_state.messages.append({"role": "user", "content": "User uploaded an image."})
            image = Image.open(image_input).convert('RGB')
        
        answer = response(image2, image)
        with st.chat_message("AI"):
            st.markdown(answer)
        
        st.session_state.messages.append({"role": "AI", "content": answer})