import g4f
import gradio as gr
from gradio import ChatInterface
from g4f.Provider import (
    Ails,
    You,
    Bing,
    Yqcloud,
    Theb,
    Aichat,
    Bard,
    Vercel,
    Forefront,
    Lockchat,
    Liaobots,
    H2o,
    ChatgptLogin,
    DeepAi,
    GetGpt,
    AItianhu,
    EasyChat,
    Acytoo,
    DfeHub,
    AiService,
    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
from models_for_langchain.memory_func import validate_memory_len

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,
    'Wewordle': Wewordle,
    'ChatgptAi': ChatgptAi,
}


with open("available_dict.txt", "r") as fp:
    # Load the dictionary from the file
    available_dict = json.load(fp)

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(history, model_name, provider_name, system_msg, agent):
    history[-1][1] = ''
    message = history[-1][0]

    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)

    if len(history)>1 and history[-2][1]!=None:
        memory.chat_memory.add_ai_message(history[-2][1]) 
    memory.chat_memory.add_user_message(history[-1][0])
    validate_memory_len(memory=memory, max_token_limit=1800)
    if len(memory.chat_memory.messages)==0:
        for c in '文本长度超过限制，请清空后再试':
            history[-1][1] += c
            yield history
    else:
        prev_memory = memory.load_memory_variables({})['chat_history']
        prompt = new_template.format(
            chat_history = prev_memory,
        )
        print(f'prompt = \n --------\n{prompt}\n --------')
        for _ in range(3):
            try:
                bot_msg = llm._call(prompt=prompt)
                break
            except:
                bot_msg = '服务器无响应，请更换提供者或者清空对话后重试。'

        for c in bot_msg:
            history[-1][1] += c
            yield history

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

def undo_fn(history):
    return history[:-1]

def retry_fn(history):
    history[-1][1] = None
    return history

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}}
    AI:"""

    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=available_dict['gpt-3.5-turbo'][0], label='提供者', min_width=20)
        agent = gr.Dropdown(['系统提示', '维基百科'], value='系统提示', label='Agent')
    system_msg = gr.Textbox(value="你是一名助手，可以解答问题。", label='系统提示')
    
    chatbot = gr.Chatbot([[None, None]], label='AI')
    with gr.Group():
        with gr.Row():
            textbox = gr.Textbox(
                                    container=False,
                                    show_label=False,
                                    label="请输入：",
                                    scale=7,
                                    autofocus=True,
                                )
            submit = gr.Button('发送', scale=1, variant="primary", min_width=150,)
    with gr.Row():
        retry = gr.Button('🔄 重试')
        undo = gr.Button('↩️ 撤销')
        clear = gr.Button("🗑️ 清空")

    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)

    textbox.submit(user, [textbox, chatbot], [textbox, chatbot], queue=False).then(
        bot, [chatbot, model_name, provider, system_msg, agent], chatbot
    ).then(lambda: gr.update(interactive=True), None, [textbox], queue=False)

    response = submit.click(user, [textbox, chatbot], [textbox, chatbot], queue=False).then(
        bot, [chatbot, model_name, provider, system_msg, agent], chatbot
    ).then(lambda: gr.update(interactive=True), None, [textbox], queue=False)

    retry.click(retry_fn, [chatbot], [chatbot]).then(
        bot, [chatbot, model_name, provider, system_msg, agent], chatbot
    )
    undo.click(undo_fn, [chatbot], [chatbot], queue=False)
    clear.click(empty_fn, None, [chatbot], queue=False)

    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()