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

st.title("Chem 210 Autograder - Manual Upload")

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 file_uploader for image input
image_input1 = st.file_uploader("Upload the solution molecule", type=["png", "jpg", "jpeg"])
image_input2 = st.file_uploader("Upload the desired molecule to be graded", type=["png", "jpg", "jpeg"])

if st.button("Submit"):
    if image_input1 and image_input2:
        # Open and convert images to RGB format
        image1 = Image.open(image_input1).convert('RGB')
        image2 = Image.open(image_input2).convert('RGB')

        # Process images using your response function
        answer = response(image1, image2)

        # Display results in chat style
        st.session_state.messages.append({"role": "user", "content": "User uploaded images."})
        with st.chat_message("AI"):
            st.markdown(answer)
        st.session_state.messages.append({"role": "AI", "content": answer})
    else:
        st.warning("Please upload two images.")