import streamlit as st
from ocr_processing import OCRProcessor
from grading_logic import GradingSystem
import os

# Initialize components
ocr_handler = OCRProcessor()
grading_handler = GradingSystem(ocr_handler)


def main():
    st.set_page_config(page_title="Grading Assistant", page_icon=":pencil:", layout="wide")

    st.sidebar.title("Upload Student Files")
    st.sidebar.info("Upload student answer sheets and marking schemes.")

    student_files = st.sidebar.file_uploader("Upload Answer Sheets", type=["pdf", "jpeg", "jpg", "png"], accept_multiple_files=True)
    marking_scheme_file = st.sidebar.file_uploader("Upload Marking Scheme (DOCX)", type=["docx"])

    st.title("📖 Automated Grading Assistant")
    st.subheader("An AI-powered tool for grading student responses.")

    if st.button("Start Grading"):
        if not student_files or not marking_scheme_file:
            st.error("Both student answers and a marking scheme are required!")
            return

        student_image_paths = []
        for student_file in student_files:
            file_path = f"./uploads/{student_file.name}"
            with open(file_path, "wb") as f:
                f.write(student_file.getvalue())

            if student_file.name.lower().endswith(".pdf"):
                try:
                    extracted_images = ocr_handler.convert_pdf_to_images(file_path)
                    student_image_paths.extend(extracted_images)
                except Exception as err:
                    st.error(f"Error processing {student_file.name}: {err}")
            else:
                student_image_paths.append(file_path)

        # Save marking scheme
        marking_scheme_path = f"./uploads/{marking_scheme_file.name}"
        with open(marking_scheme_path, "wb") as f:
            f.write(marking_scheme_file.getvalue())

        st.info("Extracting student responses...")
        extracted_answers = grading_handler.extract_answers(student_image_paths)
        st.success("Responses extracted successfully.")
        st.text_area("Extracted Responses", extracted_answers, height=200)

        st.info("Extracting marking scheme...")
        marking_criteria = grading_handler.extract_marking_criteria(marking_scheme_path)
        st.success("Marking scheme extracted.")
        st.text_area("Marking Scheme", marking_criteria, height=200)

        st.info("Evaluating responses...")
        assessment_results = grading_handler.grade_answers(extracted_answers, marking_criteria)
        st.success("Assessment complete.")
        st.text_area("Assessment Results", assessment_results, height=200)


if __name__ == "__main__":
    main()