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 test(chatbot):
    chatbot.append((None, parse_text('测试')))
    chatbot.append((parse_text('测试'), None))
    return chatbot

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'))
    if not response.status_code == 200:
        chatbot.append((parse_text("出现错误，请联系负责人"), None))
        return chatbot
    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'))
    if not response.status_code == 200:
        chatbot.append((parse_text("出现错误，请联系负责人"), None))
        return chatbot
    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'))
    if not response.status_code == 200:
        chatbot.append((parse_text("出现错误，请联系负责人"), None))
        return chatbot
    chatbot = [(None, parse_text(str(response.content, encoding="utf-8")))]
    return chatbot

with gr.Blocks() as demo:
    gr.HTML("""<h1 align="center">Generative Agents测试平台</h1>""")
    user_key = gr.Textbox(label='Key', placeholder="Input your key...", lines=1, max_lines=1).style(
        container=False)
    chatbot = gr.Chatbot([])

    with gr.Row():
        test_btn = gr.Button('测试')
        
    user_key.submit(CheckTerm, [chatbot, user_key], [chatbot])
    test_btn.click(test, [chatbot], [chatbot])
    
demo.queue().launch()