# app.py - v3.1: Manual PDF + Figure Upload with Downloadable Remixed Images (English Only, Nov 2025)
import gradio as gr
from gemini_nano import understand_image, chat, generate_image  # Import from root
from pypdf import PdfReader
import os
from dotenv import load_dotenv

load_dotenv()

pdf_context = ""  # Global for PDF text extraction

def extract_pdf_context(pdf_file):
    global pdf_context
    if pdf_file is None:
        pdf_context = ""
        return "No PDF uploaded. Upload optional PDF for context."
    try:
        reader = PdfReader(pdf_file)
        text = ""
        for page in reader.pages:  # Read all pages
            text += page.extract_text() + "\n"
        pdf_context = text  # No character limit
        return f"PDF context extracted ({len(pdf_context)} chars). Use for enhanced explanations."
    except Exception as e:
        return f"PDF extraction failed: {str(e)}"

def analyze_and_remix(figure_img, prompt, use_context=False):
    if figure_img is None:
        return "Please upload a figure screenshot.", None, None
    
    try:
        # Base description
        desc_prompt = ["Describe this academic figure/equation/table in detail for first-year students."]
        if use_context:
            desc_prompt[0] += f"\nPaper context: {pdf_context}"
        desc_prompt.append(figure_img)
        desc = understand_image(figure_img)  # Use updated function
        
        # Enhanced prompt for remix/annotation
        enhanced_prompt = prompt
        if "annotate" in prompt.lower() or "remix" in prompt.lower():
            enhanced_prompt += ("\n\nGenerate a high-resolution, colorful, annotated or remixed image. "
                                "Style: publication-quality scientific illustration, vibrant colors, clear labels for beginners. "
                                "Include layman's explanations in callouts.")
        
        # Generate response + image
        full_prompt = [enhanced_prompt, figure_img]
        if use_context:
            full_prompt[0] += f"\nPaper context: {pdf_context}"
        response_text = chat(full_prompt)  # Use updated chat function
        
        # Nano Banana Pro image generation (returns file path for download)
        img_file = generate_image(enhanced_prompt, figure_img)
        
        return desc, response_text, img_file
    except Exception as e:
        return f"Analysis failed: {str(e)}", None, None

# UI - Simplified Manual Upload with Download
with gr.Blocks(title="PaperSmith v3.1") as demo:
    gr.Markdown("# PaperSmith v3.1 – AI Figure Explainer & Remixer")
    gr.Markdown("Upload a PDF (optional, for context) + a figure screenshot → Enter request → Get description, explanation, and downloadable remixed image")
    
    with gr.Row():
        pdf_input = gr.File(label="Upload full paper PDF (optional, for context)", file_types=[".pdf"])
        figure_input = gr.Image(label="Upload figure screenshot from paper", type="pil")
    
    context_status = gr.Textbox(label="PDF Context Status", interactive=False)
    pdf_input.change(extract_pdf_context, pdf_input, context_status)
    
    with gr.Row():
        prompt_input = gr.Textbox(
            placeholder="e.g. 'Annotate this equation with colorful boxes and layman's explanation'",
            label="Enter your request (or use examples below)"
        )
        use_context = gr.Checkbox(label="Use PDF context for better accuracy", value=True)
    
    with gr.Row():
        desc_output = gr.Textbox(label="Step 1: Figure Description", lines=4)
        explanation_output = gr.Textbox(label="Step 2: Detailed Explanation/Remix", lines=6)
        remixed_file = gr.File(label="Step 3: Download Annotated/Remixed Image (PNG)", visible=True)
    
    submit_btn = gr.Button("Analyze & Remix", variant="primary")
    submit_btn.click(
        analyze_and_remix, 
        inputs=[figure_input, prompt_input, use_context], 
        outputs=[desc_output, explanation_output, remixed_file]
    )
    
    # Fixed Examples with proper chaining
    examples_component = gr.Examples(
        examples=[
            ["Annotate this equation with colorful boxes and layman's explanation"],
            ["Remix this figure to make it more intuitive and high-definition"],
            ["Convert this table into a beautiful infographic with key insights"],
            ["Explain this attention visualization for first-year students"],
            ["Add colorful callouts to this heatmap showing what each color means"],
            ["Reproduce this chart with vibrant colors and simple labels for a presentation"],
            ["Annotate this transformer architecture diagram step-by-step"],
            ["Turn this probability experiment into an interactive visual story"],
            ["Highlight key trends in this line plot with arrows and beginner notes"],
            ["Create a simplified version of this neural network diagram for slides"]
        ],
        inputs=[prompt_input]
    )

    # Chain .then() to load_input_event: Auto-populate prompt on example click
    examples_component.load_input_event.then(
        lambda ex: gr.update(value=ex[0]),  # Populate textbox with selected example
        outputs=prompt_input
    )
    
    gr.Markdown("**Powered by Gemini 2.5 Flash + Gemini 2.5 Flash Image (Nano Banana Pro)** | Group 23 – Generative AI in Creative Industries")

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