import streamlit as st
import random
from langchain_components.replier import get_context_from_vectorstore,get_vectorstore_from_postgres,prepare_prompt_and_chain_with_history,get_vectorstore_from_pinecone
import fitz

def display_pdf(pdf_path):
    try:
        pdf_document = fitz.open(pdf_path)
        num_pages = pdf_document.page_count
        st.sidebar.write(f"Total pages: {num_pages}")

        for page_num in range(num_pages):
            page = pdf_document.load_page(page_num)
            image = page.get_pixmap()
            st.sidebar.image(image.tobytes(), caption=f"Page {page_num + 1}", use_column_width=True)

    except Exception as e:
        st.sidebar.error(f"Error loading PDF: {e}")



def main():
    st.header('Interact with your PDF that includes images, tables, and graphs.')
    
    if "activate_chat" not in st.session_state:
        st.session_state.activate_chat = False

    if "messages" not in st.session_state:
        st.session_state.messages = []


    with st.sidebar:
        username = st.text_input("Please enter your name here")
        if st.button('Press Button to Start chat with your pdf...'):
            if "user_id" not in st.session_state:
                st.session_state.user_id = username
            
            if "session_id" not in st.session_state:
                random_number = random.randint(1, 1000000)
                st.session_state.session_id = str(random_number)

            if "vectorstore" not in st.session_state:
                collection_name="fy2024_chunk_2000"
                pinecone_collection_name="fy2024"
                #st.session_state.vectorstore = get_vectorstore_from_postgres(collection_name)
                st.session_state.vectorstore = get_vectorstore_from_pinecone(pinecone_collection_name)

            if "chain" not in st.session_state:
                st.session_state.chain = prepare_prompt_and_chain_with_history()

            st.session_state.activate_chat = True


        st.subheader("PDF Viewer")
        pdf_path = "fy2024.pdf"
        if st.button('Show PDF'):
            st.session_state.pdf_path = pdf_path

        if st.download_button(label="Download PDF", data=open(pdf_path, 'rb').read(), file_name=pdf_path.split("/")[-1]):
            pass

    if "pdf_path" in st.session_state:
        pdf_path = st.session_state.pdf_path
        display_pdf(pdf_path)



    for message in st.session_state.messages:
        with st.chat_message(message["role"], avatar = message['avatar']):
            st.markdown(message["content"])

    if st.session_state.activate_chat == True:
        if prompt := st.chat_input("Ask your question from the PDF? "):
            with st.chat_message("user", avatar = '👨🏻'):
                st.markdown(prompt)
            st.session_state.messages.append({"role": "user",  "avatar" :'👨🏻', "content": prompt})
            
            user_id = st.session_state.user_id
            session_id = st.session_state.session_id
            vectorstore = st.session_state.vectorstore
            chain = st.session_state.chain
            print("chain Done")

            data=get_context_from_vectorstore(vectorstore,prompt)
            ai_msg =chain.invoke({"data": data, "input": prompt}, config={"configurable": {"user_id": user_id, "session_id": session_id}})
            cleaned_response=ai_msg.content
            with st.chat_message("assistant", avatar='🤖'):
                st.markdown(cleaned_response)
            st.session_state.messages.append({"role": "assistant",  "avatar" :'🤖', "content": cleaned_response})


if __name__ == '__main__':
    main()