import streamlit as st
import PyPDF2
import io
from typing import List, Dict
from groq import Groq

# Page configuration and styling
st.set_page_config(
    page_title="PDF AI Chat Assistant", 
    page_icon="📄", 
    layout="wide"
)

hide_menu_style = """
        <style>
        #MainMenu {visibility: hidden;}
        footer {visibility: hidden;}
        </style>
        """
st.markdown(hide_menu_style, unsafe_allow_html=True)

class PDFChatAssistant:
    def __init__(self):
        # Initialize session state variables
        if 'chat_history' not in st.session_state:
            st.session_state.chat_history = []
        if 'uploaded_pdf' not in st.session_state:
            st.session_state.uploaded_pdf = None
        if 'pdf_text' not in st.session_state:
            st.session_state.pdf_text = ""

    def extract_pdf_text(self, pdf_file) -> str:
        """
        Extract text from uploaded PDF file
        """
        try:
            pdf_reader = PyPDF2.PdfReader(pdf_file)
            text = ""
            for page in pdf_reader.pages:
                text += page.extract_text() + "\n\n"
            return text
        except Exception as e:
            st.error(f"Error extracting PDF text: {e}")
            return ""

    def create_groq_client(self) -> Groq:
        """
        Create Groq client with API key
        """
        client = Groq(api_key=st.secrets["GROQ_API_KEY_DOCUMENTOR"])
        return client

    def generate_ai_response(self, client: Groq, context: str, user_query: str) -> str:
        """
        Generate AI response using Groq API
        """
        try:
            chat_completion = client.chat.completions.create(
                messages=[
                    {
                        "role": "system", 
                        "content": f"Bạn là một trợ lý PDF hữu ích. Luôn trả lời bằng tiếng Việt. Sử dụng nội dung trong file PDF đính kèm để trả lời câu hỏi sau:\n\n{context}"
                    },
                    {
                        "role": "user", 
                        "content": user_query
                    }
                ],
                model="llama3-8b-8192"
            )
            return chat_completion.choices[0].message.content
        except Exception as e:
            st.error(f"Error generating AI response: {e}")
            return "Sorry, I couldn't generate a response."

    def display_chat_interface(self):
        """
        Main chat interface with PDF upload and interaction
        """
        st.title("📄 PDF AI Chat Assistant")

        # Sidebar for PDF upload and context management
        with st.sidebar:
            st.header("PDF Management")
            uploaded_file = st.file_uploader("Upload PDF", type=['pdf'])
            
            if uploaded_file is not None:
                st.session_state.uploaded_pdf = uploaded_file
                with st.spinner('Extracting text from PDF...'):
                    st.session_state.pdf_text = self.extract_pdf_text(uploaded_file)
                    st.success("PDF text extracted successfully!")
                
                # Context actions
                col1, col2 = st.columns(2)
                with col1:
                    if st.button("Summarize PDF", key="summarize_btn"):
                        st.session_state.chat_history.append({
                            "role": "user", 
                            "content": "Provide a comprehensive summary of this PDF"
                        })

                with col2:
                    if st.button("Key Points", key="key_points_btn"):
                        st.session_state.chat_history.append({
                            "role": "user", 
                            "content": "Extract the most important key points from this PDF"
                        })

        # Main chat area
        st.header("Chat with Your PDF")
        
        # Display chat history
        for message in st.session_state.chat_history:
            with st.chat_message(message['role']):
                st.markdown(message['content'])

        # Groq Client setup
        client = self.create_groq_client()
        
        # User input
        if prompt := st.chat_input("Ask a question about your PDF"):
            # Add user message to chat history
            st.session_state.chat_history.append({
                "role": "user", 
                "content": prompt
            })

            # Display user message
            with st.chat_message("user"):
                st.markdown(prompt)

            # Generate and display AI response
            if client and st.session_state.pdf_text:
                with st.chat_message("assistant"):
                    with st.spinner('Generating response...'):
                        response = self.generate_ai_response(
                            client, 
                            st.session_state.pdf_text, 
                            prompt
                        )
                        st.markdown(response)
                
                # Add AI response to chat history
                st.session_state.chat_history.append({
                    "role": "assistant", 
                    "content": response
                })
            else:
                st.warning("Please upload a PDF.")

        # Reset button
        if st.button("Clear Chat", key="reset_btn"):
            st.session_state.chat_history = []
            st.experimental_rerun()

def main():
    assistant = PDFChatAssistant()
    assistant.display_chat_interface()

if __name__ == "__main__":
    main()