import streamlit as st
import os
import shutil
from SimpleRAG import read_doc, chunk_data, retrieve_answers
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Initialize session state
if 'docs_processed' not in st.session_state:
    st.session_state['docs_processed'] = False

# Set page config
st.set_page_config(
    page_title="Document Q&A System",
    page_icon="📚",
    layout="wide"
)

# Title and description
st.title("📚 Document Question & Answer System")
st.markdown("""
This application allows you to upload documents and ask questions about their content.
The system uses advanced RAG (Retrieval Augmented Generation) to provide accurate answers.
""")

# Check for required environment variables
if not os.environ.get('OPENAI_API_KEY'):
    st.error("⚠️ OPENAI_API_KEY is not set in the environment variables!")
    st.stop()

# Sidebar for document upload
with st.sidebar:
    st.header("Document Upload")
    uploaded_files = st.file_uploader(
        "Upload your documents (DOCX format)",
        type=['docx'],
        accept_multiple_files=True
    )
    
    if uploaded_files:
        # Create/clear documents directory
        if os.path.exists('documents'):
            shutil.rmtree('documents')
        os.makedirs('documents')
            
        # Save uploaded files
        for uploaded_file in uploaded_files:
            try:
                with open(os.path.join('documents', uploaded_file.name), 'wb') as f:
                    f.write(uploaded_file.getbuffer())
                st.success(f"✅ Successfully uploaded: {uploaded_file.name}")
            except Exception as e:
                st.error(f"❌ Error uploading {uploaded_file.name}: {str(e)}")
        
        if st.button("Process Documents"):
            try:
                with st.spinner("Processing documents..."):
                    # Read and process documents
                    documents = read_doc('documents/')
                    if not documents:
                        st.error("❌ No valid documents found in the uploaded files.")
                        st.session_state['docs_processed'] = False
                    else:
                        chunks = chunk_data(documents)
                        st.session_state['docs_processed'] = True
                        st.success(f"✅ Successfully processed {len(documents)} documents into {len(chunks)} chunks!")
            except Exception as e:
                st.error(f"❌ Error processing documents: {str(e)}")
                st.session_state['docs_processed'] = False

# Main content area
st.header("Ask Questions")

# Input for user question
user_question = st.text_input("Enter your question about the documents:", key="user_question")

# Process question
if user_question:
    if st.session_state.get('docs_processed', False):
        try:
            with st.spinner("Finding answer..."):
                answer = retrieve_answers(user_question)
                
                # Display answer in a nice format
                st.markdown("### Answer")
                st.write(answer)
        except Exception as e:
            st.error(f"❌ Error generating answer: {str(e)}")
    else:
        st.warning("⚠️ Please upload and process documents first!")