import os
import gradio as gr
import base64
from typing import List
from io import BytesIO
from PIL import Image
import zipfile
from openai import OpenAI

def encode_image(image_path):
    """Encodes an image file to base64."""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def generate_prompt(question, marking_scheme, student_response):
    """Generates the grading prompt for the OpenAI API."""
    prompt = f"""
    Question: {question}

    Marking Scheme: {marking_scheme}

    Student Response: {student_response}

    As an expert in this field, please grade the student's response based on the marking scheme provided. Provide detailed scores and feedback, and a well-tabulated breakdown of scores.
    """
    return prompt

def read_file_content(file):
    """Reads the content of a file."""
    with open(file.name, 'r') as f:
        return f.read()

def grade_student_answers(client, marking_scheme, student_answers):
    """Grades student answers using the OpenAI API."""
    prompt = f"""
    Marking Scheme:
    {marking_scheme}

    Student Answers:
    {student_answers}

    Grade the student answers based on the marking scheme. Use appropriate Checkmark (✓) and (X). Provide a detailed feedback and score as a percentage.

    The output should resemble that of a Professor!
    """
    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are an expert Quiz grader."},
            {"role": "user", "content": prompt}
        ],
        max_tokens=1048
    )
    return response.choices[0].message.content.strip()

def grade_explanatory_test_text(api_key, question_file, marking_scheme_file, student_responses):
    """Grades explanatory test text files."""
    client = OpenAI(api_key=api_key)
    output_files = []
    try:
        question = read_file_content(question_file)
        marking_scheme = read_file_content(marking_scheme_file)
        for student_file in student_responses:
            student_name = os.path.splitext(os.path.basename(student_file.name))[0]
            student_response = read_file_content(student_file)
            prompt = generate_prompt(question, marking_scheme, student_response)
            response = client.chat.completions.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant"},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=3500,
                temperature=0.0
            )
            grade = response.choices[0].message.content.strip()
            output_filename = f"{student_name}_grade.txt"
            with open(output_filename, 'w') as out_f:
                out_f.write(grade)
            output_files.append(output_filename)
        
        zip_filename = "graded_results.zip"
        with zipfile.ZipFile(zip_filename, 'w') as zip_file:
            for file in output_files:
                zip_file.write(file)
        
        return zip_filename
    except Exception as e:
        return f"An error occurred: {e}"

def extract_text_from_image(api_key, image_file):
    """Extracts text from an image using the OpenAI API."""
    try:
        client = OpenAI(api_key=api_key)
        
        with Image.open(image_file) as img:
            img = img.convert("RGB")
            img.thumbnail((1280, 1280))
            buffer = BytesIO()
            img.save(buffer, format="JPEG")
            base64_image = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        response = client.chat.completions.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": "You are a helpful assistant"},
                {"role": "user", "content": [
                    {"type":"text", "text": "Extract the text from this image. It is a student exam script, where the student is answering multiple choice questions. Write out the text in the image. Don't include any other text in your output."},
                    {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
                ]}
            ],
            max_tokens=1048
        )
        
        return response.choices[0].message.content.strip()
    except Exception as e:
        return f"An error occurred while extracting text from the image: {e}"

def grade_explanatory_test_image(api_key, question_file, marking_scheme_file, student_responses):
    """Grades explanatory test image files."""
    client = OpenAI(api_key=api_key)
    output_files = []
    try:
        question = read_file_content(question_file)
        marking_scheme = read_file_content(marking_scheme_file)
        for image_file in student_responses:
            student_name = os.path.splitext(os.path.basename(image_file.name))[0]
            student_response = extract_text_from_image(api_key, image_file)
            if "An error occurred" in student_response:
                return student_response
            prompt = generate_prompt(question, marking_scheme, student_response)
            response = client.chat.completions.create(
                model="gpt-4",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant"},
                    {"role": "user", "content": prompt}
                ],
                max_tokens=3500,
                temperature=0.0
            )
            grade = response.choices[0].message.content.strip()
            output_filename = f"{student_name}_grade.txt"
            with open(output_filename, 'w') as out_f:
                out_f.write(grade)
            output_files.append(output_filename)
        
        zip_filename = "graded_results.zip"
        with zipfile.ZipFile(zip_filename, 'w') as zip_file:
            for file in output_files:
                zip_file.write(file)
        
        return zip_filename
    except Exception as e:
        return f"An error occurred: {e}"

def grade_multiple_choice_test(api_key, marking_scheme_file, images):
    """Grades multiple choice test image files."""
    client = OpenAI(api_key=api_key)
    output_files = []
    try:
        marking_scheme = read_file_content(marking_scheme_file)
        for image_file in images:
            student_name = os.path.splitext(os.path.basename(image_file.name))[0]
            student_answers = extract_text_from_image(api_key, image_file)
            if "An error occurred" in student_answers:
                return student_answers
            
            grade = grade_student_answers(client, marking_scheme, student_answers)
            
            output_filename = f"{student_name}_grade.txt"
            with open(output_filename, 'w') as out_f:
                out_f.write(grade)
            output_files.append(output_filename)
        
        zip_filename = "graded_results.zip"
        with zipfile.ZipFile(zip_filename, 'w') as zip_file:
            for file in output_files:
                zip_file.write(file)
        
        return zip_filename
    except Exception as e:
        return f"An error occurred: {e}"

def handle_choice(choice):
    """Handles the user choice for test type and updates the UI accordingly."""
    if choice == "Explanatory Test":
        return [
            gr.update(visible=True),  # question_file
            gr.update(visible=True),  # marking_scheme_explanatory_file
            gr.update(visible=True),  # explanatory_type
            gr.update(visible=False), # marking_scheme_mcq_file
            gr.update(visible=False), # image_input_mcq
            gr.update(visible=False), # student_responses_text
            gr.update(visible=False)  # student_responses_image
        ]
    else:
        return [
            gr.update(visible=False), # question_file
            gr.update(visible=False), # marking_scheme_explanatory_file
            gr.update(visible=False), # explanatory_type
            gr.update(visible=True),  # marking_scheme_mcq_file
            gr.update(visible=True),  # image_input_mcq
            gr.update(visible=False), # student_responses_text
            gr.update(visible=False)  # student_responses_image
        ]

def handle_explanatory_type(explanatory_type):
    """Handles the explanatory test type and updates the UI accordingly."""
    if explanatory_type == "Text Files":
        return [
            gr.update(visible=True),  # student_responses_text
            gr.update(visible=False)  # student_responses_image
        ]
    else:
        return [
            gr.update(visible=False), # student_responses_text
            gr.update(visible=True)   # student_responses_image
        ]

def clear_inputs():
    """Clears all input fields."""
    return [
        "",  # API key
        None,  # choice
        gr.update(value=None, visible=False),  # explanatory_type
        gr.update(visible=False),  # question_file
        gr.update(visible=False),  # marking_scheme_explanatory_file
        gr.update(visible=False),  # student_responses_text
        gr.update(visible=False),  # student_responses_image
        gr.update(visible=False),  # marking_scheme_mcq_file
        gr.update(visible=False),  # image_input_mcq
        None   # output_file
    ]

# Gradio Interface
with gr.Blocks() as demo:
    api_key = gr.Textbox(label="OpenAI API Key", type="password")

    choice = gr.Radio(["Explanatory Test", "Multiple Choice Test"], label="Choose the type of test to grade")
    explanatory_type = gr.Radio(["Text Files", "Image Files"], label="Choose the type of explanatory test", visible=False)

    # Explanatory test inputs
    question_file = gr.File(label="Upload Question File", visible=False)
    marking_scheme_explanatory_file = gr.File(label="Upload Marking Scheme File", visible=False)
    student_responses_text = gr.File(label="Upload Student Response Text Files", file_count='multiple', visible=False)
    student_responses_image = gr.File(label="Upload Student Response Image Files", file_count='multiple', visible=False)

    # Multiple choice test inputs
    marking_scheme_mcq_file = gr.File(label="Upload Marking Scheme File", visible=False)
    image_input_mcq = gr.File(label="Upload Student Answer Images", file_count='multiple', visible=False)

    output_file = gr.File(label="Download Graded Results")

    # Handle choice to show/hide appropriate inputs
    choice.change(fn=handle_choice, inputs=choice, outputs=[question_file, marking_scheme_explanatory_file, explanatory_type, marking_scheme_mcq_file, image_input_mcq])
    explanatory_type.change(fn=handle_explanatory_type, inputs=explanatory_type, outputs=[student_responses_text, student_responses_image])

    # Submit and Clear buttons
    submit_btn = gr.Button("Submit")
    clear_btn = gr.Button("Clear")

    submit_btn.click(
        fn=lambda api_key, choice, explanatory_type, question_file, marking_scheme_explanatory_file, student_responses_text, student_responses_image, marking_scheme_mcq_file, image_input_mcq:
        grade_explanatory_test_text(api_key, question_file, marking_scheme_explanatory_file, student_responses_text) if explanatory_type == "Text Files" else
        grade_explanatory_test_image(api_key, question_file, marking_scheme_explanatory_file, student_responses_image) if choice == "Explanatory Test" else
        grade_multiple_choice_test(api_key, marking_scheme_mcq_file, image_input_mcq),
        inputs=[api_key, choice, explanatory_type, question_file, marking_scheme_explanatory_file, student_responses_text, student_responses_image, marking_scheme_mcq_file, image_input_mcq],
        outputs=output_file
    )

    clear_btn.click(
        fn=clear_inputs,
        inputs=[],
        outputs=[api_key, choice, explanatory_type, question_file, marking_scheme_explanatory_file,
                student_responses_text, student_responses_image, marking_scheme_mcq_file,
                image_input_mcq, output_file]
    )

if __name__ == "__main__":
    demo.launch()