import gradio as gr
import PyPDF2
import pytesseract
from PIL import Image
import io
import faiss
import numpy as np
from transformers import AutoTokenizer, AutoModel
from docx import Document
from docx.shared import Inches
import torch
import os
from datetime import datetime

# Initialize models
MODEL_NAME = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME)

# Helper functions
def extract_text_from_pdf(pdf_file):
    try:
        pdf_reader = PyPDF2.PdfReader(pdf_file)
        text = ""
        for page in pdf_reader.pages:
            text += page.extract_text() or ""
        if not text.strip():
            text = extract_text_with_ocr(pdf_file)
        return text
    except Exception as e:
        return f"Error extracting text: {str(e)} / متن نکالنے میں خرابی: {str(e)}"

def extract_text_with_ocr(pdf_file):
    try:
        pdf_reader = PyPDF2.PdfReader(pdf_file)
        text = ""
        for page in pdf_reader.pages:
            img = page.images[0].image if page.images else None
            if img:
                text += pytesseract.image_to_string(Image.open(io.BytesIO(img.data)), lang="eng+urd")
        return text
    except Exception as e:
        return f"OCR failed: {str(e)} / OCR ناکام: {str(e)}"

def chunk_text(text, chunk_size=400, overlap=80):
    words = text.split()
    chunks = []
    for i in range(0, len(words), chunk_size - overlap):
        chunk = " ".join(words[i:i + chunk_size])
        chunks.append(chunk)
    return chunks

def create_embeddings(chunks):
    embeddings = []
    for chunk in chunks:
        inputs = tokenizer(chunk, return_tensors="pt", truncation=True, padding=True)
        with torch.no_grad():
            embedding = model(**inputs).last_hidden_state.mean(dim=1).numpy()
        embeddings.append(embedding)
    return np.vstack(embeddings)

def setup_faiss_index(embeddings):
    dimension = embeddings.shape[1]
    index = faiss.IndexFlatL2(dimension)
    index.add(embeddings)
    return index

def retrieve_relevant_chunks(query, index, chunks, k=3):
    query_inputs = tokenizer(query, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        query_embedding = model(**query_inputs).last_hidden_state.mean(dim=1).numpy()
    distances, indices = index.search(query_embedding, k)
    return [chunks[i] for i in indices[0]]

def generate_lesson_plan_boppps(grade, subject, topic, slo, duration, context):
    doc = Document()
    doc.add_heading(f"Grade {grade} {subject} Lesson Plan: {topic}", 0)
    doc.add_paragraph(f"SLO: {slo}")
    doc.add_paragraph(f"Duration: {duration} minutes")
    doc.add_heading("BOPPPS Model", level=1)
    doc.add_heading("Bridge-in", level=2)
    doc.add_paragraph(f"Engaging activity for {topic}: [Generated activity based on {context}]")
    doc.add_heading("Outcome", level=2)
    doc.add_paragraph(f"Objective: {slo}")
    doc.add_heading("Pre-assessment", level=2)
    doc.add_paragraph("Quick quiz or question to gauge prior knowledge.")
    doc.add_heading("Participatory Learning", level=2)
    doc.add_paragraph(f"Interactive activity: [Generated from {context}]")
    doc.add_heading("Post-assessment", level=2)
    doc.add_paragraph("Evaluate SLO achievement with a short task.")
    doc.add_heading("Summary", level=2)
    doc.add_paragraph("Recap key points of the lesson.")
    buffer = io.BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    return buffer, "\n".join([p.text for p in doc.paragraphs])

def generate_lesson_plan_backward(grade, subject, topic, slo, duration, context):
    doc = Document()
    doc.add_heading(f"Grade {grade} {subject} Lesson Plan: {topic}", 0)
    doc.add_paragraph(f"SLO: {slo}")
    doc.add_paragraph(f"Duration: {duration} minutes")
    doc.add_heading("Backward Design", level=1)
    doc.add_heading("Desired Results", level=2)
    doc.add_paragraph(f"Goals: {slo}")
    doc.add_heading("Acceptable Evidence", level=2)
    doc.add_paragraph("Assessment criteria based on SLO.")
    doc.add_heading("Learning Experiences", level=2)
    doc.add_paragraph(f"Instructional strategies: [Generated from {context}]")
    buffer = io.BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    return buffer, "\n".join([p.text for p in doc.paragraphs])

def generate_flashcards(grade, subject, topic, slo, context):
    doc = Document()
    doc.add_heading(f"Grade {grade} {subject} Flashcards: {topic}", 0)
    doc.add_paragraph(f"SLO: {slo}")
    table = doc.add_table(rows=6, cols=2)
    table.style = "Table Grid"
    table.cell(0, 0).text = "Front (Question)"
    table.cell(0, 1).text = "Back (Answer)"
    for i in range(1, 6):
        table.cell(i, 0).text = f"Question {i} about {topic}?"
        table.cell(i, 1).text = f"Answer {i} based on {context}."
    buffer = io.BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    table_text = "\n".join([f"{row.cells[0].text} | {row.cells[1].text}" for row in table.rows])
    return buffer, table_text

def generate_worksheet(grade, subject, topic, slo, context):
    doc = Document()
    doc.add_heading(f"Grade {grade} {subject} Worksheet: {topic}", 0)
    doc.add_paragraph(f"SLO: {slo}")
    doc.add_heading("Instructions", level=1)
    doc.add_paragraph("Complete the following questions.")
    doc.add_heading("Multiple Choice", level=2)
    for i in range(1, 4):
        doc.add_paragraph(f"{i}. Sample MCQ about {topic}? a) Option1 b) Option2 c) Option3 d) Option4")
    doc.add_heading("Short Answer", level=2)
    for i in range(1, 3):
        doc.add_paragraph(f"{i}. Short answer question about {topic}?")
    doc.add_heading("Activity", level=2)
    doc.add_paragraph(f"Activity based on {context}.")
    buffer = io.BytesIO()
    doc.save(buffer)
    buffer.seek(0)
    return buffer, "\n".join([p.text for p in doc.paragraphs])

# Global variables for FAISS index
faiss_index = None
chunks = []
embeddings = None

def process_and_generate(pdf_file, grade, subject, topic, slo, duration, output_type, feedback=""):
    global faiss_index, chunks, embeddings
    if not pdf_file:
        return None, "Please upload a PDF file / براہ کرم پی ڈی ایف فائل اپ لوڈ کریں", None
    try:
        # Extract and process PDF
        text = extract_text_from_pdf(pdf_file)
        if "Error" in text or "failed" in text:
            return None, text, None
        if not text.strip():
            return None, "No text extracted. Please upload a valid PDF / کوئی متن نہیں نکالا گیا۔ براہ کرم ایک درست پی ڈی ایف اپ لوڈ کریں", None
        
        # Chunk and embed text
        chunks = chunk_text(text)
        embeddings = create_embeddings(chunks)
        faiss_index = setup_faiss_index(embeddings)
        
        # Retrieve relevant context
        query = f"Grade {grade} {subject} {topic} {slo}"
        relevant_chunks = retrieve_relevant_chunks(query, faiss_index, chunks)
        context = " ".join(relevant_chunks)
        
        # Generate output
        if output_type == "Lesson Plan (BOPPPS)":
            buffer, preview = generate_lesson_plan_boppps(grade, subject, topic, slo or "General SLO", duration, context)
            filename = f"Grade_{grade}_{subject}_BOPPPS_Lesson_Plan.docx"
        elif output_type == "Lesson Plan (Backward Design)":
            buffer, preview = generate_lesson_plan_backward(grade, subject, topic, slo or "General SLO", duration, context)
            filename = f"Grade_{grade}_{subject}_Backward_Design_Lesson_Plan.docx"
        elif output_type == "Flashcards":
            buffer, preview = generate_flashcards(grade, subject, topic, slo or "General SLO", context)
            filename = f"Grade_{grade}_{subject}_Flashcards.docx"
        else:
            buffer, preview = generate_worksheet(grade, subject, topic, slo or "General SLO", context)
            filename = f"Grade_{grade}_{subject}_Worksheet.docx"
        
        # Handle feedback
        if feedback:
            with open("feedback.txt", "a") as f:
                f.write(f"{datetime.now()}: {feedback}\n")
        
        return buffer, preview, filename
    except Exception as e:
        return None, f"Error: {str(e)} / خرابی: {str(e)}", None

# Gradio interface
with gr.Blocks(title="Curriculum Assistant / نصابی اسسٹنٹ") as demo:
    gr.Markdown("""
    # Curriculum Assistant / نصابی اسسٹنٹ
    Upload a curriculum PDF and generate lesson plans, flashcards, or worksheets / نصابی پی ڈی ایف اپ لوڈ کریں اور سبق کے منصوبے، فلیش کارڈز، یا ورک شیٹس بنائیں.
    *Tip*: BOPPPS is great for structured lessons; Backward Design focuses on learning goals / BOPPPS منظم اسباق کے لیے بہترین ہے؛ Backward Design سیکھنے کے اہداف پر مرکوز ہے.
    """)
    
    with gr.Row():
        with gr.Column():
            pdf_file = gr.File(label="Upload Curriculum PDF / پی ڈی ایف اپ لوڈ کریں", file_types=[".pdf"])
            grade = gr.Dropdown(choices=list(range(1, 13)), label="Grade / گریڈ")
            subject = gr.Dropdown(choices=["Math", "Science", "Social Studies", "English"], label="Subject / مضمون")
            topic = gr.Textbox(label="Topic / موضوع", placeholder="e.g., Photosynthesis / مثلاً، فوٹوسنتھیسز")
            slo = gr.Textbox(label="Specific SLO (optional) / مخصوص SLO (اختیاری)", placeholder="e.g., Understand cell structure / مثلاً، خلیے کی ساخت کو سمجھیں")
            duration = gr.Dropdown(choices=[30, 45, 60], label="Lesson Duration (minutes) / سبق کا دورانیہ (منٹ)")
            output_type = gr.Radio(choices=["Lesson Plan (BOPPPS)", "Lesson Plan (Backward Design)", "Flashcards", "Worksheet"], label="Output Type / آؤٹ پٹ کی قسم")
            feedback = gr.Textbox(label="Feedback (optional) / رائے (اختیاری)", placeholder="Report issues or suggestions / مسائل یا تجاویز کی اطلاع دیں")
            submit_button = gr.Button("Generate / بنائیں")
        
        with gr.Column():
            preview = gr.Textbox(label="Preview / پیش منظر", lines=10, interactive=False)
            download_button = gr.File(label="Download as Word / ورڈ کے طور پر ڈاؤن لوڈ کریں")
    
    submit_button.click(
        fn=process_and_generate,
        inputs=[pdf_file, grade, subject, topic, slo, duration, output_type, feedback],
        outputs=[download_button, preview]
    )

demo.launch()