import torch
import gradio as gr
from transformers import pipeline
import logging

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

# Use a pipeline as a high-level helper
device = 0 if torch.cuda.is_available() else -1
text_summary = pipeline("summarization", model="facebook/bart-large-cnn", device=device, torch_dtype=torch.bfloat16)

# Function for summarization with enhancements
def summary(input, summary_type="medium"):
    # Check for empty input
    if not input.strip():
        return "Error: Please provide some text to summarize."

    # Calculate input length
    input_length = len(input.split())
    logging.info(f"Input length: {input_length} words")

    # Handle input that's too short
    if input_length < 10:
        return "Error: Input is too short. Please provide at least 10 words."

    # Handle input that's too long for the model
    if input_length > 512:
        return "Warning: Input exceeds the model's limit of 512 tokens. Please shorten the input text."

    # Adjust max/min lengths based on the summary type
    if summary_type == "short":
        max_output_tokens = max(10, input_length // 4)
    elif summary_type == "medium":
        max_output_tokens = max(20, input_length // 2)
    elif summary_type == "long":
        max_output_tokens = max(30, (3 * input_length) // 4)
    min_output_tokens = max(10, input_length // 6)

    # Generate summary
    output = text_summary(input, max_length=max_output_tokens, min_length=min_output_tokens, truncation=True)
    return output[0]['summary_text']

# Function to save the output summary to a file
def save_summary(summary_text):
    """Save the summarized text to a file."""
    with open("summary_output.txt", "w") as file:
        file.write(summary_text)
    return "Summary saved to 'summary_output.txt'."

# Gradio interface setup
gr.close_all()

# Create the Gradio interface
demo = gr.Interface(
    fn=summary,
    inputs=[
        gr.Textbox(label="INPUT THE PASSAGE TO SUMMARIZE", lines=15, placeholder="Paste your text here."),
        gr.Dropdown(["short", "medium", "long"], label="SUMMARY LENGTH", value="medium")
    ],
    outputs=[
        gr.Textbox(label="SUMMARIZED TEXT", lines=10, placeholder="Your summarized text will appear here."),
        gr.Button("Save Summary", click=save_summary)
    ],
    title="PAVISHINI @ GenAI Project 1: Text Summarizer",
    description=(
        "This application summarizes input text. "
        "The output length can be short, medium, or long based on your selection."
    ),
    live=True
)

demo.launch()