import g4f
import gradio as gr
from g4f.Provider import (
    Ails,
    You,
    Bing,
    Yqcloud,
    Theb,
    Aichat,
    Bard,
    Vercel,
    Forefront,
    Lockchat,
    Liaobots,
    H2o,
    ChatgptLogin,
    DeepAi,
    GetGpt,
    AItianhu,
    EasyChat,
    Acytoo,
    DfeHub,
    AiService,
    BingHuan,
    Wewordle,
    ChatgptAi,
)
import os
import json
import pandas as pd
from langchain.tools.python.tool import PythonREPLTool
from langchain.agents.agent_toolkits import create_python_agent
from models_for_langchain.model import CustomLLM
from langchain.memory import ConversationBufferWindowMemory, ConversationTokenBufferMemory
from langchain import LLMChain, PromptTemplate
from langchain.prompts import (
    ChatPromptTemplate,
    PromptTemplate,
    SystemMessagePromptTemplate,
    AIMessagePromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.agents.agent_types import AgentType
from langchain.tools import WikipediaQueryRun
from langchain.utilities import WikipediaAPIWrapper
from langchain.tools import DuckDuckGoSearchRun

provider_dict = {
    'Ails': Ails,
    'You': You,
    'Bing': Bing,
    'Yqcloud': Yqcloud,
    'Theb': Theb,
    'Aichat': Aichat,
    'Bard': Bard,
    'Vercel': Vercel,
    'Forefront': Forefront,
    'Lockchat': Lockchat,
    'Liaobots': Liaobots,
    'H2o': H2o,
    'ChatgptLogin': ChatgptLogin,
    'DeepAi': DeepAi,
    'GetGpt': GetGpt,
    'AItianhu': AItianhu,
    'EasyChat': EasyChat,
    'Acytoo': Acytoo,
    'DfeHub': DfeHub,
    'AiService': AiService,
    'BingHuan': BingHuan,
    'Wewordle': Wewordle,
    'ChatgptAi': ChatgptAi,
}

available_dict = {
    'gpt-3.5-turbo':['Acytoo', 'AiService', 'Aichat', 'GetGpt', 'Wewordle'],
    'gpt-4':['ChatgptAi'],
    'falcon-7b':['H2o'],
    'falcon-13b':['H2o'],
    'llama-13b':['H2o']
}

def change_prompt_set(prompt_set_name):
    return gr.Dropdown.update(choices=list(prompt_set_list[prompt_set_name].keys()))

def change_model(model_name):
    new_choices = list(available_dict[model_name])
    return gr.Dropdown.update(choices=new_choices, value=new_choices[0])

def change_prompt(prompt_set_name, prompt_name):
    return gr.update(value=prompt_set_list[prompt_set_name][prompt_name])

def user(user_message, history):
    return gr.update(value="", interactive=False), history + [[user_message, None]]

def bot(message, history, model_name, provider_name, system_msg, agent):
    response = ''

    if len(system_msg)>3000:
        system_msg = system_msg[:2000] + system_msg[-1000:]

    global template, memory
    llm.model_name = model_name
    llm.provider_name = provider_name
    if agent == '系统提示':
        new_template = template.format(system_instruction=system_msg)
    elif agent == '维基百科':
        wikipedia = WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper())
        target = llm(f'用户的问题：```{message}```。为了回答用户的问题，你需要在维基百科上进行搜索，只有一次搜索的机会，请返回需要搜索的词汇，只需要返回一个英文词汇，不要加任何解释：')
        new_template = template.format(system_instruction=wikipedia.run(str(target)))
    elif agent == 'duckduckgo':
        search = DuckDuckGoSearchRun()
        target = llm(f'用户的问题：```{message}```。为了回答用户的问题，你需要在duckduckgo搜索引擎上进行搜索，只有一次搜索的机会，请返回需要搜索的内容，只需要返回纯英文的搜索语句，不要加任何解释：')     
        new_template = template.format(system_instruction=search.run(str(target))) 
    elif agent == 'python':
        py_agent = create_python_agent(
                llm,
                tool=PythonREPLTool(), # REPL，一种代码交互方式，类似jupyter，可以执行代码
                verbose=True,
                # agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION
                handle_parsing_errors=True, # 输出无法解析，返回给llm要求改正。
            )
        response = py_agent.run(message)
        return str(response)
    else:
        new_template = template.format(system_instruction=system_msg)
    prompt = PromptTemplate(
                        input_variables=["chat_history", "human_input"], template=new_template
                        )
    llm_chain = LLMChain(
                        llm=llm,
                        prompt=prompt,
                        verbose=True,
                        memory=memory,
                    )
    bot_msg = llm_chain.run(message)
    for c in bot_msg:
        response += c
    return response

def empty_chat():
    global memory
    memory = ConversationBufferWindowMemory(k=6, memory_key="chat_history")
    return None

prompt_set_list = {}
for prompt_file in os.listdir("prompt_set"):
    key = prompt_file
    if '.csv' in key:
        df = pd.read_csv("prompt_set/" + prompt_file)
        prompt_dict = dict(zip(df['act'], df['prompt']))
    else:
        with open("prompt_set/" + prompt_file, encoding='utf-8') as f:
            ds = json.load(f)
        prompt_dict = {item["act"]: item["prompt"] for item in ds}
    prompt_set_list[key] = prompt_dict

with gr.Blocks() as demo:
    llm = CustomLLM()

    template = """
    Chat with human based on following instructions:
    ```
    {system_instruction}
    ```
    The following is a conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.
    {{chat_history}}
    Human: {{human_input}}
    Chatbot:"""

    memory = ConversationBufferWindowMemory(k=6, memory_key="chat_history")
    with gr.Row():
        model_name = gr.Dropdown(list(available_dict.keys()), value='gpt-3.5-turbo', label='模型')
        provider = gr.Dropdown(available_dict['gpt-3.5-turbo'], value='AiService', label='提供者', min_width=20)
        agent = gr.Dropdown(['系统提示', '维基百科', 'duckduckgo'], value='系统提示', label='Agent')
    system_msg = gr.Textbox(value="你是一名助手，可以解答问题。", label='系统提示')
    gr.ChatInterface(bot,
                     additional_inputs=[
                                model_name,
                                provider,
                                system_msg,
                                agent]
                    )
    with gr.Row():
        default_prompt_set = "1 中文提示词.json"
        prompt_set_name = gr.Dropdown(prompt_set_list.keys(), value=default_prompt_set, label='提示词集合')
        prompt_name = gr.Dropdown(prompt_set_list[default_prompt_set].keys(), label='提示词', min_width=5, container=True)

    prompt_set_name.select(change_prompt_set, prompt_set_name, prompt_name)
    model_name.select(change_model, model_name, provider)
    prompt_name.select(change_prompt, [prompt_set_name, prompt_name], system_msg)

demo.title = "AI Chat"
demo.queue()
demo.launch()