from langchain.document_loaders import (
    CSVLoader,
    PDFMinerLoader,
    TextLoader,
    UnstructuredEPubLoader,
    UnstructuredHTMLLoader,
    UnstructuredMarkdownLoader,
    UnstructuredPowerPointLoader,
    UnstructuredWordDocumentLoader, JSONLoader,
)
from langchain.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
import os
import glob
from args import args
from langchain.vectorstores import FAISS
def embed_data(embeddings_model_name,original_data_path,preprocessed_data_path):

    embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)  # embedding模型

    if does_vectorstore_exist(preprocessed_data_path):
        print('++++++++++++++++++++++++++++启用本地知识库+++++++++++++++++++++++++++++++++')
        # 如果存在，加载向量库。
        #db = Chroma(persist_directory=preprocessed_data_path, embedding_function=embeddings)
        db = FAISS.load_local(preprocessed_data_path,embeddings)
    else:
        print('++++++++++++++++++++++++++++保存本地知识库+++++++++++++++++++++++++++++++++')
        # 如果不存在，就保存在本地
        texts = load_data(original_data_path)  # 文本切分
        #db = Chroma.from_documents(texts, embeddings, persist_directory=preprocessed_data_path)
        db = FAISS.from_documents(texts, embeddings)
        db.save_local(preprocessed_data_path)
    return db

def load_data(original_data_path):
    data_map = {
        ".csv": CSVLoader,
        ".json": JSONLoader,
        ".doc": UnstructuredWordDocumentLoader,
        ".docx": UnstructuredWordDocumentLoader,
        ".txt": TextLoader,
        ".md": UnstructuredMarkdownLoader,
        ".epub": UnstructuredEPubLoader,
        ".pdf": PDFMinerLoader,
        ".ppt": UnstructuredPowerPointLoader,
        ".pptx": UnstructuredPowerPointLoader,
        ".html": UnstructuredHTMLLoader,
    }
    data_all = []

    for _, (key, value) in enumerate(data_map.items()):
        if key == ".csv":
            loader_kwargs = {'encoding': 'gbk'}
            loader = DirectoryLoader(original_data_path,
                                     glob=f"*{key}",
                                     show_progress=True,  # 显示进度条
                                     use_multithreading=True,  # 多线程
                                     loader_cls=value,
                                     silent_errors=True,  # 跳过失败加载
                                     loader_kwargs=loader_kwargs
                                     )
        elif key == ".txt":
            loader_kwargs = {'autodetect_encoding': True}
            loader = DirectoryLoader(original_data_path,
                                     glob=f"*{key}",
                                     show_progress=True,  # 显示进度条
                                     use_multithreading=True,  # 多线程
                                     loader_cls=value,
                                     silent_errors=True,  # 跳过失败加载
                                     loader_kwargs=loader_kwargs
                                     )
        else:
            loader = DirectoryLoader(original_data_path,
                                     glob=f"*{key}",
                                     show_progress=True,  # 显示进度条
                                     use_multithreading=True,  # 多线程
                                     loader_cls=value,
                                     silent_errors=True,  # 跳过失败加载
                                     )
        data_one = loader.load()
        data_all.append(data_one)
    #length_function：用于计算文本块长度的方法。chunk_size：文本块的最大尺寸.chunk_overlap：文本块之间的最大重叠量。
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=args.chunk_size,chunk_overlap=args.chunk_overlap,length_function=len)
    texts = []
    for one in data_all:
        if one == []:
            continue
        else:
            text = text_splitter.split_documents(one)
            texts.extend(text)

    return texts
#更新向量库
def dBgx(original_data_path,preprocessed_data_path,embeddings_model_name):
    embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)  # embedding模型
    texts = load_data(original_data_path) # 文本切分

    # db = Chroma.from_documents(texts, embeddings, persist_directory=preprocessed_data_path)
    db = FAISS.from_documents(texts, embeddings)
    db.save_local(preprocessed_data_path)
    return db

def does_vectorstore_exist(persist_directory):
    if os.path.exists(os.path.join(persist_directory, 'index.pkl')) and os.path.exists(os.path.join(persist_directory, 'index.faiss')):
        return True
    else:
        return False