import os
import streamlit as st
from groq import Groq
from PyPDF2 import PdfReader
from sentence_transformers import SentenceTransformer
import faiss
import numpy as np

# Initialize Groq client
client = Groq(api_key="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")

# Load embedding model
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

# Initialize FAISS
embedding_dimension = 384  # Dimension of embeddings from the model
index = faiss.IndexFlatL2(embedding_dimension)

# Streamlit App - Islamic Theme
st.set_page_config(page_title="Quranic Therapy for Patients", page_icon="🕌", layout="wide")
st.markdown(
    """
    <style>
        .title {
            text-align: center;
            font-size: 3rem;
            font-weight: bold;
            color: #006400;  /* Islamic green */
        }
        .subheader {
            font-size: 1.5rem;
            font-weight: bold;
            color: #3E4E50;
        }
        .footer {
            text-align: center;
            font-size: 0.9rem;
            color: #888888;
            margin-top: 50px;
        }
        .button {
            background-color: #006400;
            color: white;
            padding: 10px 20px;
            font-size: 1rem;
            border-radius: 5px;
        }
        .button:hover {
            background-color: #004d00;
        }
        .container {
            margin-top: 20px;
        }
    </style>
    """,
    unsafe_allow_html=True,
)

# Title
st.markdown('<div class="title">🕌 Quranic Therapy for Patients</div>', unsafe_allow_html=True)
st.markdown("---")

# Sidebar for Upload
st.sidebar.header("Upload Your Quranic PDF")
uploaded_file = st.sidebar.file_uploader("Upload a PDF file containing Quranic verses", type="pdf")

if uploaded_file:
    # Step 1: Extract text from PDF
    st.markdown('<div class="subheader">1. Extracted Text from PDF</div>', unsafe_allow_html=True)
    pdf_reader = PdfReader(uploaded_file)
    pdf_text = ""
    for page in pdf_reader.pages:
        pdf_text += page.extract_text()

    if not pdf_text.strip():
        st.error("Could not extract text from the PDF. Please upload a readable PDF.")
    else:
        st.success("PDF text successfully extracted!")
        with st.expander("View Extracted Text", expanded=False):
            st.write(pdf_text[:3000] + "..." if len(pdf_text) > 3000 else pdf_text)

        # Step 2: Generate Summary (Optional Feature)
        if st.button("Generate Summary", key="summary_button"):
            st.info("Generating summary...")
            summary = embedding_model.encode([pdf_text[:1000]])  # Mock summary generation
            st.success("Summary generated!")
            st.write("🚀 **Summary:** This is a mock summary. Replace with your own summarization logic.")

        # Step 3: Split text into chunks and create embeddings
        st.markdown('<div class="subheader">2. Processing PDF Content</div>', unsafe_allow_html=True)
        with st.spinner("Splitting text into chunks and generating embeddings..."):
            chunk_size = 500  # Approximate characters per chunk
            chunks = [pdf_text[i:i + chunk_size] for i in range(0, len(pdf_text), chunk_size)]
            embeddings = embedding_model.encode(chunks, convert_to_numpy=True)
            index.add(embeddings)
        st.success(f"Successfully processed {len(chunks)} chunks and stored embeddings in FAISS!")

        # Step 4: Ask a Question
        st.markdown('<div class="subheader">3. Ask a Question</div>', unsafe_allow_html=True)
        user_query = st.text_input("Enter your question about the Quranic text:")

        if st.button("Get Answer", key="answer_button"):
            if user_query:
                # Step 4.1: Generate embedding for the query
                query_embedding = embedding_model.encode([user_query], convert_to_numpy=True)

                # Step 4.2: Search the vector database
                k = 5  # Number of chunks to retrieve
                distances, indices = index.search(query_embedding, k)
                retrieved_chunks = [chunks[i] for i in indices[0]]

                # Step 4.3: Use the retrieved chunks as context for Groq API
                context = " ".join(retrieved_chunks)
                prompt = (
                    f"Context: {context}\n\n"
                    f"Based on the above context, answer the following question:\n"
                    f"{user_query}"
                )

                # Call Groq API
                try:
                    chat_completion = client.chat.completions.create(
                        messages=[{"role": "user", "content": prompt}],
                        model="llama-3.3-70b-versatile",
                    )
                    response = chat_completion.choices[0].message.content

                    # Display the response
                    st.markdown('<div class="subheader">Response</div>', unsafe_allow_html=True)
                    st.success(response)

                except Exception as e:
                    st.error(f"Error interacting with Groq API: {e}")
            else:
                st.warning("Please enter a query before clicking 'Get Answer'.")

        # Clear Cache Button
        if st.button("Clear Cache", key="clear_cache"):
            index.reset()
            st.success("Cache cleared!")

# Footer with Islamic theme
st.markdown('<div class="footer">Made with ❤️ using Streamlit | Powered by Generative AI for Quranic Therapy</div>', unsafe_allow_html=True)