import re
import time
import io
from io import StringIO
from typing import Any, Dict, List
#Modules to Import
import openai
import streamlit as st
from langchain import LLMChain, OpenAI
from langchain.agents import AgentExecutor, Tool, ZeroShotAgent
from langchain.chains import RetrievalQA
from langchain.chains.question_answering import load_qa_chain
from langchain.docstore.document import Document
from langchain.document_loaders import PyPDFLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.memory import ConversationBufferMemory
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import VectorStore
from langchain.vectorstores.faiss import FAISS
from pypdf import PdfReader

@st.cache_data
def parse_pdf (file: io.BytesIO)-> List[str]:
    pdf = PdfReader(file)
    output = []
    for page in pdf.pages:

        text = page.extract_text()
        #Merge hyphenated words
        text = re.sub(r"(\w+)-\n(\w+)", "\1\2", text)
        # Fix newlines in the middle of sentences 
        text = re.sub(r"(?<!\n\s)\n(?!\s\n)", " ", text.strip())
        #Remove multiple newlines
        text = re.sub(r"\n\s*\n", "\n\n", text)
        
        output.append(text)
    return output

@st.cache_data
def text_to_docs(text: str) -> List [Document]:

    """Converts a string or list of strings to a list of Documents with metadata,"""

    if isinstance(text, str):
        #Take a single string as one page 
        text = [text]
    page_docs = [Document (page_content=page) for page in text]
    # Add page numbers as metadata 
    for i, doc in enumerate(page_docs): 

        doc.metadata["page"] = 1 + 1
    # Split pages into chunks 
    doc_chunks = []
    for doc in page_docs:
        text_splitter = RecursiveCharacterTextSplitter( 
            chunk_size=4000, 
            separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""],
            chunk_overlap=0,
        )
        chunks = text_splitter.split_text(doc.page_content)
        for i, chunk in enumerate(chunks):
            doc = Document(
                page_content=chunk, metadata={"page": doc.metadata["page"], "chunk": 1}
            )
            # Add sources a metadata
            doc.metadata["source"] = f"{doc.metadata['page']}-{doc.metadata['chunk']}" 
            doc_chunks.append(doc)
    return doc_chunks

 

uploaded_file = st.sidebar.file_uploader(":blue[Upload]", type=["pdf"])
if uploaded_file:

    doc = parse_pdf(uploaded_file)

    pages = text_to_docs(doc)
    # pages
    if pages:
        with st.expander('Show page contents', expanded=False):
            page_sel =st.number_input(
                label="selected page", min_value=1, max_value=len(pages), step=1
            )
            st.write(pages[page_sel-1])
            api = st.sidebar.text_input(
                "Open api key",
                type="password",
                placeholder="sk-",
                help="https://platform.openai.com/account/api-keys",
            )
            if api:
                embeddings = OpenAIEmbeddings(openai_api_key = api)
                # Indexing
                # Save in a Vector DB_
                with st.spinner("It's indexing. .."):

                    index = FAISS.from_documents(pages, embeddings)

                qa = RetrievalQA.from_chain_type(
                    llm = OpenAI(openai_api_key = api),
                    chain_type = "stuff",
                    retriever = index.as_retriever()
                )

                # our tool
                tools = [
                    Tool(
                        name="State of Union QA System",
                        func=qa.run,
                        description="Useful for when you need to answer questions about the aspects asked. Input may be a partial or fully formed question."
                    )
                ]
                prefix=""""Have a conversation with a human, answering the following questions as best you can based on the context and memory available. 
                        You have access to a single tool:"""
                suffix="""Begin!"
                {chat_history}
                Question: {input}
                {agent_scratchpad}"""
                prompt = ZeroShotAgent.create_prompt(
                    tools,
                    prefix=prefix,
                    suffix=suffix,
                    input_variables=["input", "chat_history", "agent_scratchpad"],
                )

                if "memory" not in st.session_state:
                    st.session_state.memory = ConversationBufferMemory(memory_key ="chat_history")

                #Chain
                # ZeroShotAgent

                llm_chain = LLMChain(
                    llm=OpenAI(
                    temperature=0, openai_api_key=api, model_name="gpt-3.5-turbo"
                    ),
                    prompt=prompt,
                )
                agent = ZeroShotAgent (llm_chain=llm_chain, tools=tools, verbose=True) 
                agent_chain = AgentExecutor.from_agent_and_tools(
                    agent=agent, tools=tools, verbose=True, memory=st.session_state.memory
                )
container = st.container()
with container:
    st.title("🤖 AI ChatBot")
                
# Initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages = []
# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])    
              
if query := st.chat_input("Hey yo !!! Wazzups!"):
       
    
    st.chat_message("user").markdown(query)
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": query})
   
    # response=llm_chain.memory.chat_memory.add_user_message(prompt)
    with st.spinner("It's indexing. .."):
        response = agent_chain.run(query)
    # st.write(response)
     # #f"Echo: {prompt}" get_completion(template_string) #
    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        st.markdown(response)
    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": response})
# with st.expander("History/Memory"):
# st.write(st.session_state.memory)