import streamlit as st
import google.generativeai as genai
from langchain.document_loaders import PyPDFDirectoryLoader
import os
import shutil

# Configuration
GOOGLE_API_KEY = st.secrets["GOOGLE_API_KEY"]

# Page configuration
st.set_page_config(page_title="Chat with PDFs", page_icon="📚")

def initialize_session_state():
    """Initialize session state variables"""
    session_state_vars = {
        "messages": [],
        "loaded_files": False,
        "pdf_content": None,
        "chat": None
    }
    
    for var, value in session_state_vars.items():
        if var not in st.session_state:
            st.session_state[var] = value

def load_pdfs(pdf_folder):
    """Load PDFs and return their content"""
    if not os.path.exists(pdf_folder):
        os.makedirs(pdf_folder)
        
    loader = PyPDFDirectoryLoader(pdf_folder)
    documents = loader.load()
    
    # Concatenate all documents content
    content = "\n\n".join([doc.page_content for doc in documents])
    return content

def initialize_chat(pdf_content):
    """Initialize Gemini chat with PDF content"""
    genai.configure(api_key=GOOGLE_API_KEY)
    
    generation_config = {
        "temperature": 0.7,
        "top_p": 0.95,
        "top_k": 40,
        "max_output_tokens": 8192,
    }
    
    model = genai.GenerativeModel(
        model_name="gemini-1.5-pro",
        generation_config=generation_config,
    )
    
    # Start chat with context
    context_prompt = f"""You are a helpful assistant that answers questions based on the following document content:

{pdf_content}

Please use this content to answer user questions. If the answer cannot be found in the content, say so."""

    chat = model.start_chat(history=[])
    # Send initial context
    chat.send_message(context_prompt)
    return chat

def main():
    initialize_session_state()
    
    st.title("💬 Chat with PDFs")
    
    # Sidebar for PDF upload
    with st.sidebar:
        st.header("Upload Documents")
        uploaded_files = st.file_uploader(
            "Upload your PDFs",
            type=["pdf"],
            accept_multiple_files=True
        )

        if uploaded_files and not st.session_state.loaded_files:
            # Create pdfs directory if it doesn't exist
            if not os.path.exists("pdfs"):
                os.makedirs("pdfs")
            
            # Clean up old PDF files
            for file in os.listdir("pdfs"):
                os.remove(os.path.join("pdfs", file))
            
            # Save uploaded files
            for file in uploaded_files:
                with open(f"pdfs/{file.name}", "wb") as f:
                    f.write(file.getvalue())
            
            # Load PDF content
            with st.spinner("Processing PDFs..."):
                try:
                    pdf_content = load_pdfs("pdfs")
                    st.session_state.pdf_content = pdf_content
                    st.session_state.loaded_files = True
                    
                    # Initialize chat with content
                    st.session_state.chat = initialize_chat(pdf_content)
                except Exception as e:
                    st.error(f"Error processing PDFs: {str(e)}")
                    return

    # Main chat interface
    if st.session_state.loaded_files:
        # Display chat messages
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

        # Chat input
        if prompt := st.chat_input("Ask a question about your PDFs:"):
            # Add user message to chat history
            st.session_state.messages.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.markdown(prompt)

            with st.chat_message("assistant"):
                response_placeholder = st.empty()
                try:
                    # Get response from Gemini
                    if not st.session_state.chat:
                        st.session_state.chat = initialize_chat(st.session_state.pdf_content)
                    
                    response = st.session_state.chat.send_message(prompt)
                    response_text = response.text
                    
                    response_placeholder.markdown(response_text)
                    
                    # Add assistant response to chat history
                    st.session_state.messages.append({"role": "assistant", "content": response_text})
                except Exception as e:
                    response_placeholder.error(f"Error generating response: {str(e)}")

    else:
        st.info("Please upload PDFs to start chatting.")

if __name__ == "__main__":
    main()