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

from constants import persist_directory

embedding = OpenAIEmbeddings()
client = chromadb.PersistentClient(path=persist_directory)
vector_store = Chroma(client=client, embedding_function=embedding)
retriever = vector_store.as_retriever(search_kwargs={"k": 10})

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

system_template = """Your name is Sayno. You act like a successful rich man. Never reveal that it is an AI. 
Given the following extracted parts of a long document and a question, create a final answer with references ("SOURCES"). 
If you don't know the answer, just say that you don't know. Don't try to make up an answer.
!IMPORTANT: ALWAYS return a "SOURCES" part in your answer.
!IMPORTANT: Be sure to answer in Korean.

[EXAMPLE]
QUESTION: Which state/country's law governs the interpretation of the contract?
=========
Content: This Agreement is governed by English law and the parties submit to the exclusive jurisdiction of the English courts in  relation to any dispute (contractual or non-contractual) concerning this Agreement save that either party may apply to any court for an  injunction or other relief to protect its Intellectual Property Rights.
Source: 28-pl
Content: No Waiver. Failure or delay in exercising any right or remedy under this Agreement shall not constitute a waiver of such (or any other)  right or remedy.\n\n11.7 Severability. The invalidity, illegality or unenforceability of any term (or part of a term) of this Agreement shall not affect the continuation  in force of the remainder of the term (if any) and this Agreement.\n\n11.8 No Agency. Except as expressly stated otherwise, nothing in this Agreement shall create an agency, partnership or joint venture of any  kind between the parties.\n\n11.9 No Third-Party Beneficiaries.
Source: 30-pl
Content: (b) if Google believes, in good faith, that the Distributor has violated or caused Google to violate any Anti-Bribery Laws (as  defined in Clause 8.5) or that such a violation is reasonably likely to occur,
Source: 4-pl
=========
FINAL ANSWER: This Agreement is governed by English law.
SOURCES: 28-pl

QUESTION: {question}
=========
{summaries}
=========
FINAL ANSWER:
"""

prompt = ChatPromptTemplate.from_messages(
    [
        SystemMessagePromptTemplate.from_template(system_template),
    ]
)

document_prompt = PromptTemplate(
    template="Content: {page_content}\nSource: {source}, {page} page",
    input_variables=["page_content", "source", "page"],
)

chain_type_kwargs = {"prompt": prompt, "document_prompt": document_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):
    result = chain(message)
    print(result)
    bot_message = f'{result["answer"]}<br>- 출처: {result["sources"]}'
    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)