import os
import numpy as np
import streamlit as st
import PyPDF2
from transformers import AutoTokenizer, AutoModel
import torch
import faiss
from groq import Groq  # Make sure this import is correct

# Hardcoded API Key (for demonstration only, use caution with hardcoding keys)
API_KEY = "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"

# Initialize the Groq client
client = Groq(api_key=API_KEY)

# Load the transformer model and tokenizer (this can take some time initially)
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model.eval()  # Put the model in evaluation mode

def extract_text_from_pdfs(pdf_files):
    """Extract text from a list of PDF file paths."""
    text_data = []
    for pdf_path in pdf_files:
        with open(pdf_path, 'rb') as file:
            reader = PyPDF2.PdfReader(file)
            for page in reader.pages:
                text = page.extract_text()
                if text:
                    text_data.append(text)
    return text_data

def generate_embeddings(text_list):
    """Generate embeddings for a list of texts using the transformer model."""
    inputs = tokenizer(text_list, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1).numpy()  # Simple pooling
    return embeddings

def create_vector_db(embeddings):
    """Create a FAISS index for fast similarity search."""
    dim = embeddings.shape[1]
    index = faiss.IndexFlatL2(dim)
    index.add(embeddings)
    return index

def query_pdf_database(user_query, index, text_data, k=3):
    """Query the database to retrieve top-k relevant texts based on user input."""
    query_embedding = generate_embeddings([user_query])[0].reshape(1, -1)
    _, I = index.search(query_embedding, k)
    relevant_texts = [text_data[i] for i in I[0]]
    return relevant_texts

# Streamlit app setup
st.title("PDF-based Question Answering with RAG and Groq")

# Prompt for question input
user_input = st.text_input("Ask a question:")

if user_input:
    # Assume the PDFs are loaded and preprocessed into embeddings (for demonstration)
    pdf_files = ["The Constitution of the Islamic Republic of Pakistan.pdf", "data law.pdf"]  # Replace with actual file paths
    text_data = extract_text_from_pdfs(pdf_files)
    embeddings = generate_embeddings(text_data)
    index = create_vector_db(np.array(embeddings))

    # Query the database
    relevant_texts = query_pdf_database(user_input, index, text_data)

    # Use relevant texts as context for Groq API
    context = " ".join(relevant_texts)

    # Generate a response using the Groq API
    try:
        chat_completion = client.chat.completions.create(
            messages=[{"role": "user", "content": f"{user_input}\n\nContext: {context}"}],
            model="llama3-8b-8192",
        )
        st.write("Response:")
        st.write(chat_completion.choices[0].message.content)
    except Exception as e:
        st.error(f"An error occurred: {e}")