import os
import streamlit as st
from transformers import AutoTokenizer, AutoModel, pipeline
import faiss
import pickle
import pdfplumber
from docx import Document

# Initialize variables
INDEX_FILE = "simple_rag_index.pkl"
EMBEDDING_DIM = 384  # Dimension for "all-MiniLM-L6-v2"

# Initialize FAISS index and document store
if os.path.exists(INDEX_FILE):
    with open(INDEX_FILE, "rb") as f:
        document_index, doc_store = pickle.load(f)
else:
    document_index = faiss.IndexFlatL2(EMBEDDING_DIM)
    doc_store = []

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

# Initialize a question-answering pipeline
qa_pipeline = pipeline("question-answering", model="distilbert-base-uncased", tokenizer="distilbert-base-uncased")

# Streamlit UI
st.title("Simple RAG Chatbot ")

# Function to extract text from PDF
def extract_pdf_text(file):
    with pdfplumber.open(file) as pdf:
        text = ""
        for page in pdf.pages:
            text += page.extract_text()
    return text

# Function to extract text from DOCX
def extract_docx_text(file):
    doc = Document(file)
    text = ""
    for para in doc.paragraphs:
        text += para.text + "\n"
    return text

# Step 1: Upload Document
uploaded_file = st.file_uploader("Upload a document (txt, pdf, docx)", type=["txt", "pdf", "docx"])
if uploaded_file:
    file_extension = uploaded_file.name.split('.')[-1].lower()
    
    if file_extension == "txt":
        document_content = uploaded_file.read().decode("utf-8")
    
    elif file_extension == "pdf":
        document_content = extract_pdf_text(uploaded_file)
    
    elif file_extension == "docx":
        document_content = extract_docx_text(uploaded_file)
    
    else:
        st.error("Unsupported file type!")
        document_content = None
    
    if document_content:
        st.write("Document uploaded successfully!")
        doc_store.append(document_content)

        # Generate embedding for the document
        tokens = tokenizer(document_content, return_tensors="pt", truncation=True, max_length=512)
        embedding = model(**tokens).last_hidden_state.mean(dim=1).detach().numpy()

        # Verify the embedding dimension matches FAISS index
        assert embedding.shape[1] == EMBEDDING_DIM, "Embedding dimension mismatch!"

        # Add embedding to FAISS index
        document_index.add(embedding)

        # Save the updated index and document store
        with open(INDEX_FILE, "wb") as f:
            pickle.dump((document_index, doc_store), f)

# Step 2: Ask a Question
user_query = st.text_input("Ask a question about the document:")
if user_query and len(doc_store) > 0:
    # Generate embedding for the query
    query_tokens = tokenizer(user_query, return_tensors="pt", truncation=True, max_length=512)
    query_embedding = model(**query_tokens).last_hidden_state.mean(dim=1).detach().numpy()

    # Retrieve the most relevant document
    distances, indices = document_index.search(query_embedding, k=1)
    closest_doc = doc_store[indices[0][0]]

    st.write("Retrieved Document Context:", closest_doc[:500], "...")  # Show a snippet of the document

    # Use QA pipeline to answer the query based on the retrieved document
    answer = qa_pipeline(question=user_query, context=closest_doc)
    st.write("Answer:", answer["answer"])