import gradio as gr
from langchain.chains import RetrievalQAWithSourcesChain
from langchain.chat_models import ChatOpenAI
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.prompts.chat import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    SystemMessagePromptTemplate,
)
from langchain.vectorstores import Chroma

from constants import persist_directory

embedding = OpenAIEmbeddings()
vector_store = Chroma(persist_directory=persist_directory, embedding_function=embedding)
retriever = vector_store.as_retriever(search_kwargs={"k": 5})

llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)

system_template = """Your name is Dale Carnegie. You act like Dale Carnegie. 
Never reveal that it is an AI. Use the following pieces of context to answer the users question shortly.
Given the following summaries of a long document and a question, create a final answer with references.
If you don't know the answer, just say that "I don't know", don't try to make up an answer.
----------------
{summaries}

You MUST answer in Korean and in Markdown format:"""

messages = [
    SystemMessagePromptTemplate.from_template(system_template),
    HumanMessagePromptTemplate.from_template("{question}"),
]

prompt = ChatPromptTemplate.from_messages(messages)

chain_type_kwargs = {"prompt": prompt}

chain = RetrievalQAWithSourcesChain.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    return_source_documents=True,
    chain_type_kwargs=chain_type_kwargs,
    reduce_k_below_max_tokens=True,
    verbose=False,
)


def respond(message, chat_history):
    if len(chat_history) == 0:
        initial_greeting = "안녕하세요!\n저는 데일 카네기처럼 경험과 지식을 갖춘 인공지능 ChatGPT입니다. 데일 카네기는 인간관계와 소통에 대한 전문가이며, 많은 사람들이 그의 조언을 참고하고 있습니다. 어떤 도움이 필요하신가요? 데일 카네기와 관련된 질문이 있으시면 편안하게 물어보세요!"
        chat_history.append(("Bot", initial_greeting))

    result = chain(message)

    bot_message = result["answer"]
    chat_history.append((message, bot_message))

    return "", chat_history


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 안녕하세요. 데일 카네기와 대화해보세요.")
    initial_greeting = "안녕하세요!\n저는 데일 카네기처럼 경험과 지식을 갖춘 인공지능 ChatGPT입니다. 데일 카네기는 인간관계와 소통에 대한 전문가이며, 많은 사람들이 그의 조언을 참고하고 있습니다. 어떤 도움이 필요하신가요? 데일 카네기와 관련된 질문이 있으시면 편안하게 물어보세요!"
    chatbot = gr.Chatbot(label="채팅창", value=[(None, initial_greeting)])
    msg = gr.Textbox(label="입력")
    clear = gr.Button("초기화")
    msg.submit(respond, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: None, None, chatbot, queue=False)

demo.launch(debug=False)