import streamlit as st
import os
from utils.pdf_utils import PDFProcessor
from utils.embeddings_utils import EmbeddingsManager
from utils.qa_utils import QASystem
from dotenv import load_dotenv
import openai

def initialize_session_state():
    if 'pdf_processor' not in st.session_state:
        st.session_state['pdf_processor'] = None
    if 'embeddings_manager' not in st.session_state:
        st.session_state['embeddings_manager'] = None
    if 'qa_system' not in st.session_state:
        st.session_state['qa_system'] = None
    if 'processed_pdfs' not in st.session_state:
        st.session_state['processed_pdfs'] = set()
    if 'all_text_chunks' not in st.session_state:
        st.session_state['all_text_chunks'] = []

def main():
    load_dotenv()
    st.set_page_config(page_title="AI-Powered PDF Assistant", layout="wide")

    initialize_session_state()

    # Header Section
    st.markdown(
        """
        <style>
        .main-header {
            font-size: 2.5rem;
            color: #1F77B4;
            text-align: center;
            margin-bottom: 1rem;
        }
        .sub-header {
            font-size: 1.25rem;
            color: #555;
            text-align: center;
            margin-bottom: 2rem;
        }
        </style>
        <div class="main-header">📘 AI-Powered PDF Assistant</div>
        <div class="sub-header">Upload, Analyze, and Interact with Your Documents</div>
        """,
        unsafe_allow_html=True
    )

    # Navigation Menu
    selected_page = st.sidebar.radio(
        "Navigate", ["Upload PDFs", "Ask Questions", "About"]
    )

    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        st.sidebar.error("OpenAI API key not found in .env file!")
        return

    openai.api_key = api_key

    if not st.session_state['pdf_processor']:
        st.session_state['pdf_processor'] = PDFProcessor()
    if not st.session_state['embeddings_manager']:
        st.session_state['embeddings_manager'] = EmbeddingsManager(api_key)
    if not st.session_state['qa_system']:
        st.session_state['qa_system'] = QASystem(api_key)

    if selected_page == "Upload PDFs":
        st.header("📤 Upload PDFs")
        st.markdown(
            """<p style='font-size: 1.1rem;'>Drag and drop your PDF files below to extract and process content for analysis.</p>""",
            unsafe_allow_html=True
        )

        uploaded_files = st.file_uploader(
            "Upload PDF files", type=['pdf'], accept_multiple_files=True
        )

        if uploaded_files:
            new_files = [f for f in uploaded_files if f.name not in st.session_state['processed_pdfs']]
            if new_files:
                with st.spinner("Processing PDFs..."):
                    for pdf_file in new_files:
                        try:
                            pages = st.session_state['pdf_processor'].extract_text(pdf_file)
                            for page_text in pages.values():
                                chunks = st.session_state['pdf_processor'].chunk_text(page_text)
                                st.session_state['all_text_chunks'].extend(chunks)
                            st.session_state['processed_pdfs'].add(pdf_file.name)
                        except Exception as e:
                            st.error(f"Error processing {pdf_file.name}: {str(e)}")
                            continue

                    with st.spinner("Generating embeddings..."):
                        try:
                            st.session_state['embeddings_manager'].generate_embeddings(
                                st.session_state['all_text_chunks']
                            )
                            st.success("✅ Documents processed successfully!")
                        except Exception as e:
                            st.error(f"Error generating embeddings: {str(e)}")

    elif selected_page == "Ask Questions":
        st.header("❓ Ask Questions")
        st.markdown(
            """<p style='font-size: 1.1rem;'>Query your uploaded documents and get precise answers backed by AI-powered analysis.</p>""",
            unsafe_allow_html=True
        )

        if st.session_state['all_text_chunks']:
            question = st.text_input("Enter your question:")

            if question:
                try:
                    with st.spinner("Finding relevant information..."):
                        relevant_chunks = st.session_state['embeddings_manager'].find_relevant_chunks(
                            question, k=3
                        )
                        answer = st.session_state['qa_system'].generate_answer(
                            question, relevant_chunks
                        )

                        st.markdown("### 🤖 Answer")
                        st.write(answer)

                        with st.expander("🔍 View Source Context"):
                            for i, chunk in enumerate(relevant_chunks, 1):
                                st.markdown(f"**Context {i}:**")
                                st.write(chunk)
                                st.markdown("---")
                except openai.error.RateLimitError:
                    st.error("Rate limit exceeded. Please try again later.")
                except Exception as e:
                    st.error(f"Error: {str(e)}")
        else:
            st.warning("Please upload and process documents in the 'Upload PDFs' section first.")

    elif selected_page == "About":
        st.header("ℹ️ About This App")
        st.markdown(
            """
            <p style='font-size: 1.1rem;'>
            <b>AI-Powered PDF Assistant</b> is a smart solution for extracting and querying information from PDF files. With powerful AI integrations, 
            this tool allows seamless document analysis and interaction.
            </p>

            <h3>🔑 Key Features</h3>
            <ul>
                <li>Upload and process multiple PDF files</li>
                <li>Generate embeddings for precise content retrieval</li>
                <li>Query documents and receive context-aware answers</li>
            </ul>

            <h3>🛠️ Technologies Used</h3>
            <ul>
                <li>Streamlit for interactive UI</li>
                <li>OpenAI GPT API for Q&A</li>
                <li>Custom PDF processing and embedding tools</li>
            </ul>

            <p style='text-align: center;'>
            Built with ❤️ by [Your Name]
            </p>
            """,
            unsafe_allow_html=True
        )

if __name__ == "__main__":
    main()