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


# Load your embedding model (e.g., all-MiniLM)
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

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

# Helper: Extract text from uploaded PDF
def extract_text_from_pdf(file):
    pdf = PdfReader(file)
    text = ""
    for page in pdf.pages:
        text += page.extract_text()
    return text

# Helper: Chunk text
def chunk_text(text, chunk_size=500):
    words = text.split()
    return [" ".join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]

# Helper: Create FAISS vector store
def create_faiss_index(chunks):
    embeddings = embedding_model.encode(chunks)
    dim = embeddings.shape[1]
    index = faiss.IndexFlatL2(dim)
    index.add(embeddings)
    return index, embeddings, chunks

# Helper: Retrieve top-k similar chunks
def retrieve_similar_chunks(query, index, chunks, embeddings, k=3):
    query_vector = embedding_model.encode([query])
    _, I = index.search(query_vector, k)
    return [chunks[i] for i in I[0]]

# Streamlit UI
st.title("🧠 PDF-based RAG App using Groq")
uploaded_file = st.file_uploader("Upload a PDF", type="pdf")

if uploaded_file:
    text = extract_text_from_pdf(uploaded_file)
    chunks = chunk_text(text)
    index, embeddings, stored_chunks = create_faiss_index(chunks)
    st.success("PDF processed and vector DB created.")

    user_query = st.text_input("Ask a question based on the document:")
    if user_query:
        top_chunks = retrieve_similar_chunks(user_query, index, stored_chunks, embeddings)

        # Build context for Groq
        context = "\n".join(top_chunks)
        prompt = f"Answer the question based on the context below:\n\nContext:\n{context}\n\nQuestion:\n{user_query}"

        # Send to Groq
        response = client.chat.completions.create(
            model="llama3-8b-8192",
            messages=[
                {"role": "user", "content": prompt}
            ]
        )

        st.subheader("Answer:")
        st.write(response.choices[0].message.content)