import os

import gradio as gr

from bot.chatbot import ChatBot
from bot.config import special_token_list

bot:ChatBot = None

def get_skill_list() -> list:
    path:str = os.path.split( os.path.realpath(__file__) )[0]
    file_list:list[str] = os.listdir( path + "/bot/skills/" )
    plugin_list:list[str] = []
    for file in file_list:
        if file.endswith(".py"):
            plugin_list.append( file[:-3] )
    return plugin_list

def general(input_txt:str, state:dict = {}):
    global bot

    history_list:list = state.get("history", [])
    role_card:dict[str, str] = state.get("role_card", {
        "<NAME>": "Winnie",
        "<GENDER>": "女",
        "<YEAROFBIRTH>":"1995",
        "<MONTHOFBIRTH>":"5",
        "<DAYOFBIRTH>":"6",
        "<ZODIAC>":"金牛座",
        "<AGE>":"27"
        }
    )

    output_txt:str = None

    for skill_name in get_skill_list():
        if output_txt is None:
            plugin = __import__("bot.skills."+skill_name, fromlist=[skill_name])
            plugin_class = getattr(plugin, "Skill")
            p = plugin_class()
            p.set_tokenizer( bot.tokenizer )
            output_txt, history_list, role_card = p.process(input_txt, history_list, role_card)

    if output_txt is None:
        res, history_list = bot.chat(input_txt, history_list, role_card=role_card)
        output_txt = "".join(res)

    state["history"] = history_list
    state["role_card"] = role_card

    return output_txt, state

def main() -> None:
    global bot

    bot = ChatBot.get_chat_bot("lewiswu1209/Winnie", special_token_list=special_token_list)

    title:str = "使用中文和Winnie聊天"

    description:str = "输入任意文字，Winnie会和你对话<br>"
    description += "输入ERASE MEMORY，会清空Winnie的记忆<br>"
    description += "输入\"&lt;TAG&gt;=&lt;VALUE&gt;\"，可以修改Winnie的角色信息<br>"
    description += "例如：&lt;NAME&gt;=Vicky，会修改Winnie的名字<br>"
    description += "可以修改的角色信息有：<br>"
    description += "&lt;NAME&gt;, &lt;GENDER&gt;, &lt;YEAROFBIRTH&gt;, &lt;MONTHOFBIRTH&gt;, &lt;DAYOFBIRTH&gt;, &lt;ZODIAC&gt;, &lt;AGE&gt;<br>"
    description += "输入“上联：XXXXXXX”，Winnie会和你对对联<br>"
    description += "输入“写诗：XXXXXXX”，Winnie会以XXXXXXX为开头写诗<br>"
    description += "以\"请问\"开头并以问号结尾，Winnie会回答该问题"

    gr.Interface(
        fn = general,
        title = title,
        description = description,
        inputs = ["text", "state"],
        outputs = ["text", "state"]
    ).launch()

if __name__ == "__main__":
    main()