import gradio as gr
from langchain.document_loaders import TextLoader
from agents.tools.python_code_tool import generate_and_excute_python_code
from chains import HumanFeedBackChain, contextRewriteChain
from embedding import CustomEmbedding
from memories import HumenFeedbackBufferMemory
from agents.code_generate_agent import code_agent_executor, code_agent_tools


baMemory = HumenFeedbackBufferMemory(
            input_key="input", human_prefix="Answer", ai_prefix="AI")
baChain = HumanFeedBackChain(verbose=True, memory=baMemory)

"""读取document/business_context.py文件内容作为context"""
context_path = "./documents/bussiness_context/business_context.md"


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


def clearMemory(chatbot):
    chatbot.clear()
    if baMemory != None:
        baMemory.clear()
    return chatbot, ""

def loadContext():
    textloader = TextLoader(context_path)
    return textloader.load()[0].page_content


def saveContext(context):
    with open(context_path, 'w') as f:
        f.write(context)

def feedBack(context, story, chatbot=[], input=""):
    if len(input) > 0:
        context += (f"\n\n {input}")
        saveContext(context)
    response = baChain.run(
        input=(input if len(input) == 0 else input), context=context, story=story, stop="\nAnswer:")
    chatbot[-1][1] = response
    return chatbot, "", context


customerEmbedding = CustomEmbedding()

faqChain = customerEmbedding.getFAQChain()

code_agent_executor = code_agent_executor()
def faqFromLocal(input, chatbot=[]):
    response = faqChain({"question": f"{input}"})
    chatbot.append((input, response["answer"]))
    return chatbot, ""


def generateEmbeddings(chatbot=[]):
    response = customerEmbedding.calculateEmbedding()
    chatbot.append((None, response))
    return chatbot


def generateCode(input: str, chatbot=[], returnCode=False):
    if len(input) <=0:
        chatbot[-1][1] = None
        return chatbot, ""
    response = code_agent_executor.run(
        input=(input if len(input) == 0 else input))
    chatbot[-1][1] = response
    return chatbot, ""

def generateCodeByMultiPart(context: str, relateCode: str, toolName: str, chatbot=[]):
    input = f"请根据如下信息{toolName}:\n{context}\n\n{relateCode}"
    return generateCode(input, chatbot)

def sendMessageByMultiPart(chatbot, context: str, relateCode: str, toolName: str):
    input = f"请根据如下信息{toolName}:\n{context}\n\n{relateCode}"
    chatbot.append((input, None))
    return chatbot


def rewriteContext(input, chatbot):
    response = contextRewriteChain.run(input=input, verbose=True)
    chatbot.append((input, response))
    return chatbot, response

def generateCodeAndExcute(input, chatbot=[]):
    result = generate_and_excute_python_code.run(input)
    chatbot.append((input, result))
    return chatbot

toolTextBox = []
with gr.Blocks() as demo:
    with gr.Row():
            with gr.Tab("Business"):
                with gr.Row():
                    with gr.Column():
                        chatbot = gr.Chatbot().style()
                        with gr.Row():
                            txt = gr.Textbox(show_label=False, placeholder="Enter text and press enter").style(
                                container=False)
                    with gr.Column():
                        with gr.Row():
                            context = gr.Textbox(show_label=True, label="Context", placeholder="Enter Context").style(
                                container=False)
                        with gr.Row():
                            story = gr.Textbox(show_label=True, label="User Story", placeholder="Enter User Story").style(
                                container=False)
                        with gr.Row():
                            gr.Button("Generate Scenarios").click(clearMemory, [chatbot], [chatbot, txt]).then(sendMessage, [chatbot, txt], [chatbot]).then(
                                feedBack, [context, story, chatbot], [chatbot, txt])
                        with gr.Row():
                            with gr.Column(scale=5):
                                gr.Button("Rewrite Context").click(rewriteContext, [context, chatbot], [chatbot, context])
                            with gr.Column(scale=1):
                                gr.Button("Revert").click(loadContext, [], [context])
                        with gr.Row():
                            gr.Button("Save Context").click(saveContext, [context], [])

            with gr.Tab("Tech"):
                with gr.Row():
                    with gr.Column():
                        code_chatbot = gr.Chatbot().style()
                        with gr.Row():
                            code = gr.Textbox(show_label=False, label="Code Generate", placeholder="Enter text and press enter").style(
                            container=False)
                    with gr.Column():
                        with gr.Row():
                            code_context = gr.Textbox(show_label=True, label="Context", placeholder="Enter Context").style(
                                container=False)
                        with gr.Row():
                            relateCode = gr.Textbox(show_label=True, label="Relate Code", placeholder="Enter Relate Code").style(
                                container=False)
                        for index, tool in enumerate(code_agent_tools):
                            with gr.Row():
                                toolTextBox.append(gr.Textbox(show_label=False, visible=False, label=tool.name, value=tool.name).style())
                                gr.Button(tool.name).click(
                                    sendMessageByMultiPart, [code_chatbot, code_context, relateCode, toolTextBox[index]], [code_chatbot]).then(
                                    generateCodeByMultiPart, [code_context, relateCode, toolTextBox[index], code_chatbot], [code_chatbot, code])
            with gr.Tab("FAQ"):
                faq_chatbot = gr.Chatbot().style()
                with gr.Row():
                    faq = gr.Textbox(show_label=False, placeholder="Enter text and press enter").style(
                        container=False)
                with gr.Row():
                    gr.Button("Regenerate embedding").click(generateEmbeddings,[faq_chatbot], [faq_chatbot])
            with gr.Tab("TOOL"):
                with gr.Row():
                    with gr.Column():
                        tool_request = gr.Textbox(show_label=False, placeholder="Enter your tool Request").style(
                            container=False, show_copy_button=True)
                        tool_button = gr.Button("Generate Code and Execute")
                    with gr.Column():
                        tool_chatbot = gr.Chatbot(elem_id="chatbot").style(container=False)
                    tool_button.click(generateCodeAndExcute,[tool_request, tool_chatbot], [tool_chatbot])

        
    txt.submit(sendMessage, [chatbot, txt], [chatbot]).then(
        feedBack, [context, story, chatbot, txt], [chatbot, txt, context])
    
    code.submit(sendMessage, [code_chatbot, code], [code_chatbot]).then(
        generateCode, [code, code_chatbot], [code_chatbot, code])
    
    faq.submit(faqFromLocal, [faq, faq_chatbot], [faq_chatbot, faq])
    
    demo.load(loadContext, [], [context])    
demo.launch()