# -*- coding: utf-8 -*-
"""app.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1XblbxoRxB4XOHixjGij789FPD9KjKdhi
"""

import os
import PyPDF2
import gradio as gr
from langchain_groq.chat_models import ChatGroq

# Set Groq API key securely
GROQ_API_KEY = os.getenv("GROQ_API_KEY")  # Fetch from environment variables

# Ensure API key is available
if not GROQ_API_KEY:
    raise ValueError("GROQ_API_KEY is not set. Add it in Hugging Face Secrets.")

# Initialize LLM (Mistral-8x7B)
llm = ChatGroq(model_name="mixtral-8x7b-32768")

def extract_text_from_pdf(pdf_file):
    """Extract text from a PDF file."""
    text = ""
    reader = PyPDF2.PdfReader(pdf_file)
    for page in reader.pages:
        page_text = page.extract_text()
        if page_text:
            text += page_text + "\n"
    return text

def summarize_text(text, length='Medium', style='Concise Paragraph'):
    """Summarize the text with CoT and adjustable format."""

    # Adjust summary length
    length_map = {
        'Short': 'Summarize in 3-4 lines.',
        'Medium': 'Summarize in 6-8 lines.',
        'Long': 'Provide a detailed summary in multiple paragraphs.'
    }

    # Adjust summary style
    style_map = {
        'Bulleted List': 'Format the summary as a list of key points.',
        'Key Takeaways': 'Extract the most important insights as key takeaways.',
        'Concise Paragraph': 'Write the summary as a structured paragraph.'
    }

    prompt = f"""
    Step 1: Identify the main topics covered in the document.
    Step 2: Extract key facts, arguments, and conclusions.
    Step 3: Generate a structured summary based on extracted information.
    {length_map[length]} {style_map[style]}

    Document:
    {text[:10000]}
    """

    response = llm.predict(prompt)
    return response

def process_pdf(file, length, style):
    """Extract text and summarize PDF using Mistral-8x7B with customization."""
    if file is None:
        return "No file uploaded."

    text = extract_text_from_pdf(file)
    summary = summarize_text(text, length, style)
    return summary

# Create Gradio Interface
interface = gr.Interface(
    fn=process_pdf,
    inputs=[
        gr.File(label="Upload a PDF"),
        gr.Radio(["Short", "Medium", "Long"], value="Medium", label="Summary Length"),
        gr.Radio(["Bulleted List", "Key Takeaways", "Concise Paragraph"], value="Concise Paragraph", label="Summary Style")
    ],
    outputs="text",
    title="📄 AI-Powered PDF Summarizer",
    description="Upload a PDF file and customize the summary format and length."
)

# Run the app
interface.launch()