import gradio as gr
from langchain.chains import LLMChain
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import TextLoader
from embedding import CustomEmbedding
from memories import HumenFeedbackBufferMemory
from langchain.memory import ConversationBufferMemory
from promopts import FEEDBACK, FEEDBACK_PROMPT

llm = ChatOpenAI(temperature=0.1)

memory = HumenFeedbackBufferMemory(
    input_key="input", human_prefix="Answer", ai_prefix="AI")

llmchain = LLMChain(
    llm=llm,
    prompt=FEEDBACK_PROMPT,
    memory=memory,
    verbose=True
)

"""读取document/business_context.py文件内容作为context"""
context_path = "./documents/bussiness_context/business_context.md"
textloader = TextLoader(context_path)
CONTEXT = textloader.load()[0].page_content


def sendMessage(chatbot, input):
    chatbot.append((
        (None if len(input) == 0 else input), None))
    return chatbot


def clearMemory(chatbot):
    chatbot.clear()
    memory.clear()
    return chatbot, ""


def feedBack(context, story, chatbot=[], input=""):
    if len(input) > 0:
        context += (f"\n- {input}")
        with open(context_path, 'w') as f:
            f.write(context)

    response = llmchain.run(
        input=(input if len(input) == 0 else ("Answer: " + input)), context=context, story=story, stop="\nAnswer:")
    chatbot[-1][1] = response
    return chatbot, "", context


customerEmbedding = CustomEmbedding()

faqChain = customerEmbedding.getFAQChain()


def faqFromLocal(input, chatbot=[]):
    response = faqChain.run(input)
    chatbot.append((input, response))
    return chatbot, ""


def generateEmbeddings():
    customerEmbedding.calculateBusinessContextEmbedding()
    customerEmbedding.calculateNotionEmbedding()


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            chatbot = gr.Chatbot(elem_id="chatbot").style()

            with gr.Row():
                txt = gr.Textbox(show_label=False, placeholder="Enter text and press enter").style(
                    container=False)
            with gr.Row():
                faq = gr.Textbox(show_label=False, placeholder="A&Q from local context").style(
                    container=False)
        with gr.Column():
            with gr.Row():
                context = gr.Textbox(show_label=True, label="Context", value=CONTEXT, placeholder="Enter Context").style(
                    container=False)

            with gr.Row():
                story = gr.Textbox(show_label=True, label="Story", placeholder="Enter Story").style(
                    container=False)

            with gr.Row():
                button = gr.Button("Generate Scenarios")

            with gr.Blocks():
                with gr.Row():
                    generateEmbedding = gr.Button("Generate embedding")

    button.click(clearMemory, [chatbot], [chatbot, txt]).then(sendMessage, [chatbot, txt], [chatbot]).then(
        feedBack, [context, story, chatbot], [chatbot, txt])
    txt.submit(sendMessage, [chatbot, txt], [chatbot]).then(
        feedBack, [context, story, chatbot, txt], [chatbot, txt, context])
    faq.submit(faqFromLocal, [faq, chatbot], [chatbot, faq])
    generateEmbedding.click(generateEmbeddings)

demo.launch()