import gradio as gr
import mdtex2html
import os
import json
import requests

def postprocess(self, y):
    if y is None:
        return []
    for i, (message, response) in enumerate(y):
        y[i] = (
            None if message is None else mdtex2html.convert((message)),
            None if response is None else mdtex2html.convert(response),
        )
    return y


gr.Chatbot.postprocess = postprocess


def parse_text(text):
    """copy from https://github.com/GaiZhenbiao/ChuanhuChatGPT/"""
    lines = text.split("\n")
    lines = [line for line in lines if line != ""]
    count = 0
    for i, line in enumerate(lines):
        if "```" in line:
            count += 1
            items = line.split('`')
            if count % 2 == 1:
                lines[i] = f'<pre><code class="language-{items[-1]}">'
            else:
                lines[i] = f'<br></code></pre>'
        else:
            if i > 0:
                if count % 2 == 1:
                    line = line.replace("`", "\`")
                    line = line.replace("<", "&lt;")
                    line = line.replace(">", "&gt;")
                    line = line.replace(" ", "&nbsp;")
                    line = line.replace("*", "&ast;")
                    line = line.replace("_", "&lowbar;")
                    line = line.replace("-", "&#45;")
                    line = line.replace(".", "&#46;")
                    line = line.replace("!", "&#33;")
                    line = line.replace("(", "&#40;")
                    line = line.replace(")", "&#41;")
                    line = line.replace("$", "&#36;")
                lines[i] = "<br>"+line
    text = "".join(lines)
    return text


def SubmitKey(chatbot, key):
    data = {'AgentStart': True, 'key': key}
    response=requests.post(os.environ.get("URL"), data=json.dumps(data, ensure_ascii=False).encode('utf-8'), timeout=500)
    msg = str(response.content, encoding="utf-8")
    if not response.status_code == 200:
        chatbot = [(None, parse_text(msg))]
        return chatbot, gr.update(value=msg), gr.update(value=msg)
    response = json.loads(msg)
    chatbot = response['chatbot']
    return chatbot, gr.update(value=response['story']), gr.update(value=response['answer'])

def CheckTrue(chatbot, key):
    data = {'AgentCheck': True, 'key': key}
    response=requests.post(os.environ.get("URL"), data=json.dumps(data, ensure_ascii=False).encode('utf-8'), timeout=500)
    chatbot.append((parse_text("是。"), parse_text(str(response.content, encoding="utf-8"))))
    return chatbot

def CheckFalse(chatbot, key):
    data = {'AgentCheck': False, 'key': key}
    response=requests.post(os.environ.get("URL"), data=json.dumps(data, ensure_ascii=False).encode('utf-8'), timeout=500)
    chatbot.append((parse_text("否。"), parse_text(str(response.content, encoding="utf-8"))))
    return chatbot

def CheckIrrelevant(chatbot, key):
    data = {'AgentCheck': 'Irrelevant', 'key': key}
    response=requests.post(os.environ.get("URL"), data=json.dumps(data, ensure_ascii=False).encode('utf-8'), timeout=500)
    chatbot.append((parse_text("无关。"), parse_text(str(response.content, encoding="utf-8"))))
    return chatbot

def CheckTerm(chatbot, key):
    data = {'AgentFinish': True, 'key': key}
    response=requests.post(os.environ.get("URL"), data=json.dumps(data, ensure_ascii=False).encode('utf-8'), timeout=500)
    msg = str(response.content, encoding="utf-8")
    if not response.status_code == 200:
        chatbot = [(None, parse_text(msg))]
        return chatbot, gr.update(value=msg), gr.update(value=msg)
    response = json.loads(msg)
    chatbot = [(None, parse_text(response['response']))]
    return chatbot, gr.update(value=response['story']), gr.update(value=response['answer'])

with gr.Blocks() as demo:
    gr.HTML("""<h1 align="center">Generative Agents测试平台</h1>""")
    user_key = gr.Textbox(label='密钥', placeholder="Input your key, press enter to apply", lines=1, max_lines=1).style(
        container=False)
    with gr.Row():
        with gr.Column(scale=1):
            story = gr.Textbox(label='汤面', lines=3, max_lines=3).style(
                container=False)
        with gr.Column(scale=1):
            answer = gr.Textbox(label='汤底', lines=3, max_lines=3).style(
                container=False)
    chatbot = gr.Chatbot([])

    with gr.Row():
        with gr.Column(scale=1):
            trueBtn = gr.Button('是')
        with gr.Column(scale=1):
            falseBtn = gr.Button('否')
        with gr.Column(scale=1):
            irreBtn = gr.Button('无关')
        with gr.Column(scale=1):
            termBtn = gr.Button('达到终止条件')
        
    user_key.submit(SubmitKey, [chatbot, user_key], [chatbot, story, answer])
    trueBtn.click(CheckTrue, [chatbot, user_key], [chatbot])
    falseBtn.click(CheckFalse, [chatbot, user_key], [chatbot])
    irreBtn.click(CheckIrrelevant, [chatbot, user_key], [chatbot])
    termBtn.click(CheckTerm, [chatbot, user_key], [chatbot, story, answer])
    
demo.queue().launch()