import os
import gradio as gr
import google.generativeai as genai

# Configure Gemini API (key must be set in Hugging Face Space secrets)
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

# ---------- PROMPTS ----------
TRANSCRIPTION_PROMPT = """
Persona:
You are an expert transcriptionist specializing in scientific and mathematical documents. 
Your primary goal is to convert handwritten mathematical work into a perfectly formatted, 
machine-readable Markdown document using LaTeX for all mathematical notation.

Rules:
- Transcribe exactly what is written, do not correct errors.
- Use $...$ for inline math, $$...$$ for block math.
- Ignore struck-through text.
- Preserve structure: bold for Q numbers (**1.**), step-by-step math with \\begin{align*}.
- If a symbol is ambiguous, mark as [x?].
Output must be a clean Markdown string.
"""

GRADING_PROMPT = """
You are an official examiner. Grade the student transcription using the question paper 
and the official marking scheme.

Rules:
1. Apply marks exactly as per the markscheme (M1, A1, etc.).
2. M marks must be earned before A marks.
3. Justify each awarded or withheld mark with clear reasoning.
4. Classify all errors as Conceptual Error, Silly Mistake, or None.
5. Follow dependency between M and A strictly.
6. Do not give marks outside the markscheme.
Output must be a structured grading report with reasoning.
"""

# ---------- STEP 1: TRANSCRIPTION ----------
def transcribe(ans_file):
    try:
        ans_uploaded = genai.upload_file(path=ans_file.name, display_name="Answer Sheet")

        model = genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
        resp = model.generate_content([TRANSCRIPTION_PROMPT, ans_uploaded])

        transcription = getattr(resp, "text", None) or resp.candidates[0].content.parts[0].text
        return transcription
    except Exception as e:
        return f"❌ Error during transcription: {e}"

# ---------- STEP 2: GRADING ----------
def grade(qp_file, ms_file, transcription):
    try:
        qp_uploaded = genai.upload_file(path=qp_file.name, display_name="Question Paper")
        ms_uploaded = genai.upload_file(path=ms_file.name, display_name="Marking Scheme")

        model = genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
        resp = model.generate_content([GRADING_PROMPT, qp_uploaded, ms_uploaded, transcription])

        grading = getattr(resp, "text", None) or resp.candidates[0].content.parts[0].text
        return grading
    except Exception as e:
        return f"❌ Error during grading: {e}"

# ---------- GRADIO APP ----------
with gr.Blocks(title="📘 AI Teacher Assistant") as demo:
    gr.Markdown("## 📘 AI Teacher Assistant\nUpload exam documents to transcribe and grade student answers step by step.")

    with gr.Row():
        qp_file = gr.File(label="Upload Question Paper (PDF)", type="filepath")
        ms_file = gr.File(label="Upload Mark Scheme (PDF)", type="filepath")
        ans_file = gr.File(label="Upload Student Answer Sheet (PDF)", type="filepath")

    # Step 1: Transcription
    transcribe_btn = gr.Button("Step 1: Transcribe Answer Sheet")
    transcription_out = gr.Markdown(label="📄 Student Transcription")

    # Step 2: Grading
    grade_btn = gr.Button("Step 2: Grade the Student")
    grading_out = gr.Textbox(label="✅ Grading Report (Step-by-Step)", lines=20)

    # Button Logic
    transcribe_btn.click(fn=transcribe, inputs=[ans_file], outputs=[transcription_out])
    grade_btn.click(fn=grade, inputs=[qp_file, ms_file, transcription_out], outputs=[grading_out])

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