import streamlit as st
import os
import google.generativeai as genai
from PyPDF2 import PdfReader
from fpdf import FPDF

genai.configure(api_key="AIzaSyCHtsaOe5k-MXYO0RDnytc6iX7FpSKCPmE")
#genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))

generation_config = {
    "temperature": 2,
    "top_p": 0.95,
    "top_k": 64,
    "max_output_tokens": 65536,
    "response_mime_type": "text/plain",
}

model = genai.GenerativeModel(
    model_name="gemini-2.0-flash-thinking-exp-01-21",
    generation_config=generation_config,
)

def extract_text_from_pdf(file):
    """Extract text from an uploaded PDF file."""
    reader = PdfReader(file)
    text = ""
    for page in reader.pages:
        text += page.extract_text() + "\n"
    return text

def stream_response(prompt):
    """Stream response from the generative model."""
    chat_session = model.start_chat(
        history=[{"role": "user", "parts": [prompt]}]
    )
    for response_chunk in chat_session.send_message("Processing input...", stream=True):
        if response_chunk.text:
            yield response_chunk.text

def create_pdf(content):
    """Create a well-formatted PDF file from the given content."""
    pdf = FPDF()
    pdf.add_page()

    pdf.set_font("Arial", size=12)

    lines = content.split("\n")
    for line in lines:
        formatted_line = line.replace("**", "")  # Remove markdown-style bolding
        pdf.multi_cell(0, 10, formatted_line)

    return pdf

# Streamlit UI
st.title("Text Summarization and Insight Extraction")

# Session state to persist generated text
if "generated_text" not in st.session_state:
    st.session_state.generated_text = ""

uploaded_file = st.file_uploader("Upload a document (PDF) or paste text below:", type=["pdf"])
input_text = st.text_area("Or paste your article here:", placeholder="Paste your article or text here...")

if st.button("Generate Summary and Insights"):
    if uploaded_file:
        input_text = extract_text_from_pdf(uploaded_file)

    if input_text.strip():
        st.markdown("### Generating Results...")
        response_container = st.empty()
        st.session_state.generated_text = ""  # Clear previous text
        for chunk in stream_response(f"Summarize the following text and extract key insights:\n{input_text}"):
            st.session_state.generated_text += chunk
            response_container.markdown(st.session_state.generated_text)
    else:
        st.warning("Please upload a file or enter text to process.")

if st.session_state.generated_text:
    st.markdown("### Summary and Insights")
    st.text_area("Results:", st.session_state.generated_text, height=300)

    # PDF Download
    pdf = create_pdf(st.session_state.generated_text)
    pdf_file = "summary_insights.pdf"
    pdf.output(pdf_file)
    with open(pdf_file, "rb") as f:
        st.download_button(
            label="Download PDF",
            data=f,
            file_name="summary_insights.pdf",
            mime="application/pdf"
        )