import gradio as gr
import PyPDF2
from transformers import pipeline
import re

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

def extract_text_from_pdf(pdf_file):
    """Extract text from uploaded PDF file"""
    try:
        pdf_reader = PyPDF2.PdfReader(pdf_file)
        text = ""
        for page in pdf_reader.pages:
            text += page.extract_text()
        return text
    except Exception as e:
        return f"Error reading PDF: {str(e)}"

def extract_key_elements(text, summary):
    """Extract key elements from text and summary"""
    summary_sentences = re.split(r'[.!?]+', summary.strip())
    
    key_elements = {
        "Main Summary": summary,
        "Total Words Original": len(text.split()),
        "Total Words Summary": len(summary.split()),
        "Compression Ratio": f"{(len(summary.split()) / len(text.split()) * 100):.1f}%",
        "Key Points": "\n".join([s.strip() for s in summary_sentences if s.strip()][:5])
    }
    return key_elements

def summarize_text(text_input, max_length=150, min_length=50):
    """Summarize text with key elements"""
    if not text_input or len(text_input.strip()) < 50:
        return "Error: Please provide text with at least 50 characters", {}
    
    try:
        # Split long text into chunks
        words = text_input.split()
        chunk_size = 1024
        chunks = [' '.join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]
        
        summaries = []
        for chunk in chunks:
            if len(chunk.split()) > 30:
                result = summarizer(chunk, max_length=max_length, min_length=min_length, do_sample=False)
                summaries.append(result[0]['summary_text'])
        
        final_summary = " ".join(summaries)
        key_elements = extract_key_elements(text_input, final_summary)
        
        return final_summary, key_elements
    except Exception as e:
        return f"Error: {str(e)}", {}

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📄 PDF & Text Summarizer")
    gr.Markdown("Upload PDF or paste text to generate summaries with key elements using BART model.")
    
    with gr.Tabs():
        # Text Tab
        with gr.Tab("Text Summarizer"):
            with gr.Row():
                text_input = gr.Textbox(label="Enter Text", placeholder="Paste text here...", lines=8)
            
            with gr.Row():
                max_len = gr.Slider(50, 300, value=150, step=10, label="Max Length")
                min_len = gr.Slider(20, 100, value=50, step=10, label="Min Length")
            
            summarize_btn = gr.Button("🚀 Summarize", variant="primary")
            summary_output = gr.Textbox(label="Summary", lines=6)
            key_elements_output = gr.JSON(label="Key Elements")
            
            summarize_btn.click(
                fn=summarize_text,
                inputs=[text_input, max_len, min_len],
                outputs=[summary_output, key_elements_output]
            )
        
        # PDF Tab
        with gr.Tab("PDF Summarizer"):
            pdf_input = gr.File(label="Upload PDF", file_count="single", file_types=[".pdf"])
            
            with gr.Row():
                pdf_max_len = gr.Slider(50, 300, value=150, step=10, label="Max Length")
                pdf_min_len = gr.Slider(20, 100, value=50, step=10, label="Min Length")
            
            pdf_summarize_btn = gr.Button("📤 Process PDF", variant="primary")
            pdf_summary_output = gr.Textbox(label="Summary", lines=6)
            pdf_key_elements_output = gr.JSON(label="Key Elements")
            
            pdf_summarize_btn.click(
                fn=lambda pdf, max_l, min_l: summarize_text(extract_text_from_pdf(pdf), max_l, min_l),
                inputs=[pdf_input, pdf_max_len, pdf_min_len],
                outputs=[pdf_summary_output, pdf_key_elements_output]
            )
    
    gr.Markdown("---")
    gr.Markdown("✨ Features: PDF & Text support | BART-large-cnn model | Key element extraction")

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