import os
import pandas as pd
import gradio
from datetime import datetime

import models
import vdb

def refresh_csv_knowledge_to_vector_db(knowledges_path: str):
    embedding = models.model_embedding()
    count     = 0
    ids       = []
    questions = []
    answers   = []
    vectors   = []
    source    = []
    for fl in os.listdir(knowledges_path):
        if fl.endswith(".csv"):
            with open(knowledges_path + os.sep + f"{fl}", "r") as csv_file:
                for index, row in pd.read_csv(csv_file).iterrows():
                    count += 1
                    ids.append(count)
                    questions.append(row['question'])
                    answers.append(row['answer'])
                    vectors.append(embedding.encode(f"{row['question']}\n\n{row['answer']}"))
                    source.append(fl.split(".csv")[0])
                print(f"fetch {len(questions)} lines from file `{fl}`. ")   
    if len(questions) == 0:
        print("no csv knowledge to process")
    entity = [ids, questions, answers, source, vectors]
    db = vdb.Milvus()
    db.delete_collection()
    db.create_collection(dim=len(vectors[0]))
    db.upsert_document(entities=entity)

if __name__ == "__main__":
    knowledges_path = os.path.dirname(os.path.realpath(__file__)) + os.sep + 'knowledges'
    refresh_csv_knowledge_to_vector_db(knowledges_path=knowledges_path)

    with gradio.Blocks() as app:
        with gradio.Tab(label="个人助理"):
            chatbot = gradio.Chatbot(label="Robot", show_copy_button=True)
            llm = models.model_llm()
            llm_modes = llm.get_aviable_mode_names()
            llm_selected = gradio.State(value=llm_modes[1])
            with gradio.Row():
                text_input = gradio.Textbox(label="问题", placeholder="请输入您想问的问题")
                with gradio.Column():
                    btn_submit = gradio.Button(value="提交", interactive=True)
                    btn_clear = gradio.ClearButton(components=[text_input, chatbot], value="清空")
                with gradio.Column():
                    drop_mode = gradio.Dropdown(choices=llm_modes, value=llm_modes[1], label="大语言模型选择:")

            def response(message, chat_history):
                bot_message = llm.retrival_inference_answer(message, chat_history, llm_selected.value)
                chat_history.append((message, bot_message))

                return "", chat_history
            
            def mode_selected(mode):
                llm_selected.value = mode

            drop_mode.change(fn=mode_selected, inputs=[drop_mode], outputs=[llm_selected])
            btn_submit.click(fn=response, inputs=[text_input, chatbot], outputs=[text_input, chatbot])
        
        with gradio.Tab(label="知识库管理"):
            datas = []
            for fl in os.listdir(knowledges_path):
                if fl.endswith(".csv"):
                    datas.append([fl,
                                  round(os.path.getsize(knowledges_path + os.sep + fl)/1024, 2), 
                                  datetime.fromtimestamp(os.path.getmtime(knowledges_path + os.sep + fl)).strftime('%Y-%m-%d %H:%M:%S')
                    ])
            if len(datas) > 0:
                df_csv = gradio.DataFrame(
                    headers=["文件名", "文件大小(KB)", "更新时间"], datatype=["str", "number", "str"],
                    row_count=len(datas), col_count=(3, "fixed"), value=datas
                )
                with gradio.Row():
                    btn_re_vector = gradio.Button(value="重新向量化到数据库", interactive=True)
                    btn_file_upload = gradio.UploadButton(label="上传文件到知识库", file_count="multiple")

    app.launch(server_name="0.0.0.0")