import gradio as gr
from main import grade_submission
import pdfplumber
import docx


def extract_text(file):
    if file is None:
        return ""

    name = file.name.lower()

    if name.endswith(".pdf"):
        text = ""
        with pdfplumber.open(file) as pdf:
            for page in pdf.pages:
                text += page.extract_text() + "\n"
        return text

    if name.endswith(".docx"):
        doc = docx.Document(file)
        return "\n".join(p.text for p in doc.paragraphs)

    return file.read().decode("utf-8", errors="ignore")


def resolve_input(file, text):
    return extract_text(file) if file else text


def grade_ui(
    model,
    q_file, q_text,
    r_file, r_text,
    s_file, s_text,
    instruction
):
    question_paper = resolve_input(q_file, q_text)
    rubric = resolve_input(r_file, r_text)
    student_answer = resolve_input(s_file, s_text)

    return grade_submission(
        model,
        question_paper,
        rubric,
        student_answer,
        instruction
    )


with gr.Blocks() as demo:
    gr.Markdown("## AutoGrader")

    model = gr.Dropdown(
        ["Phi-3-mini", "Mistral-7B-Instruct"],
        value="Phi-3-mini",
        label="Model"
    )

    with gr.Tab("Question Paper"):
        q_file = gr.File(
            label="Upload Question Paper (PDF/DOCX/TXT)",
            file_types=[".pdf", ".docx", ".txt"]
        )
        q_text = gr.Textbox(
            label="Or paste Question Paper",
            lines=5
        )

    with gr.Tab("Rubric"):
        r_file = gr.File(
            label="Upload Rubric (PDF/DOCX/TXT)",
            file_types=[".pdf", ".docx", ".txt"]
        )
        r_text = gr.Textbox(
            label="Or paste Rubric",
            lines=5
        )

    with gr.Tab("Student Submission"):
        s_file = gr.File(
            label="Upload Student Submission (PDF/DOCX/TXT)",
            file_types=[".pdf", ".docx", ".txt"]
        )
        s_text = gr.Textbox(
            label="Or paste Student Submission",
            lines=6
        )

    instruction = gr.Textbox(
        label="Grading Instruction",
        placeholder="e.g. Grade only Q2 out of 20"
    )

    output = gr.Textbox(
        label="Grading Output (JSON)",
        lines=12
    )

    gr.Button("Grade").click(
        grade_ui,
        inputs=[
            model,
            q_file, q_text,
            r_file, r_text,
            s_file, s_text,
            instruction
        ],
        outputs=output
    )

demo.launch()