import streamlit as st
import os
import tempfile
import pickle
import faiss
import numpy as np
from helper import extract_text_from_pdf, chunk_text, embedding_function, embedding_model, query_llm_with_context
import logging

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Set page configuration
st.set_page_config(
    page_title="PDF RAG System",
    page_icon="📚",
    layout="wide"
)

# Title and description
st.title("📚 PDF RAG System")
st.markdown("""
This application allows you to upload a PDF file, ask questions about its content, and get AI-generated answers based on the document.
""")

# File upload section
st.header("1. Upload PDF")
uploaded_file = st.file_uploader("Choose a PDF file", type="pdf", key="pdf_uploader")

# Initialize session state variables
if 'pdf_processed' not in st.session_state:
    st.session_state.pdf_processed = False
if 'index' not in st.session_state:
    st.session_state.index = None
if 'chunks' not in st.session_state:
    st.session_state.chunks = None
if 'pdf_path' not in st.session_state:
    st.session_state.pdf_path = None

# Process the uploaded PDF
if uploaded_file is not None and not st.session_state.pdf_processed:
    with st.spinner("Processing PDF..."):
        # Create a temporary file to save the uploaded PDF
        with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
            tmp_file.write(uploaded_file.getvalue())
            st.session_state.pdf_path = tmp_file.name
        
        # Extract text from PDF
        pdf_text = extract_text_from_pdf(st.session_state.pdf_path)
        
        # Chunk the text
        chunks = chunk_text(pdf_text, chunk_size=1000, chunk_overlap=100)
        st.session_state.chunks = chunks
        
        # Create embeddings
        embeddings = embedding_function(chunks)
        
        # Convert embeddings to numpy array if they aren't already
        if not isinstance(embeddings, np.ndarray):
            embeddings = np.array(embeddings).astype('float32')
        
        # Get the dimension of the embeddings
        dimension = embeddings.shape[1]
        
        # Initialize FAISS index
        index = faiss.IndexFlatL2(dimension)
        
        # Add vectors to the index
        index.add(embeddings)
        
        # Save the index and chunks
        faiss.write_index(index, "./faiss_index")
        with open("./document_chunks.pkl", 'wb') as f:
            pickle.dump(chunks, f)
        
        # Update session state
        st.session_state.index = index
        st.session_state.pdf_processed = True
        
        st.success(f"PDF processed successfully! {len(chunks)} chunks created.")

# Query section
st.header("2. Ask a Question")
query = st.text_input("Enter your question about the PDF content:", key="query_input")

# Add a button to submit the query
if st.button("Get Answer", key="get_answer_button") and query and st.session_state.pdf_processed:
    with st.spinner("Retrieving relevant information and generating answer..."):
        try:
            # Generate embedding for the query
            query_embedding = embedding_model.encode([query], convert_to_numpy=True).astype('float32')
            
            # Search the index
            n_results = 5
            distances, indices = st.session_state.index.search(query_embedding, n_results)
            
            # Get the documents
            documents = [st.session_state.chunks[i] for i in indices[0]]
            
            # Convert distances to similarity scores (L2 distance: lower is better)
            # Normalize distances to [0, 1] range where 1 is most similar
            max_distance = np.max(distances)
            similarity_scores = [1 - (dist / max_distance) for dist in distances[0]]
            
            # Create context from retrieved documents
            context = (documents, similarity_scores)
            
            # Query the LLM with context
            answer = query_llm_with_context(query, context, top_n=3)
            
            # Display the answer
            st.header("3. Answer")
            st.write(answer)
            
            # Display the retrieved documents
            with st.expander("View Retrieved Documents", expanded=False):
                for i, (doc, score) in enumerate(zip(documents, similarity_scores)):
                    st.markdown(f"**Document {i+1}** (Relevance: {score:.4f})")
                    st.text(doc[:500] + "..." if len(doc) > 500 else doc)
                    st.markdown("---")
        
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
            logger.exception("Error during query processing")

# Add a reset button
if st.button("Reset and Upload New PDF", key="reset_button"):
    # Clean up temporary files
    if st.session_state.pdf_path and os.path.exists(st.session_state.pdf_path):
        os.unlink(st.session_state.pdf_path)
    
    # Reset session state
    st.session_state.pdf_processed = False
    st.session_state.index = None
    st.session_state.chunks = None
    st.session_state.pdf_path = None
    
    # Reload the page
    st.experimental_rerun()

# Footer
st.markdown("---")
st.markdown("Built with Streamlit, FAISS, and Hugging Face API")