import os
import gradio as gr
from PyPDF2 import PdfReader
import requests
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Get the Hugging Face API token
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")

def summarize_text(text, instructions):
    API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B-Instruct"
    headers = {"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"}
    
    payload = {
        "inputs": f"{instructions}\n\nText to summarize:\n{text}",
        "parameters": {"max_length": 500}
    }
    
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()[0]["generated_text"]

def process_pdf(pdf_file, chunk_instructions, final_instructions):
    # Read PDF
    reader = PdfReader(pdf_file)
    text = ""
    for page in reader.pages:
        text += page.extract_text() + "\n\n"
    
    # Chunk the text (simple splitting by pages for this example)
    chunks = text.split("\n\n")
    
    # Agent 1: Summarize each chunk
    agent1_summaries = []
    for chunk in chunks:
        summary = summarize_text(chunk, chunk_instructions)
        agent1_summaries.append(summary)
    
    # Concatenate Agent 1 summaries
    concatenated_summary = "\n\n".join(agent1_summaries)
    
    # Agent 2: Final summarization
    final_summary = summarize_text(concatenated_summary, final_instructions)
    
    return final_summary

def pdf_summarizer(pdf_file, chunk_instructions, final_instructions):
    if pdf_file is None:
        return "Please upload a PDF file."
    
    try:
        summary = process_pdf(pdf_file.name, chunk_instructions, final_instructions)
        return summary
    except Exception as e:
        return f"An error occurred: {str(e)}"

# Gradio interface
iface = gr.Interface(
    fn=pdf_summarizer,
    inputs=[
        gr.File(label="Upload PDF"),
        gr.Textbox(label="Chunk Instructions", placeholder="Instructions for summarizing each chunk"),
        gr.Textbox(label="Final Instructions", placeholder="Instructions for final summarization")
    ],
    outputs=gr.Textbox(label="Summary"),
    title="PDF Earnings Summary Generator",
    description="Upload a PDF of an earnings summary or transcript to generate a concise summary."
)

iface.launch()