import gradio as gr
import PyPDF2
from pptx import Presentation
from PIL import Image
import io
from google import genai
from jinja2 import Template
import fitz  # PyMuPDF
import os
import logging
import re
import time

# Set up logging
logging.basicConfig(level=logging.INFO)

# Function to extract content from PDF
def extract_content_from_pdf(file_path):
    try:
        text = ""
        images = []
        doc = fitz.open(file_path)
        for page in doc:
            text += page.get_text() + "\n"
            for img in page.get_images():
                xref = img[0]
                base_image = doc.extract_image(xref)
                image_bytes = base_image["image"]
                image = Image.open(io.BytesIO(image_bytes))
                images.append(image)
        return text, images
    except Exception as e:
        logging.error(f"Error extracting content from PDF: {e}")
        return "", []

# Function to extract content from PPTX
def extract_content_from_pptx(file_path):
    try:
        text = ""
        images = []
        prs = Presentation(file_path)
        for slide in prs.slides:
            for shape in slide.shapes:
                if hasattr(shape, 'text'):
                    text += shape.text + "\n"
                if shape.shape_type == 13:  # Picture
                    image = shape.image
                    image_bytes = image.blob
                    image = Image.open(io.BytesIO(image_bytes))
                    images.append(image)
        return text, images
    except Exception as e:
        logging.error(f"Error extracting content from PPTX: {e}")
        return "", []

# Function to process file
def process_file(file_path):
    if file_path is None:
        return "No file uploaded", []

    try:
        if file_path.lower().endswith('.pdf'):
            return extract_content_from_pdf(file_path)
        elif file_path.lower().endswith('.pptx'):
            return extract_content_from_pptx(file_path)
        else:
            return "Unsupported file format", []
    except Exception as e:
        logging.error(f"Error processing file: {e}")
        return f"An error occurred while processing the file: {str(e)}", []

# Function to clean response
def clean_response(response_text):
    # Remove code block markers if present
    cleaned = re.sub(r'```python|```', '', response_text).strip()
    # Handle newlines and indentation
    cleaned = re.sub(r'\n\s*\n', '\n\n', cleaned)
    return cleaned

# Function to understand text and images using Gemini V2 API
def understand_content(api_key, text, images, progress=gr.Progress()):
    try:
        # Initialize the Gemini client
        client = genai.Client(api_key=api_key)

        progress(0.3, desc="Preparing content...")
        # Prepare content for Gemini
        content = [text]
        for image in images[:10]:  # Limit to 10 images
            content.append(image)

        progress(0.5, desc="Generating unit plan...")
        # Generate response from Gemini
        prompt = Template("""
You are an expert instructional designer.
Below are materials shared by a teacher. Your role is to reverse-engineer the unit planner for this content.
To do so:
1) Read the content carefully
2) Create a unit planner for this content that follows this exact structure:
{{ unit_plan_structure }}
""").render(unit_plan_structure="""
# Standards
# Transfer Goal
# Essential Questions
# Enduring Understandings
# Students will know
# Students will be able to
# Formative Assessments
# Summative Assessments
# Scope and Sequence
# Unit Overview
# Potential barriers
# Connections
## ISP Core Values
## IB Theory of Knowledge
## IB Approaches to Learning
# Authentic Assessment
""")

        response = client.models.generate_content(
            model="gemini-2.0-flash-exp",
            contents=prompt + text
        )

        progress(0.8, desc="Finalizing output...")
        # Log the raw response for debugging
        response_text = response.text
        logging.info(f"Raw response from Gemini: {response_text}")

        # Clean the response
        cleaned_response = clean_response(response_text)
        logging.info(f"Cleaned response: {cleaned_response}")

        progress(1.0, desc="Complete!")
        return cleaned_response

    except Exception as e:
        logging.error(f"Error in content understanding: {e}")
        return f"Error in processing the content: {str(e)}. Please check your API key and try again."

# Function to reverse engineer unit plan
def generate_elt_plan(api_key, file, progress=gr.Progress()):
    if not api_key:
        return "Please enter your Gemini API key", None

    try:
        progress(0.1, desc="Starting file processing...")
        logging.info(f"Processing file: {file.name}")
        content, images = process_file(file.name)
        if isinstance(content, str) and content.startswith("An error occurred"):
            return content, None
        logging.info(f"Extracted content length: {len(content)}, Number of images: {len(images)}")
        progress(0.2, desc="Content extracted, generating plan...")
        elt_plan = understand_content(api_key, content, images, progress)

        # Create a downloadable text file
        timestamp = time.strftime("%Y%m%d-%H%M%S")
        filename = f"unit_plan_{timestamp}.txt"
        with open(filename, "w", encoding="utf-8") as f:
            f.write(elt_plan)

        return elt_plan, filename
    except Exception as e:
        logging.error(f"Error in generate_elt_plan: {e}")
        return f"An error occurred: {str(e)}", None

# Set up Gradio Blocks
with gr.Blocks() as demo:
    gr.Markdown("# 🔄 Reverse Unit Planner")

    with gr.Row():
        with gr.Column():
            api_key = gr.Textbox(label="Enter your Gemini API key", type="password")
            file_input = gr.File(label="Upload PPTX or PDF")
            submit_btn = gr.Button("Reverse engineer a unit plan", variant="primary")

        with gr.Column():
            output = gr.Markdown(label="Draft Unit Planner")
            copy_btn = gr.Button("📋 Copy to Clipboard")
            download_btn = gr.File(label="Download Unit Plan")

    # Handle the main processing
    result = submit_btn.click(
        generate_elt_plan,
        inputs=[api_key, file_input],
        outputs=[output, download_btn]
    )

    # Add copy to clipboard functionality
    copy_btn.click(
        fn=None,
        inputs=output,
        outputs=None,
        js="""
        async (text) => {
            await navigator.clipboard.writeText(text);
            await new Promise(resolve => {
                const notification = document.createElement('div');
                notification.textContent = 'Copied to clipboard!';
                notification.style.position = 'fixed';
                notification.style.bottom = '20px';
                notification.style.left = '50%';
                notification.style.transform = 'translateX(-50%)';
                notification.style.backgroundColor = '#4CAF50';
                notification.style.color = 'white';
                notification.style.padding = '10px 20px';
                notification.style.borderRadius = '5px';
                notification.style.zIndex = '1000';
                document.body.appendChild(notification);
                setTimeout(() => {
                    notification.remove();
                    resolve();
                }, 2000);
            });
        }
    """)

# Launch the app
demo.launch()