import os.path
from server.knowledge_base.utils import (get_file_path, get_vs_path,
                                         refresh_vs_cache, load_embeddings)
from configs.model_config import (embedding_model_dict, EMBEDDING_MODEL, EMBEDDING_DEVICE)
from langchain.vectorstores import FAISS
from server.utils import torch_gc


class KnowledgeFile:
    def __init__(
            self,
            filename: str,
            knowledge_base_name: str
    ):
        self.knowledge_base_name = knowledge_base_name
        self.knowledge_base_type = "faiss"
        self.filename = filename
        self.ext = os.path.splitext(filename)[-1]
        self.filepath = get_file_path(knowledge_base_name, filename)
        self.docs = None

    def file2text(self):
        if self.ext in []:
            from langchain.document_loaders import UnstructuredFileLoader
            loader = UnstructuredFileLoader(self.filepath)
        elif self.ext in []:
            pass

        from langchain.text_splitter import CharacterTextSplitter
        text_splitter = CharacterTextSplitter(chunk_size=500, chunk_overlap=200)
        self.docs = loader.load_and_split(text_splitter)
        return True

    def docs2vs(self):
        vs_path = get_vs_path(self.knowledge_base_name)
        embeddings = load_embeddings(embedding_model_dict[EMBEDDING_MODEL], EMBEDDING_DEVICE)
        if os.path.exists(vs_path) and "index.faiss" in os.listdir(vs_path):
            vector_store = FAISS.load_local(vs_path, embeddings)
            vector_store.add_documents(self.docs)
            torch_gc()
        else:
            if not os.path.exists(vs_path):
                os.makedirs(vs_path)
            vector_store = FAISS.from_documents(self.docs, embeddings)  # docs 为Document列表
            torch_gc()
        vector_store.save_local(vs_path)
        refresh_vs_cache(self.knowledge_base_name)
        return True