import streamlit as st
from src.main import process_essays
from PIL import Image
import os
from openpyxl import load_workbook, Workbook
from io import BytesIO
import openai

def save_workbook_to_bytes(wb):
    # Save the workbook into a BytesIO object (in memory, not on disk)
    byte_io = BytesIO()
    wb.save(byte_io)
    byte_io.seek(0)  # Go to the beginning of the BytesIO buffer
    return byte_io.getvalue()


openai_api_key = os.getenv("OPENAI_API_KEY")

# Set the OpenAI API key
openai.api_key = openai_api_key

st.title("AutoAssess: Student Essay Transcription and Assessment")

st.title("AutoAssess")
st.write("If you see this, the basic app is loading correctly!")

# Upload folder of images
uploaded_files = st.file_uploader("Upload a folder of student essays (images)", type=['jpg', 'jpeg', 'png'], accept_multiple_files=True)

# replace uploaded files with files loading from directory
# image_dir = "data/images"
# uploaded_files = []
# for file in os.listdir(image_dir):
#     with open(image_dir + '/' + file, "rb") as image_file:
#         uploaded_files.append(image_file.read())

# Text inputs for question and criteria
# essay_question = st.text_input("Enter the essay question:")
# grading_criteria = st.text_area("Enter grading criteria or relevant marking information:")

essay_question = "What is beauty?"
grading_criteria = "1. Introduction\n2. Body\n3. Conclusion\n4. Grammar\n5. Spelling\n6. Punctuation\n7. Originality\n8. Creativity"

# Upload Excel file with student IDs and page count
excel_file = st.file_uploader("Upload Excel file with student IDs and page count", type=["xlsx"])
# excel_file = "data/essays.xlsx"

if st.button("Process Essays"):
    if not uploaded_files or not essay_question or not grading_criteria or not excel_file:
        st.warning("Please upload all required files and enter necessary information.")
    else:
        uploaded_files = sorted(uploaded_files, key=lambda x: x.name)
        # Process student info file
        workbook = load_workbook(excel_file)


        new_workbook = process_essays(uploaded_files,essay_question,grading_criteria,workbook)


        # Optional: Save results to the output folder
        # output_file = "output/results.xlsx"
        # new_workbook.save(output_file)
        # st.success(f"All essays processed. Results saved to {output_file}")

        # Convert the workbook to bytes
        excel_file = save_workbook_to_bytes(new_workbook)

        # Display the download button
        st.download_button(
            label="Download the Excel file",
            data=excel_file,
            file_name="results.xlsx",
            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
)