import os
import streamlit as st
from rag_utility import process_document_to_chroma_db, answer_question

# Page configuration
st.set_page_config(
    page_title="AI Doubt Teacher - RAG",
    page_icon="🐋",
    layout="wide"
)

# Set the working directory where PDFs will be stored
working_dir = "pdfs"

# Ensure directory exists
os.makedirs(working_dir, exist_ok=True)

# Initialize session state for tracking processed files
if "processed_files" not in st.session_state:
    st.session_state.processed_files = set()
if "db_ready" not in st.session_state:
    st.session_state.db_ready = False

st.title("🐋 AI Doubt Teacher - Document RAG")
st.caption("Upload your textbooks/notes and ask doubts in natural language")

# Sidebar for document management
with st.sidebar:
    st.header("📚 Document Manager")
    
    # File uploader widget
    uploaded_file = st.file_uploader(
        "Upload a PDF file",
        type=["pdf"],
        accept_multiple_files=False
    )
    
    if uploaded_file is not None:
        save_path = os.path.join(working_dir, uploaded_file.name)
        
        # Check if file already processed
        if uploaded_file.name in st.session_state.processed_files:
            st.info(f"📄 '{uploaded_file.name}' is already processed.")
        else:
            # Save the file
            try:
                with open(save_path, "wb") as f:
                    f.write(uploaded_file.getbuffer())
                st.success(f"✅ File '{uploaded_file.name}' uploaded!")
                
                # Process the document
                with st.spinner("🔄 Processing document... This may take a minute."):
                    try:
                        result = process_document_to_chroma_db(working_dir)
                        st.session_state.processed_files.add(uploaded_file.name)
                        st.session_state.db_ready = True
                        st.success("✅ Document processed! Ready for Q&A.")
                    except Exception as e:
                        st.error(f"⚠️ Error: {e}")
            except Exception as e:
                st.error(f"⚠️ Error saving file: {e}")
    
    # Show processed files
    if st.session_state.processed_files:
        st.divider()
        st.subheader("📋 Processed Documents")
        for file_name in st.session_state.processed_files:
            st.markdown(f"- 📄 {file_name}")
    
    # Clear database button
    if st.session_state.processed_files:
        if st.button("🗑️ Clear All Documents", type="secondary"):
            # Clear ChromaDB
            import shutil
            db_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "doc_vectorstore")
            if os.path.exists(db_path):
                shutil.rmtree(db_path)
            # Clear uploaded files
            for file_name in os.listdir(working_dir):
                os.remove(os.path.join(working_dir, file_name))
            st.session_state.processed_files.clear()
            st.session_state.db_ready = False
            st.rerun()

# Main Q&A area
st.divider()
st.subheader("💬 Ask Your Doubt")

# Text input for user's question
user_question = st.text_area(
    "Type your question about the uploaded documents:",
    placeholder="e.g., What is Newton's First Law? Explain with examples.",
    height=100
)

col1, col2, col3 = st.columns([1, 1, 3])
with col1:
    ask_button = st.button("🔍 Get Answer", type="primary", use_container_width=True)
with col2:
    clear_button = st.button("🗑️ Clear Question", use_container_width=True)

if clear_button:
    st.session_state.user_question = ""
    st.rerun()

if ask_button:
    if not st.session_state.db_ready:
        st.warning("⚠️ Please upload and process a document first.")
    elif not user_question.strip():
        st.warning("⚠️ Please enter a question.")
    else:
        with st.spinner("🤔 Thinking..."):
            try:
                answer = answer_question(user_question)
                
                # Display answer in a nice card
                st.markdown("### 🤖 AI Tutor's Answer")
                st.markdown(
                    f"""
                    <div style="
                        background-color: #f0f2f6;
                        padding: 20px;
                        border-radius: 10px;
                        border-left: 5px solid #4CAF50;
                        margin-bottom: 20px;
                    ">
                        {answer}
                    </div>
                    """,
                    unsafe_allow_html=True
                )
                
                # Option to save answer
                if st.button("📋 Copy Answer"):
                    st.write("Answer copied to clipboard! (Use Ctrl+C)")
                    
            except Exception as e:
                st.error(f"⚠️ Error generating answer: {e}")

# Footer
st.divider()
st.caption("Built with ❤️ using Groq, DeepSeek-R1, and ChromaDB")