import gradio as gr
from pdf2image import convert_from_path
from PIL import Image
import pytesseract
from transformers import pipeline
import tempfile
import os

# Load Hugging Face summarization pipeline
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

# Chunk text into smaller pieces for model input limit
def chunk_text(text, max_tokens=1000):
    words = text.split()
    for i in range(0, len(words), max_tokens):
        yield " ".join(words[i:i+max_tokens])

# Main function: PDF upload → OCR → Summarization
def summarize_image_pdf(pdf_file):
    try:
        # Save uploaded file temporarily
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
            tmp.write(pdf_file.read())
            tmp_path = tmp.name

        # Convert PDF to images (Poppler is pre-installed on Spaces)
        images = convert_from_path(tmp_path, dpi=300)

        # OCR: Extract text from images
        extracted_text = ""
        for i, img in enumerate(images):
            text = pytesseract.image_to_string(img)
            extracted_text += f"\n\n--- Page {i+1} ---\n{text}"

        os.remove(tmp_path)  # Clean up temp file

        # Summarize text
        summaries = []
        for chunk in chunk_text(extracted_text):
            summary = summarizer(chunk, max_length=130, min_length=30, do_sample=False)[0]['summary_text']
            summaries.append(summary)

        return "\n\n".join(summaries)

    except Exception as e:
        return f"❌ Error: {str(e)}"

# Gradio UI
interface = gr.Interface(
    fn=summarize_image_pdf,
    inputs=gr.File(label="📄 Upload a scanned/image PDF", file_types=[".pdf"]),
    outputs=gr.Textbox(label="📝 Summary"),
    title="📘 OCR PDF Summarizer",
    description="Upload a scanned or image-based PDF. The app will extract text using OCR and summarize it using Hugging Face's BART model.",
)

interface.launch()