import logging
import os
import re
import shutil

import gradio as gr
import openai
import pandas as pd
from backoff import on_exception, expo
from sqlalchemy import create_engine

from tools.doc_qa import DocQAPromptAdapter
from tools.web.overwrites import postprocess, reload_javascript
from tools.web.presets import (
    small_and_beautiful_theme,
    title,
    description,
    description_top,
    CONCURRENT_COUNT
)
from tools.web.utils import (
    convert_to_markdown,
    shared_state,
    reset_textbox,
    cancel_outputing,
    transfer_input,
    reset_state,
    delete_last_conversation
)

logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s",
)


openai.api_key = "xxx"
doc_adapter = DocQAPromptAdapter()


def add_llm(model_name, api_base, models):
    """ 添加模型 """
    models = models or {}
    if model_name and api_base:
        models.update(
            {
                model_name: api_base
            }
        )
    choices = [m[0] for m in models.items()]
    return "",  "", models, gr.Dropdown.update(choices=choices, value=choices[0] if choices else None)


def set_openai_env(api_base):
    """ 配置接口地址 """
    openai.api_base = api_base
    doc_adapter.embeddings.openai_api_base = api_base


def get_file_list():
    """ 获取文件列表 """
    if not os.path.exists("doc_store"):
        return []
    return os.listdir("doc_store")


file_list = get_file_list()


def upload_file(file):
    """ 上传文件 """
    if not os.path.exists("doc_store"):
        os.mkdir("docs")

    if file is not None:
        filename = os.path.basename(file.name)
        shutil.move(file.name, f"doc_store/{filename}")
        file_list = get_file_list()
        file_list.remove(filename)
        file_list.insert(0, filename)
        return gr.Dropdown.update(choices=file_list, value=filename)


def add_vector_store(filename, model_name, models, chunk_size, chunk_overlap):
    """ 将文件转为向量数据存储 """
    api_base = models[model_name]
    set_openai_env(api_base)
    doc_adapter.chunk_size = chunk_size
    doc_adapter.chunk_overlap = chunk_overlap

    if filename is not None:
        vs_path = f"vector_store/{filename.split('.')[0]}-{filename.split('.')[-1]}"
        if not os.path.exists(vs_path):
            doc_adapter.create_vector_store(f"doc_store/{filename}", vs_path=vs_path)
            msg = f"Successfully added vector store for {filename}!"
        else:
            doc_adapter.reset_vector_store(vs_path=vs_path)
            msg = f"Successfully loaded vector store for {filename}!"
    else:
        msg = "Please select a file!"
    return msg


def add_db(db_user, db_password, db_host, db_port, db_name, databases):
    """ 添加数据库 """
    databases = databases or {}
    if db_user and db_password and db_host and db_port and db_name:
        databases.update(
            {
                db_name: {
                    "user": db_user,
                    "password": db_password,
                    "host": db_host,
                    "port": int(db_port)
                }
            }
        )
    choices = [m[0] for m in databases.items()]
    return "", "", "", "", "", databases, gr.Dropdown.update(choices=choices, value=choices[0] if choices else None)


def get_table_names(select_database, databases):
    """ 获取数据库表名 """
    if select_database:
        db_config = databases[select_database]
        con = create_engine(f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{select_database}")
        tables = pd.read_sql("show tables;", con=con).values
        tables = [t[0] for t in tables]
        return gr.Dropdown.update(choices=tables, value=[tables[0]])


def get_sql_result(x, con):
    q = r"sql\n(.+?);\n"
    sql = re.findall(q, x, re.DOTALL)[0] + ";"
    df = pd.read_sql(sql, con=con).iloc[:10, :]
    return df.to_markdown(numalign="center", stralign="center")


@on_exception(expo, openai.error.RateLimitError, max_tries=5)
def chat_completions_create(params):
    """ chat接口 """
    return openai.ChatCompletion.create(**params)


def predict(
    model_name,
    models,
    text,
    chatbot,
    history,
    top_p,
    temperature,
    max_tokens,
    memory_k,
    is_kgqa,
    single_turn,
    is_dbqa,
    select_database,
    select_table,
    databases,
):
    api_base = models[model_name]
    set_openai_env(api_base)

    if text == "":
        yield chatbot, history, "Empty context."
        return

    if history is None:
        history = []

    messages = []
    if is_dbqa:
        temperature = 0.0
        db_config = databases[select_database]
        con = create_engine(f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{select_database}")
        table_schema = ""
        for t in select_table:
            table_schema += pd.read_sql(f"show create table {t};", con=con)["Create Table"][0] + "\n\n"
        table_schema = table_schema.replace("DEFAULT NULL", "")
        messages.append(
            {
                "role": "system",
                "content": f"你现在是一名SQL助手，能够根据用户的问题生成准确的SQL查询。已知SQL的建表语句为：{table_schema}根据上述数据库信息，回答相关问题。"
            },
        )
    else:
        if not single_turn:
            for h in history[-memory_k:]:
                messages.extend(
                    [
                        {
                            "role": "user",
                            "content": h[0]
                        },
                        {
                            "role": "assistant",
                            "content": h[1]
                        }
                    ]
                )

    messages.append(
        {
            "role": "user",
            "content": doc_adapter(text) if is_kgqa else text
        }
    )

    params = dict(
        stream=True,
        messages=messages,
        model=model_name,
        top_p=top_p,
        temperature=temperature,
        max_tokens=max_tokens
    )

    res = chat_completions_create(params)
    x = ""
    for openai_object in res:
        delta = openai_object.choices[0]["delta"]
        if "content" in delta:
            x += delta["content"]

        a, b = [[y[0], convert_to_markdown(y[1])] for y in history] + [
            [text, convert_to_markdown(x)]
        ], history + [[text, x]]

        yield a, b, "Generating..."

    if shared_state.interrupted:
        shared_state.recover()
        try:
            yield a, b, "Stop: Success"
            return
        except:
            pass

    if is_dbqa:
        try:
            res = get_sql_result(x, con)
            a[-1][-1] += "\n\n" + convert_to_markdown(res)
            b[-1][-1] += "\n\n" + convert_to_markdown(res)
        except:
            pass

    try:
        yield a, b, "Generate: Success"
    except:
        pass


def retry(
    model_name,
    models,
    text,
    chatbot,
    history,
    top_p,
    temperature,
    max_tokens,
    memory_k,
    is_kgqa,
    single_turn,
    is_dbqa,
    select_database,
    select_table,
    databases,
):
    logging.info("Retry...")
    if len(history) == 0:
        yield chatbot, history, "Empty context."
        return
    chatbot.pop()
    inputs = history.pop()[0]
    for x in predict(
        model_name,
        models,
        inputs,
        chatbot,
        history,
        top_p,
        temperature,
        max_tokens,
        memory_k,
        is_kgqa,
        single_turn,
        is_dbqa,
        select_database,
        select_table,
        databases,
    ):
        yield x


gr.Chatbot.postprocess = postprocess

with open("assets/custom.css", "r", encoding="utf-8") as f:
    customCSS = f.read()

with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
    history = gr.State([])
    user_question = gr.State("")

    with gr.Row():
        gr.HTML(title)
        status_display = gr.Markdown("Success", elem_id="status_display")

    gr.Markdown(description_top)

    with gr.Row(scale=1).style(equal_height=True):
        with gr.Column(scale=5):
            with gr.Row():
                chatbot = gr.Chatbot(elem_id="chuanhu_chatbot").style(height="100%")
            with gr.Row():
                with gr.Column(scale=12):
                    user_input = gr.Textbox(
                        show_label=False, placeholder="Enter text"
                    ).style(container=False)
                with gr.Column(min_width=70, scale=1):
                    submitBtn = gr.Button("发送")
                with gr.Column(min_width=70, scale=1):
                    cancelBtn = gr.Button("停止")
            with gr.Row():
                emptyBtn = gr.Button(
                    "🧹 新的对话",
                )
                retryBtn = gr.Button("🔄 重新生成")
                delLastBtn = gr.Button("🗑️ 删除最旧对话")

        with gr.Column():
            with gr.Column(min_width=50, scale=1):
                with gr.Tab(label="模型"):
                    model_name = gr.Textbox(
                        placeholder="chatglm",
                        label="模型名称",
                    )
                    api_base = gr.Textbox(
                        placeholder="https://0.0.0.0:80/v1",
                        label="模型接口地址",
                    )
                    add_model = gr.Button("添加模型")
                    with gr.Accordion(open=False, label="所有模型配置"):
                        models = gr.Json()
                    single_turn = gr.Checkbox(label="使用单轮对话", value=False)
                    select_model = gr.Dropdown(
                        choices=[m[0] for m in models.value.items()] if models.value else [],
                        value=[m[0] for m in models.value.items()][0] if models.value else None,
                        label="选择模型",
                        interactive=True,
                    )

                with gr.Tab(label="知识库"):
                    is_kgqa = gr.Checkbox(
                        label="使用知识库问答",
                        value=False,
                        interactive=True,
                    )
                    gr.Markdown("""**基于本地知识库生成更加准确的回答！**""")
                    select_file = gr.Dropdown(
                        choices=file_list,
                        label="选择文件",
                        interactive=True,
                        value=file_list[0] if len(file_list) > 0 else None
                    )
                    file = gr.File(
                        label="上传文件",
                        visible=True,
                        file_types=['.txt', '.md', '.docx', '.pdf']
                    )
                    add_vs = gr.Button(value="添加到知识库")

                with gr.Tab(label="数据库"):
                    with gr.Accordion(open=False, label="数据库配置"):
                        db_user = gr.Textbox(
                            placeholder="root",
                            label="用户名",
                        )
                        db_password = gr.Textbox(
                            placeholder="password",
                            label="密码",
                            type="password"
                        )
                        db_host = gr.Textbox(
                            placeholder="0.0.0.0",
                            label="主机",
                        )
                        db_port = gr.Textbox(
                            placeholder="3306",
                            label="端口",
                        )
                        db_name = gr.Textbox(
                            placeholder="test",
                            label="数据库名称",
                        )
                    add_database = gr.Button("添加数据库")

                    with gr.Accordion(open=False, label="所有数据库配置"):
                        databases = gr.Json()
                    select_database = gr.Dropdown(
                        choices=[d[0] for d in databases.value.items()] if databases.value else [],
                        value=[d[0] for d in databases.value.items()][0] if databases.value else None,
                        interactive=True,
                        label="选择数据库"
                    )
                    select_table = gr.Dropdown(label="选择表", interactive=True, multiselect=True)
                    is_dbqa = gr.Checkbox(
                        label="使用数据库问答",
                        value=False,
                        interactive=True,
                    )

                with gr.Tab(label="参数"):
                    top_p = gr.Slider(
                        minimum=-0,
                        maximum=1.0,
                        value=0.95,
                        step=0.05,
                        interactive=True,
                        label="Top-p",
                    )
                    temperature = gr.Slider(
                        minimum=0.1,
                        maximum=2.0,
                        value=1,
                        step=0.1,
                        interactive=True,
                        label="Temperature",
                    )
                    max_tokens = gr.Slider(
                        minimum=0,
                        maximum=512,
                        value=512,
                        step=8,
                        interactive=True,
                        label="Max Generation Tokens",
                    )
                    memory_k = gr.Slider(
                        minimum=0,
                        maximum=10,
                        value=5,
                        step=1,
                        interactive=True,
                        label="Max Memory Window Size",
                    )
                    chunk_size = gr.Slider(
                        minimum=100,
                        maximum=1000,
                        value=200,
                        step=100,
                        interactive=True,
                        label="Chunk Size",
                    )
                    chunk_overlap = gr.Slider(
                        minimum=0,
                        maximum=100,
                        value=0,
                        step=10,
                        interactive=True,
                        label="Chunk Overlap",
                    )

    gr.Markdown(description)

    add_model.click(
        add_llm,
        inputs=[model_name, api_base, models],
        outputs=[model_name, api_base, models, select_model],
    )

    add_database.click(
        add_db,
        inputs=[db_user, db_password, db_host, db_port, db_name, databases],
        outputs=[db_user, db_password, db_host, db_port, db_name, databases, select_database],
    )

    select_database.change(
        get_table_names,
        inputs=[select_database, databases],
        outputs=select_table,
    )

    file.upload(
        upload_file,
        inputs=file,
        outputs=select_file,
    )

    add_vs.click(
        add_vector_store,
        inputs=[select_file, select_model, models, chunk_size, chunk_overlap],
        outputs=status_display,
    )

    predict_args = dict(
        fn=predict,
        inputs=[
            select_model,
            models,
            user_question,
            chatbot,
            history,
            top_p,
            temperature,
            max_tokens,
            memory_k,
            is_kgqa,
            single_turn,
            is_dbqa,
            select_database,
            select_table,
            databases,
        ],
        outputs=[chatbot, history, status_display],
        show_progress=True,
    )
    retry_args = dict(
        fn=retry,
        inputs=[
            select_model,
            models,
            user_question,
            chatbot,
            history,
            top_p,
            temperature,
            max_tokens,
            memory_k,
            is_kgqa,
            single_turn,
            is_dbqa,
            select_database,
            select_table,
            databases,
        ],
        outputs=[chatbot, history, status_display],
        show_progress=True,
    )

    reset_args = dict(fn=reset_textbox, inputs=[], outputs=[user_input, status_display])

    cancelBtn.click(cancel_outputing, [], [status_display])
    transfer_input_args = dict(
        fn=transfer_input,
        inputs=[user_input],
        outputs=[user_question, user_input, submitBtn, cancelBtn],
        show_progress=True,
    )

    user_input.submit(**transfer_input_args).then(**predict_args)

    submitBtn.click(**transfer_input_args).then(**predict_args)

    emptyBtn.click(
        reset_state,
        outputs=[chatbot, history, status_display],
        show_progress=True,
    )
    emptyBtn.click(**reset_args)

    retryBtn.click(**retry_args)

    delLastBtn.click(
        delete_last_conversation,
        [chatbot, history],
        [chatbot, history, status_display],
        show_progress=True,
    )

demo.title = "OpenLLM Chatbot 🚀 "

if __name__ == "__main__":
    reload_javascript()
    demo.queue(concurrency_count=CONCURRENT_COUNT).launch()