import streamlit as st
import os
import time
import hashlib
from io import BytesIO
from langchain_groq import ChatGroq
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.chains import create_retrieval_chain, create_history_aware_retriever
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.vectorstores import FAISS
from langchain_community.document_loaders import PyPDFLoader
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.documents import Document

# ==========================================
# 10. SYSTEM CONFIGURATIONS & CONSTANTS
# ==========================================
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
GROQ_MODEL = "llama-3.3-70b-versatile"
CHUNK_SIZE = 700
CHUNK_OVERLAP = 100
MAX_HISTORY_PAIRS = 10  # Keeps up to 20 messages total (10 user, 10 assistant)

# Retriever Configuration (MMR)
SEARCH_TYPE = "mmr"
K_DOCS = 5
FETCH_K = 20

# ==========================================
# 11. SECURE API KEY MANAGEMENT
# ==========================================
# Hugging Face Spaces automatically exposes repo secrets to os.environ
groq_api_key = os.environ.get("GROQ_API_KEY")

if not groq_api_key:
    # Fallback for local testing if not deployed to HF Spaces yet
    from dotenv import load_dotenv
    load_dotenv()
    groq_api_key = os.getenv("GROQ_API_KEY")

# Initialize LLM
if groq_api_key:
    llm = ChatGroq(groq_api_key=groq_api_key, model_name=GROQ_MODEL)
else:
    llm = None

st.set_page_config(page_title="Production RAG Chatbot", layout="wide")
st.title("📚 Advanced RAG Document Q&A with Chat History")

# ==========================================
# 2. CACHED EMBEDDING MODEL LOADING
# ==========================================
@st.cache_resource
def load_embedding_model():
    """Loads and caches the embedding model globally across sessions."""
    return HuggingFaceEmbeddings(
        model_name=EMBEDDING_MODEL,
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True}
    )

embeddings = load_embedding_model()

# ==========================================
# 1. FILE HASHING & PROCESSING
# ==========================================
def calculate_files_hash(uploaded_files):
    """Generates a unique SHA-256 hash combining all uploaded files' contents."""
    hasher = hashlib.sha256()
    for uploaded_file in sorted(uploaded_files, key=lambda x: x.name):
        # Read file chunks to handle large files efficiently without memory spikes
        bytes_data = uploaded_file.getvalue()
        hasher.update(bytes_data)
    return hasher.hexdigest()

def process_uploaded_files(uploaded_files):
    """Loads, splits, and embeds all pages from uploaded PDFs."""
    all_docs = []
    
    for uploaded_file in uploaded_files:
        # Streamlit files exist in memory; write temporarily to parse via PyPDFLoader
        temp_filename = f"temp_{uploaded_file.name}"
        with open(temp_filename, "wb") as f:
            f.write(uploaded_file.getvalue())
        
        try:
            loader = PyPDFLoader(temp_filename)
            # 4. Process ALL documents (removed [:50] restriction)
            for page in loader.load():
                # Add human-readable metadata for tracking source info
                page.metadata["source"] = uploaded_file.name
                all_docs.append(page)
        finally:
            if os.path.exists(temp_filename):
                os.remove(temp_filename)
                
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=CHUNK_SIZE, 
        chunk_overlap=CHUNK_OVERLAP
    )
    final_documents = text_splitter.split_documents(all_docs)
    
    # Create the FAISS Vector Database
    vector_store = FAISS.from_documents(final_documents, embeddings)
    return vector_store

# ==========================================
# 3. CHAIN INITIALIZATION (BUILD ONCE)
# ==========================================
def build_retrieval_chain(vector_store):
    """Builds and caches the entire LangChain execution pipeline."""
    # 5. Configure Retriever with Maximal Marginal Relevance (MMR)
    retriever = vector_store.as_retriever(
        search_type=SEARCH_TYPE,
        search_kwargs={"k": K_DOCS, "fetch_k": FETCH_K}
    )
    
    contextualize_q_prompt = ChatPromptTemplate.from_messages([
        ("system", "Given a chat history and the latest user question which might reference context in the chat history, formulate a standalone question which can be understood without the chat history. Do NOT answer the question, just reformulate it if needed and otherwise return it as is."),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}")
    ])
    
    qa_prompt = ChatPromptTemplate.from_messages([
        ("system", "Answer the user's questions based ONLY on the provided context. Provide the most accurate response.\n\nContext:\n{context}"),
        MessagesPlaceholder(variable_name="chat_history"),
        ("human", "{input}")
    ])
    
    history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt)
    document_chain = create_stuff_documents_chain(llm, qa_prompt)
    
    return create_retrieval_chain(history_aware_retriever, document_chain)

# ==========================================
# STREAMLIT STATE & INTERFACE MANAGEMENT
# ==========================================
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "current_file_hash" not in st.session_state:
    st.session_state.current_file_hash = None
if "retreival_chain" not in st.session_state:
    st.session_state.retreival_chain = None

# Sidebar for file upload
with st.sidebar:
    st.header("Document Management")
    uploaded_files = st.file_uploader(
        "Upload your Research Papers (PDF)", 
        type=["pdf"], 
        accept_multiple_files=True
    )
    
    if not groq_api_key:
        st.error("GROQ_API_KEY is missing. Please configure it in your Space Secrets environment.")

# Handle Vector Database compilation conditionally
if uploaded_files and llm:
    new_hash = calculate_files_hash(uploaded_files)
    
    # Trigger processing only if files have changed or chain hasn't been built yet
    if st.session_state.current_file_hash != new_hash:
        with st.sidebar:
            with st.spinner("Processing new documents..."):
                try:
                    vs = process_uploaded_files(uploaded_files)
                    # Compile and store the chain once
                    st.session_state.retreival_chain = build_retrieval_chain(vs)
                    st.session_state.current_file_hash = new_hash
                    st.success("Vector Store & Chains Ready!")
                except Exception as e:
                    st.error(f"Error parsing documents: {e}")

# Display historic chat items
for message in st.session_state.chat_history:
    role = "user" if isinstance(message, HumanMessage) else "assistant"
    with st.chat_message(role):
        st.write(message.content)

# Chat execution block
if user_prompt := st.chat_input("Ask something about your uploaded documents..."):
    if not uploaded_files:
        st.warning("Please upload at least one PDF document first.")
    elif not st.session_state.retreival_chain:
        st.error("Retrieval pipeline is not initialized.")
    else:
        # User input visual representation
        with st.chat_message("user"):
            st.write(user_prompt)
            
        with st.chat_message("assistant"):
            with st.spinner("Analyzing context..."):
                # 7. Use perf_counter to accurately capture real elapsed turnaround time
                start_time = time.perf_counter()
                
                response = st.session_state.retreival_chain.invoke({
                    'input': user_prompt,
                    'chat_history': st.session_state.chat_history
                })
                
                elapsed_time = time.perf_counter() - start_time
                
                st.write(response['answer'])
                st.caption(f"Response compiled in {elapsed_time:.2f} seconds")
                
                # 9. Improved Reference Tracker with metadata source & page tracking
                with st.expander("View Reference Passages"):
                    for i, doc in enumerate(response['context']):
                        source_file = doc.metadata.get("source", "Unknown Document")
                        # Page index typically starts at 0 internally; add 1 for user reading clarity
                        page_num = doc.metadata.get("page", 0) + 1
                        
                        st.markdown(f"**Source {i+1}:** {source_file} (Page {page_num})")
                        st.write(doc.page_content)
                        st.divider()
                        
        # 8. Slide window memory to prevent unbounded chat history growth
        st.session_state.chat_history.append(HumanMessage(content=user_prompt))
        st.session_state.chat_history.append(AIMessage(content=response['answer']))
        
        if len(st.session_state.chat_history) > (MAX_HISTORY_PAIRS * 2):
            st.session_state.chat_history = st.session_state.chat_history[-(MAX_HISTORY_PAIRS * 2):]