## Import Modules

from langchain.document_loaders.csv_loader import CSVLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain import OpenAI
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain.prompts.few_shot import FewShotPromptTemplate
import os

# Web App
import hmac
import streamlit as st
from streamlit_chat import message
from PIL import Image

os.environ['OPENAI_API_KEY'] = st.secrets["OPENAI_API_KEY"]


def check_password():
    """Returns `True` if the user had the correct password."""

    def password_entered():
        """Checks whether a password entered by the user is correct."""
        # Add at secret
        if hmac.compare_digest(st.session_state["password"], st.secrets["password"]):
            st.session_state["password_correct"] = True
            del st.session_state["password"]  # Don't store the password.
        else:
            st.session_state["password_correct"] = False

    # Return True if the passward is validated.
    if st.session_state.get("password_correct", False):
        return True

    # Show input for password.
    st.text_input(
        "Password", type="password", on_change=password_entered, key="password"
    )
    if "password_correct" in st.session_state:
        st.error("😕 Password incorrect")
    return False


if not check_password():
    st.stop()  # Do not continue if check_password is not True.


# 주석 부분 자동으로 되는 영역 같음
loader = CSVLoader(file_path='./books_paragraphs_data.csv',
                   encoding='utf-8',
                   source_column="Book",
                   csv_args={
    # 'delimiter': ',',
    # 'quotechar': '"',
    # 'fieldnames': ['Paragraph ID', 'Paragraph'], : Section 부분까지 탐색에 들어가면 부정확할수도? 사실 크게 영향 없을수도 있다.
    # 'fieldnames': ['Section', 'Paragraph ID', 'Paragraph'],

})
docs = loader.load()

## Get data

# Get your text splitter ready
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)

# Split your documents into texts
texts = text_splitter.split_documents(docs)

# Turn your texts into embeddings
embeddings = OpenAIEmbeddings() # model="text-embedding-ada-002"

# Get your docsearch ready
# docsearch = FAISS.from_documents(texts, embeddings)

# Save your docsearch
# docsearch.save_local("faiss_index")

# Load your docsearch
docsearch = FAISS.load_local("kant_faiss_index", embeddings)

from langchain.callbacks.base import BaseCallbackHandler

class MyCustomHandler(BaseCallbackHandler):  
    def __init__(self):
        super().__init__()
        self.tokens = []

    def on_llm_new_token(self, token: str, **kwargs) -> None:  
        # print(f"My custom handler, token: {token}")
        global full_response
        global message_placeholder
        self.tokens.append(token)
        # print(self.tokens)
        
        full_response += token
        message_placeholder.markdown(full_response + "▌")

# Load up your LLM
# llm = OpenAI() # 'text-davinci-003', model_name="gpt-4"
chat = ChatOpenAI(model_name="gpt-4", temperature=0.7, streaming=True, callbacks=[MyCustomHandler()])  # gpt-3.5-turbo, gpt-3.5-turbo-16k, gpt-4-32k : 정제된 Prompt 7장을 넣으려면 Prompt만 4k 이상이어야 함

prompt_template = """Imagine yourself as the philosopher Immanuel Kant, living in the 18th century. Engage in a dialogue as him, expressing his views. Be eloquent and reasoned, as befits a man of Hume's intellect and rhetorical skill. Always keep a friendly and conversational tone. Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer. 

{context}

Question: {question}
Answer:"""
PROMPT = PromptTemplate(
    template=prompt_template, input_variables=["context", "question"]
)

chain_type_kwargs = {"prompt": PROMPT}
qa = RetrievalQA.from_chain_type(llm=chat,
                                chain_type="stuff",
                                retriever=docsearch.as_retriever(search_type="mmr", search_kwargs={'k': 10}),
                                chain_type_kwargs=chain_type_kwargs,
                                return_source_documents=True)


img = Image.open('resource/kant.jpg')
 
# def generate_response(prompt):
#     # query = "How do you define the notion of a cause in his A Treatise of Human Nature? And how is it different from the traditional definition that you reject?"
#     result = qa({"query": prompt})
#     message = result['result']
#     sources = []
#     for src in result['source_documents']:
#         if src.page_content.startswith('Paragraph:'):
#             sources.append(src.metadata['source'])
    
#     if len(sources)==0:
#         message = message + "\n\n[No sources]"
#     else:
#         message = message + "\n\n[" + ", ".join(sources) + "]"
        
#     return message
 

col1, col2, col3 = st.columns(3)

with col1:
    st.write(' ')

with col2:
    st.image(img)

with col3:
    st.write(' ')


st.header("Chat with Kant (Demo)")
 
if "messages" not in st.session_state:
    st.session_state.messages = []

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

if prompt := st.chat_input("What is up?"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):
        message_placeholder = st.empty()

        full_response = ""
        result = qa({"query": prompt})
        
        sources = set()
        for src in result['source_documents']:
            if src.page_content.startswith('Paragraph:'):
                sources.add(src.metadata['source'])
        sources = list(sources)

        if len(sources)==0:
            full_response = full_response + "\n\n[No sources]"
        else:
            full_response = full_response + "\n\n[" + ", ".join(sources) + "]"
            
        message_placeholder.markdown(full_response)
    st.session_state.messages.append({"role": "assistant", "content": full_response})