import streamlit as st
import os
import PyPDF2
from groq import Groq

# Set up your Groq client
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# Function to extract text from a PDF
def extract_text_from_pdf(pdf_file):
    pdf_reader = PyPDF2.PdfReader(pdf_file)
    text = ""
    for page in pdf_reader.pages:
        text += page.extract_text()
    return text

# Function to query Groq for simplified content or Q&A
def query_groq(prompt):
    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
        model="llama-3.3-70b-versatile",
    )
    return chat_completion.choices[0].message.content

# Persistent history
if "history" not in st.session_state:
    st.session_state["history"] = []

# Streamlit App
st.set_page_config(
    page_title="🌍 Next-Gen Scientific Paper Translator",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Apply custom CSS for animations and background
st.markdown(
    """
    <style>
    body {
        background-image: url('https://source.unsplash.com/1920x1080/?science,technology');
        background-size: cover;
        background-attachment: fixed;
        color: #FFFFFF;
    }
    .css-18e3th9 {
        background: rgba(0, 0, 0, 0.7) !important;
        border-radius: 10px;
        padding: 20px;
    }
    button {
        transition: 0.3s ease;
    }
    button:hover {
        background-color: #1f78d1 !important;
        color: white !important;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

st.title("📄 **Next-Gen Scientific Paper Translator**")
st.markdown(
    "Unlock the power of **scientific papers** with simplified explanations, extracted citations, and more!"
)

# Sidebar with user options
with st.sidebar:
    st.markdown("## ⚙️ **Settings**")
    theme_mode = st.radio("Choose Theme", ["Light Mode", "Dark Mode"])
    if theme_mode == "Dark Mode":
        st.markdown("""
            <style>
                body {
                    background-color: #1E1E1E;
                }
            </style>
        """, unsafe_allow_html=True)

# Upload PDF
st.markdown("### 📤 **Upload Your PDF**")
pdf_file = st.file_uploader("Drop your scientific paper (PDF only)", type=["pdf"])

# Left column for preview, right column for output
if pdf_file is not None:
    col1, col2 = st.columns([1, 2])

    # PDF Preview
    with col1:
        st.markdown("### 📑 **PDF Preview**")
        text = extract_text_from_pdf(pdf_file)
        st.text_area("Preview of Content", text[:1500], height=400)

    # Right column: Options for processing
    with col2:
        st.markdown("### 🧠 **What would you like to do?**")
        task = st.radio(
            "Select a task",
            ["Simplify the Content", "Extract Title", "Extract Citation", "Ask a Question"]
        )

        if task == "Simplify the Content":
            prompt = f"Simplify this for a beginner:\n\n{text[:1000]}"
            if st.button("Simplify"):
                with st.spinner("Processing..."):
                    response = query_groq(prompt)
                    st.session_state["history"].append({"task": "Simplify", "output": response})
                    st.success("Simplified!")
                    st.write(response)

        elif task == "Extract Title":
            prompt = f"Extract the title:\n\n{text[:1000]}"
            if st.button("Extract Title"):
                with st.spinner("Processing..."):
                    response = query_groq(prompt)
                    st.session_state["history"].append({"task": "Title", "output": response})
                    st.success("Title extracted!")
                    st.write(response)

        elif task == "Extract Citation":
            prompt = f"Extract the citation:\n\n{text[:1000]}"
            if st.button("Extract Citation"):
                with st.spinner("Processing..."):
                    response = query_groq(prompt)
                    st.session_state["history"].append({"task": "Citation", "output": response})
                    st.success("Citation extracted!")
                    st.write(response)

        elif task == "Ask a Question":
            user_question = st.text_input("Enter your question:")
            if st.button("Get Answer"):
                prompt = f"Answer this question:\n\n{text[:1000]}\n\nQuestion: {user_question}"
                with st.spinner("Processing..."):
                    response = query_groq(prompt)
                    st.session_state["history"].append({"task": "Q&A", "output": response})
                    st.success("Answer generated!")
                    st.write(response)

# History Section
st.markdown("---")
st.markdown("### 🕒 **History**")
for entry in st.session_state["history"]:
    st.markdown(f"**Task:** {entry['task']} | **Output:** {entry['output']}")

# Footer
st.markdown("---")
st.markdown(
    "<div style='text-align: center;'>Built with ❤️ using Streamlit and Groq</div>",
    unsafe_allow_html=True,
)