import os
import gradio as gr
import google.generativeai as genai
from markdown_pdf import MarkdownPdf, Section
import subprocess

# ---------- PROMPTS ----------
PROMPTS = {
    "ALIGNMENT_PROMPT": {
        "role": "system",
        "content": """Developer: Role: Expert examiner and transcription specialist. 
Your objective is to align three sources per question/sub-question:
- Question Paper (QP)
- Markscheme (MS)
- Student Answer Sheet (AS)

## Instructions
1. Carefully parse all documents and align content per question and sub-question.
2. For each question/sub-question, create a structured Markdown block as follows:
---
## Question X [and sub-question if applicable, e.g., ### (b)(ii)]
*QP:* [Exact question text or [Not found]]
*MS:* [Relevant markscheme section or [Not found]]
*AS:* [Final cleaned student answer; use fenced code for mathematics; insert [illegible] or [No response] as required]
---
3. Formatting requirements:
- Use '##' for main questions, '###' for sub-questions.
- Maintain section order: QP | MS | AS (always in that sequence).
- Enclose all mathematical expressions in Markdown fenced code blocks (``` triple backticks).
- If a diagram/graph is omitted, write [Graph omitted] in its place.
- For unreadable portions of the student's answer, insert [illegible]; if the answer is wholly unreadable, set AS to [illegible].
- If a question is skipped or unanswered, AS must be exactly [No response].
- Keep MS annotations (e.g., M1, A1, R1) verbatim.
- Diagrams/graphs are not to be recreated.
- If any QP, MS, or AS content is missing, specify [Not found] for that section.
- Ensure consistency and determinism in formatting so subsequent models can grade directly from this aligned format.
- List all main questions and sub-questions in their original order, clearly denoting sub-questions (e.g., '### (b)(i)', '### (b)(ii)').

After each alignment action, briefly validate that the content for QP, MS, and AS matches expectations and alignments are correct. If validation fails, self-correct or flag the issue.

## Example
---
## Question 1
*QP:* Expand (1+x)^3
*MS:* M1 for binomial expansion, A1 for coefficients, A1 for final form
*AS:* x^3 + 3x^2 + 3x + 1
---
"""
    },
    "GRADING_PROMPT": {
        "role": "system",
        "content": """Developer: You are an official examiner. Apply the following grading rules precisely.

## Grading Checklist
- Assess each question part against the provided markscheme.
- Award marks for correct methods (M), accurate answers (A), and clear reasoning (R).
- Use Follow Through (FT) for correctly applied subsequent working using a previous error.
- Always state BOTH:
  1. What was wrong (the error).
  2. What is right (the correct method/answer from markscheme).
- Summarize total marks and classify error types.
- End with an Examiner’s Report table.

### Abbreviations:
- **M**: Method
- **A**: Accuracy/Answer
- **R**: Reasoning
- **AG**: Answer given (no marks)
- **FT**: Follow Through

---
## Grading Instructions
1. Award marks using official annotations (M1, A1, etc.).
2. A marks generally require valid M marks.
3. Allow FT unless result is nonsensical.
4. Accept valid alternative forms.
5. Apply accuracy requirements (default 3 s.f. if not stated).
6. Ignore crossed-out work unless requested otherwise.
7. Mark only the first full solution unless otherwise indicated.
8. Assume graphs/diagrams are correct if required.

---
## Output Format
Produce a GitHub-flavored Markdown table:

| Student wrote | Marks Awarded | Reason |
|---------------|---------------|--------|

Rules:
- Each row matches a markable step.
- For blanks, write “(no answer)” and indicate lost mark(s).
- Lost marks: wrap in red with `<span style="color:red">A0</span>` (or M0, R0) and make Reason column red. Always also show the correct method/answer.
- Awarded marks remain plain text.
- For partial awards (M1A0A1), highlight only lost marks.
- After each question, show total in square brackets: `[2/4]`.

---
### Examiner’s Report
At the very end, provide a summary table:

Codes:
- A : All Good
- B : Silly Mistake
- C : Conceptual Error
- D : Hard Question
- E : Not Applicable

| Question Number | Marks | Remark |
|-----------------|-------|--------|
| 1               | 6/9   | C      |
| 2               | 7/7   | A      |
| 3               | 8/14  | D      |
| …               | …     | …      |

Then show total clearly:  
`Total: 40/61`

Optionally, if reasons are available, extend with:

| Question Number | Marks | Remark | Reason |
|-----------------|-------|--------|--------|

⚠️ Do NOT add any "Validation" or meta commentary. End the output after Examiner’s Report.
"""
    }
}

# -------------------- CONFIG --------------------
genai.configure(api_key=os.getenv("GEMINI_API_KEY"))

# ---------- HELPER: Save to PDF ----------
def save_as_pdf(text, filename="output.pdf"):
    pdf = MarkdownPdf()
    pdf.add_section(Section(text, toc=False))
    pdf.save(filename)
    return filename

# ---------- HELPER: Compress PDF ----------
def compress_pdf(input_path, output_path=None, max_size=20*1024*1024):
    if output_path is None:
        base, ext = os.path.splitext(input_path)
        output_path = f"{base}_compressed{ext}"

    if os.path.getsize(input_path) <= max_size:
        return input_path

    try:
        gs_cmd = [
            "gs", "-sDEVICE=pdfwrite",
            "-dCompatibilityLevel=1.4",
            "-dPDFSETTINGS=/ebook",
            "-dNOPAUSE", "-dQUIET", "-dBATCH",
            f"-sOutputFile={output_path}", input_path
        ]
        subprocess.run(gs_cmd, check=True)
        if os.path.getsize(output_path) <= max_size:
            print(f"✅ Compressed {input_path} → {output_path}")
            return output_path
        else:
            print(f"⚠️ Compression failed to reduce below {max_size/1024/1024} MB")
            return input_path
    except Exception as e:
        print(f"⚠️ Compression error: {e}")
        return input_path

# ---------- HELPER: Create Model with Fallback ----------
def create_model():
    try:
        print("⚡ Using gemini-2.5-pro model")
        return genai.GenerativeModel("gemini-2.5-pro", generation_config={"temperature": 0})
    except Exception:
        print("⚡ Falling back to gemini-2.5-flash model")
        return genai.GenerativeModel("gemini-2.5-flash", generation_config={"temperature": 0})

# ---------- PIPELINE: ALIGN + GRADE ----------
def align_and_grade(qp_file, ms_file, ans_file):
    try:
        qp_file = compress_pdf(qp_file, "qp_compressed.pdf")
        ms_file = compress_pdf(ms_file, "ms_compressed.pdf")
        ans_file = compress_pdf(ans_file, "ans_compressed.pdf")

        qp_uploaded = genai.upload_file(path=qp_file, display_name="Question Paper")
        ms_uploaded = genai.upload_file(path=ms_file, display_name="Markscheme")
        ans_uploaded = genai.upload_file(path=ans_file, display_name="Answer Sheet")

        model = create_model()

        resp = model.generate_content([
            PROMPTS["ALIGNMENT_PROMPT"]["content"],
            qp_uploaded,
            ms_uploaded,
            ans_uploaded
        ])
        aligned_text = getattr(resp, "text", None)
        if not aligned_text and resp.candidates:
            aligned_text = resp.candidates[0].content.parts[0].text
        aligned_pdf_path = save_as_pdf(aligned_text, "aligned_qp_ms_as.pdf")

        response = model.generate_content([
            PROMPTS["GRADING_PROMPT"]["content"],
            aligned_text
        ])
        grading = getattr(response, "text", None)
        if not grading and response.candidates:
            grading = response.candidates[0].content.parts[0].text

        base_name = os.path.splitext(os.path.basename(ans_file))[0]
        grading_pdf_path = save_as_pdf(grading, f"{base_name}_graded.pdf")

        return aligned_text, aligned_pdf_path, grading, grading_pdf_path

    except Exception as e:
        return f"❌ Error: {e}", None, None, None

# ---------- GRADIO APP ----------
with gr.Blocks(title="LeadIB AI Grading (Alignment + Auto-Grading)") as demo:
    gr.Markdown("## LeadIB AI Grading\nUpload Question Paper, Markscheme, and Student Answer Sheet.\nThe system will align and grade automatically.")

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

    run_btn = gr.Button("Start Alignment + Auto-Grading")

    with gr.Row():
        aligned_out = gr.Textbox(label="📄 Aligned QP | MS | AS", lines=20)
        aligned_pdf = gr.File(label="⬇️ Download Aligned (PDF)")

    with gr.Row():
        grading_out = gr.Textbox(label="✅ Grading Report", lines=20)
        grading_pdf = gr.File(label="⬇️ Download Grading Report (PDF)")

    run_btn.click(
        fn=align_and_grade,
        inputs=[qp_file, ms_file, ans_file],
        outputs=[aligned_out, aligned_pdf, grading_out, grading_pdf],
        show_progress=True
    )

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