import os
import streamlit as st
import fitz  # PyMuPDF
import faiss
import numpy as np
import pickle
from sentence_transformers import SentenceTransformer
import tiktoken
from groq import Groq

# Initialize embedding model
embed_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

# Function to extract text from PDF
def extract_text_from_pdf(pdf_file):
    doc = fitz.open(stream=pdf_file.read(), filetype="pdf")
    text = "\n".join([page.get_text("text") for page in doc])
    return text

# Function to split text into chunks
def chunk_text(text, chunk_size=512):
    tokenizer = tiktoken.get_encoding("cl100k_base")
    tokens = tokenizer.encode(text)
    chunks = [tokens[i:i+chunk_size] for i in range(0, len(tokens), chunk_size)]
    return ["".join(tokenizer.decode(chunk)) for chunk in chunks]

# Function to generate embeddings
def generate_embeddings(chunks):
    return embed_model.encode(chunks, convert_to_numpy=True)

# Function to store embeddings in FAISS
def store_in_faiss(embeddings, chunks):
    dimension = embeddings.shape[1]
    index = faiss.IndexFlatL2(dimension)
    index.add(embeddings)
    with open("faiss_index.pkl", "wb") as f:
        pickle.dump((index, chunks), f)
    return index

# Function to load FAISS index
def load_faiss():
    with open("faiss_index.pkl", "rb") as f:
        index, chunks = pickle.load(f)
    return index, chunks

# Function to search FAISS
def search_faiss(query, top_k=3):
    query_embedding = embed_model.encode([query])
    index, chunks = load_faiss()
    _, indices = index.search(query_embedding, top_k)
    results = [chunks[i] for i in indices[0]]
    return results

# Function to interact with Groq API
def query_groq(query):
    client = Groq(api_key=os.getenv("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"))
    response = client.chat.completions.create(
        messages=[{"role": "user", "content": query}],
        model="llama-3.3-70b-versatile"
    )
    return response.choices[0].message.content

# Streamlit UI
st.title("RAG-based PDF Q&A App")

uploaded_file = st.file_uploader("Upload a PDF", type="pdf")
if uploaded_file:
    st.write("Processing PDF...")
    text = extract_text_from_pdf(uploaded_file)
    chunks = chunk_text(text)
    embeddings = generate_embeddings(chunks)
    store_in_faiss(embeddings, chunks)
    st.success("PDF processed and indexed!")

query = st.text_input("Ask a question:")
if query:
    retrieved_chunks = search_faiss(query)
    context = " ".join(retrieved_chunks)
    response = query_groq(f"Context: {context} \n Question: {query}")
    st.write("### Answer:")
    st.write(response)