import streamlit as st
import os
from PyPDF2 import PdfReader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_groq import ChatGroq
from langchain_classic.chains import create_retrieval_chain
from langchain_classic.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from transformers import AutoTokenizer

# Automatically fetch the key from Hugging Face Secrets
GROQ_API_KEY = os.getenv("RAGpdf")

# ==========================================
# Core Pipeline Functions
# ==========================================

def extract_text_from_pdf(pdf_file):
    pdf_reader = PdfReader(pdf_file)
    text = ""
    for page in pdf_reader.pages:
        if page.extract_text():
            text += page.extract_text()
    return text

def tokenize_and_chunk_text(text):
    model_name = "sentence-transformers/all-MiniLM-L6-v2"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    text_splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
        tokenizer, chunk_size=500, chunk_overlap=50
    )
    return text_splitter.split_text(text)

def create_embeddings_and_vectorstore(text_chunks):
    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
    return FAISS.from_texts(texts=text_chunks, embedding=embeddings)

def query_and_generate_response(user_query, vectorstore, api_key):
    llm = ChatGroq(groq_api_key=api_key, model_name="llama-3.1-8b-instant", temperature=0.3)
    prompt = ChatPromptTemplate.from_template(
        """
        Answer the question based only on the provided context. 
        If the answer is not in the context, say "I cannot answer this based on the provided document."
        
        Context:
        {context}
        
        Question: {input}
        """
    )
    document_chain = create_stuff_documents_chain(llm, prompt)
    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
    retrieval_chain = create_retrieval_chain(retriever, document_chain)
    return retrieval_chain.invoke({"input": user_query})["answer"]

# ==========================================
# Streamlit User Interface
# ==========================================

st.title("📄 RAG App with Groq & Llama-3")
st.write("Upload a PDF, process it, and ask questions!")

if not GROQ_API_KEY:
    st.error("🚨 GROQ_API_KEY is missing! Please set it in your Space Settings under Secrets.")
    st.stop()

pdf_file = st.file_uploader("Upload your PDF document", type="pdf")

if pdf_file:
    if st.button("Process Document"):
        with st.spinner("Extracting text..."):
            raw_text = extract_text_from_pdf(pdf_file)
            
        with st.spinner("Tokenizing and Chunking..."):
            chunks = tokenize_and_chunk_text(raw_text)
            
        with st.spinner("Creating Embeddings..."):
            st.session_state.vectorstore = create_embeddings_and_vectorstore(chunks)
            st.success("Document processed and vector store built successfully!")

if "vectorstore" in st.session_state:
    st.divider()
    user_query = st.text_input("Ask a question about your document:")
    
    if user_query:
        with st.spinner("Generating answer..."):
            answer = query_and_generate_response(
                user_query, 
                st.session_state.vectorstore, 
                GROQ_API_KEY
            )
            st.write("**Response:**")
            st.write(answer)