import os
import re
from typing import List, Optional, Any

from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from loguru import logger
from tqdm import tqdm

from src.config import local_embedding, retrieve_proxy, chunk_overlap, chunk_size, hf_emb_model_name
from src import shared
from src.utils import excel_to_string, get_files_hash, load_pkl, save_pkl

pwd_path = os.path.abspath(os.path.dirname(__file__))


class ChineseRecursiveTextSplitter(RecursiveCharacterTextSplitter):
    """Recursive text splitter for Chinese text.
    copy from: https://github.com/chatchat-space/Langchain-Chatchat/tree/master
    """

    def __init__(
            self,
            separators: Optional[List[str]] = None,
            keep_separator: bool = True,
            is_separator_regex: bool = True,
            **kwargs: Any,
    ) -> None:
        """Create a new TextSplitter."""
        super().__init__(keep_separator=keep_separator, **kwargs)
        self._separators = separators or [
            "\n\n",
            "\n",
            "。|！|？",
            "\.\s|\!\s|\?\s",
            "；|;\s",
            "，|,\s"
        ]
        self._is_separator_regex = is_separator_regex

    @staticmethod
    def _split_text_with_regex_from_end(
            text: str, separator: str, keep_separator: bool
    ) -> List[str]:
        # Now that we have the separator, split the text
        if separator:
            if keep_separator:
                # The parentheses in the pattern keep the delimiters in the result.
                _splits = re.split(f"({separator})", text)
                splits = ["".join(i) for i in zip(_splits[0::2], _splits[1::2])]
                if len(_splits) % 2 == 1:
                    splits += _splits[-1:]
            else:
                splits = re.split(separator, text)
        else:
            splits = list(text)
        return [s for s in splits if s != ""]

    def _split_text(self, text: str, separators: List[str]) -> List[str]:
        """Split incoming text and return chunks."""
        final_chunks = []
        # Get appropriate separator to use
        separator = separators[-1]
        new_separators = []
        for i, _s in enumerate(separators):
            _separator = _s if self._is_separator_regex else re.escape(_s)
            if _s == "":
                separator = _s
                break
            if re.search(_separator, text):
                separator = _s
                new_separators = separators[i + 1:]
                break

        _separator = separator if self._is_separator_regex else re.escape(separator)
        splits = self._split_text_with_regex_from_end(text, _separator, self._keep_separator)

        # Now go merging things, recursively splitting longer texts.
        _good_splits = []
        _separator = "" if self._keep_separator else separator
        for s in splits:
            if self._length_function(s) < self._chunk_size:
                _good_splits.append(s)
            else:
                if _good_splits:
                    merged_text = self._merge_splits(_good_splits, _separator)
                    final_chunks.extend(merged_text)
                    _good_splits = []
                if not new_separators:
                    final_chunks.append(s)
                else:
                    other_info = self._split_text(s, new_separators)
                    final_chunks.extend(other_info)
        if _good_splits:
            merged_text = self._merge_splits(_good_splits, _separator)
            final_chunks.extend(merged_text)
        return [re.sub(r"\n{2,}", "\n", chunk.strip()) for chunk in final_chunks if chunk.strip() != ""]


def get_documents(file_paths):
    text_splitter = ChineseRecursiveTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)

    documents = []
    logger.debug("Loading documents...")
    logger.debug(f"file_paths: {file_paths}")
    for file in file_paths:
        filepath = file.name
        filename = os.path.basename(filepath)
        file_type = os.path.splitext(filename)[1]
        logger.info(f"loading file: {filename}")
        texts = None
        try:
            if file_type == ".pdf":
                import PyPDF2
                logger.debug("Loading PDF...")
                try:
                    from src.pdf_func import parse_pdf
                    from src.config import advance_docs

                    two_column = advance_docs["pdf"].get("two_column", False)
                    pdftext = parse_pdf(filepath, two_column).text
                except:
                    pdftext = ""
                    with open(filepath, "rb") as pdfFileObj:
                        pdfReader = PyPDF2.PdfReader(pdfFileObj)
                        for page in tqdm(pdfReader.pages):
                            pdftext += page.extract_text()
                texts = [Document(page_content=pdftext,
                                  metadata={"source": filepath})]
            elif file_type == ".docx":
                logger.debug("Loading Word...")
                from langchain.document_loaders import UnstructuredWordDocumentLoader
                loader = UnstructuredWordDocumentLoader(filepath)
                texts = loader.load()
            elif file_type == ".pptx":
                logger.debug("Loading PowerPoint...")
                from langchain.document_loaders import UnstructuredPowerPointLoader
                loader = UnstructuredPowerPointLoader(filepath)
                texts = loader.load()
            elif file_type == ".epub":
                logger.debug("Loading EPUB...")
                from langchain.document_loaders import UnstructuredEPubLoader
                loader = UnstructuredEPubLoader(filepath)
                texts = loader.load()
            elif file_type == ".xlsx":
                logger.debug("Loading Excel...")
                text_list = excel_to_string(filepath)
                texts = []
                for elem in text_list:
                    texts.append(Document(page_content=elem,
                                          metadata={"source": filepath}))
            else:
                logger.debug("Loading text file...")
                from langchain_community.document_loaders import TextLoader
                loader = TextLoader(filepath, "utf8")
                texts = loader.load()
            logger.debug(f"text size: {len(texts)}, text top3: {texts[:3]}")
        except Exception as e:
            logger.error(f"Error loading file: {filename}, {e}")

        if texts is not None:
            texts = text_splitter.split_documents(texts)
            documents.extend(texts)
    logger.debug(f"Documents loaded. documents size: {len(documents)}, top3: {documents[:3]}")
    return documents


def construct_index(
        api_key,
        files,
        load_from_cache_if_possible=True,
):
    from langchain_community.vectorstores import FAISS
    from langchain.embeddings.huggingface import HuggingFaceEmbeddings
    if api_key:
        os.environ["OPENAI_API_KEY"] = api_key
    else:
        os.environ["OPENAI_API_KEY"] = "sk-xxxxxxx"
    index_name = get_files_hash(files)
    index_dir = os.path.join(pwd_path, '../index')
    index_path = f"{index_dir}/{index_name}"
    doc_file = f"{index_path}/docs.pkl"
    if local_embedding:
        embeddings = HuggingFaceEmbeddings(model_name=hf_emb_model_name)
    else:
        from langchain_community.embeddings import OpenAIEmbeddings
        if os.environ.get("OPENAI_API_TYPE", "openai") == "openai":
            embeddings = OpenAIEmbeddings(
                openai_api_base=shared.state.openai_api_base,
                openai_api_key=os.environ.get("OPENAI_EMBEDDING_API_KEY", api_key)
            )
        else:
            embeddings = OpenAIEmbeddings(
                deployment=os.environ["AZURE_EMBEDDING_DEPLOYMENT_NAME"],
                openai_api_key=os.environ["AZURE_OPENAI_API_KEY"],
                model=os.environ["AZURE_EMBEDDING_MODEL_NAME"],
                openai_api_base=os.environ["AZURE_OPENAI_API_BASE_URL"],
                openai_api_type="azure"
            )
    if os.path.exists(index_path) and load_from_cache_if_possible:
        logger.info("找到了缓存的索引文件，加载中……")
        index = FAISS.load_local(index_path, embeddings)
        documents = load_pkl(doc_file)
        return index, documents
    else:
        try:
            documents = get_documents(files)
            logger.info("构建索引中……")
            with retrieve_proxy():
                index = FAISS.from_documents(documents, embeddings)
            logger.debug("索引构建完成！")
            os.makedirs(index_dir, exist_ok=True)
            index.save_local(index_path)
            logger.debug("索引已保存至本地!")
            save_pkl(documents, doc_file)
            logger.debug("索引文档已保存至本地!")
            return index, documents
        except Exception as e:
            logger.error(f"索引构建失败！error: {e}")
            return None

langchain.document_loaders.UnstructuredWordDocumentLoader langchain.embeddings.huggingface.HuggingFaceEmbeddings langchain_community.vectorstores.FAISS.from_documents langchain.document_loaders.UnstructuredPowerPointLoader langchain.document_loaders.UnstructuredEPubLoader langchain.schema.Document langchain_community.vectorstores.FAISS.load_local langchain_community.document_loaders.TextLoader langchain_community.embeddings.OpenAIEmbeddings
[((440, 465), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (455, 465), False, 'import os\n'), ((3874, 3910), 'loguru.logger.debug', 'logger.debug', (['"""Loading documents..."""'], {}), "('Loading documents...')\n", (3886, 3910), False, 'from loguru import logger\n'), ((3915, 3956), 'loguru.logger.debug', 'logger.debug', (['f"""file_paths: {file_paths}"""'], {}), "(f'file_paths: {file_paths}')\n", (3927, 3956), False, 'from loguru import logger\n'), ((7187, 7208), 'src.utils.get_files_hash', 'get_files_hash', (['files'], {}), '(files)\n', (7201, 7208), False, 'from src.utils import excel_to_string, get_files_hash, load_pkl, save_pkl\n'), ((7225, 7259), 'os.path.join', 'os.path.join', (['pwd_path', '"""../index"""'], {}), "(pwd_path, '../index')\n", (7237, 7259), False, 'import os\n'), ((4033, 4059), 'os.path.basename', 'os.path.basename', (['filepath'], {}), '(filepath)\n', (4049, 4059), False, 'import os\n'), ((4118, 4158), 'loguru.logger.info', 'logger.info', (['f"""loading file: {filename}"""'], {}), "(f'loading file: {filename}')\n", (4129, 4158), False, 'from loguru import logger\n'), ((7390, 7441), 'langchain.embeddings.huggingface.HuggingFaceEmbeddings', 'HuggingFaceEmbeddings', ([], {'model_name': 'hf_emb_model_name'}), '(model_name=hf_emb_model_name)\n', (7411, 7441), False, 'from langchain.embeddings.huggingface import HuggingFaceEmbeddings\n'), ((8186, 8212), 'os.path.exists', 'os.path.exists', (['index_path'], {}), '(index_path)\n', (8200, 8212), False, 'import os\n'), ((8254, 8285), 'loguru.logger.info', 'logger.info', (['"""找到了缓存的索引文件，加载中……"""'], {}), "('找到了缓存的索引文件，加载中……')\n", (8265, 8285), False, 'from loguru import logger\n'), ((8302, 8342), 'langchain_community.vectorstores.FAISS.load_local', 'FAISS.load_local', (['index_path', 'embeddings'], {}), '(index_path, embeddings)\n', (8318, 8342), False, 'from langchain_community.vectorstores import FAISS\n'), ((8363, 8381), 'src.utils.load_pkl', 'load_pkl', (['doc_file'], {}), '(doc_file)\n', (8371, 8381), False, 'from src.utils import excel_to_string, get_files_hash, load_pkl, save_pkl\n'), ((2427, 2454), 're.search', 're.search', (['_separator', 'text'], {}), '(_separator, text)\n', (2436, 2454), False, 'import re\n'), ((2626, 2646), 're.escape', 're.escape', (['separator'], {}), '(separator)\n', (2635, 2646), False, 'import re\n'), ((4080, 4106), 'os.path.splitext', 'os.path.splitext', (['filename'], {}), '(filename)\n', (4096, 4106), False, 'import os\n'), ((7531, 7574), 'os.environ.get', 'os.environ.get', (['"""OPENAI_API_TYPE"""', '"""openai"""'], {}), "('OPENAI_API_TYPE', 'openai')\n", (7545, 7574), False, 'import os\n'), ((7829, 8098), 'langchain_community.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'deployment': "os.environ['AZURE_EMBEDDING_DEPLOYMENT_NAME']", 'openai_api_key': "os.environ['AZURE_OPENAI_API_KEY']", 'model': "os.environ['AZURE_EMBEDDING_MODEL_NAME']", 'openai_api_base': "os.environ['AZURE_OPENAI_API_BASE_URL']", 'openai_api_type': '"""azure"""'}), "(deployment=os.environ['AZURE_EMBEDDING_DEPLOYMENT_NAME'],\n    openai_api_key=os.environ['AZURE_OPENAI_API_KEY'], model=os.environ[\n    'AZURE_EMBEDDING_MODEL_NAME'], openai_api_base=os.environ[\n    'AZURE_OPENAI_API_BASE_URL'], openai_api_type='azure')\n", (7845, 8098), False, 'from langchain_community.embeddings import OpenAIEmbeddings\n'), ((8494, 8516), 'loguru.logger.info', 'logger.info', (['"""构建索引中……"""'], {}), "('构建索引中……')\n", (8505, 8516), False, 'from loguru import logger\n'), ((8632, 8655), 'loguru.logger.debug', 'logger.debug', (['"""索引构建完成！"""'], {}), "('索引构建完成！')\n", (8644, 8655), False, 'from loguru import logger\n'), ((8668, 8705), 'os.makedirs', 'os.makedirs', (['index_dir'], {'exist_ok': '(True)'}), '(index_dir, exist_ok=True)\n', (8679, 8705), False, 'import os\n'), ((8759, 8784), 'loguru.logger.debug', 'logger.debug', (['"""索引已保存至本地!"""'], {}), "('索引已保存至本地!')\n", (8771, 8784), False, 'from loguru import logger\n'), ((8797, 8826), 'src.utils.save_pkl', 'save_pkl', (['documents', 'doc_file'], {}), '(documents, doc_file)\n', (8805, 8826), False, 'from src.utils import excel_to_string, get_files_hash, load_pkl, save_pkl\n'), ((8839, 8866), 'loguru.logger.debug', 'logger.debug', (['"""索引文档已保存至本地!"""'], {}), "('索引文档已保存至本地!')\n", (8851, 8866), False, 'from loguru import logger\n'), ((1595, 1627), 're.split', 're.split', (['f"""({separator})"""', 'text'], {}), "(f'({separator})', text)\n", (1603, 1627), False, 'import re\n'), ((1837, 1862), 're.split', 're.split', (['separator', 'text'], {}), '(separator, text)\n', (1845, 1862), False, 'import re\n'), ((2320, 2333), 're.escape', 're.escape', (['_s'], {}), '(_s)\n', (2329, 2333), False, 'import re\n'), ((4275, 4305), 'loguru.logger.debug', 'logger.debug', (['"""Loading PDF..."""'], {}), "('Loading PDF...')\n", (4287, 4305), False, 'from loguru import logger\n'), ((6526, 6578), 'loguru.logger.error', 'logger.error', (['f"""Error loading file: {filename}, {e}"""'], {}), "(f'Error loading file: {filename}, {e}')\n", (6538, 6578), False, 'from loguru import logger\n'), ((8534, 8550), 'src.config.retrieve_proxy', 'retrieve_proxy', ([], {}), '()\n', (8548, 8550), False, 'from src.config import local_embedding, retrieve_proxy, chunk_overlap, chunk_size, hf_emb_model_name\n'), ((8576, 8619), 'langchain_community.vectorstores.FAISS.from_documents', 'FAISS.from_documents', (['documents', 'embeddings'], {}), '(documents, embeddings)\n', (8596, 8619), False, 'from langchain_community.vectorstores import FAISS\n'), ((8946, 8980), 'loguru.logger.error', 'logger.error', (['f"""索引构建失败！error: {e}"""'], {}), "(f'索引构建失败！error: {e}')\n", (8958, 8980), False, 'from loguru import logger\n'), ((4910, 4971), 'langchain.schema.Document', 'Document', ([], {'page_content': 'pdftext', 'metadata': "{'source': filepath}"}), "(page_content=pdftext, metadata={'source': filepath})\n", (4918, 4971), False, 'from langchain.schema import Document\n'), ((5062, 5093), 'loguru.logger.debug', 'logger.debug', (['"""Loading Word..."""'], {}), "('Loading Word...')\n", (5074, 5093), False, 'from loguru import logger\n'), ((5205, 5245), 'langchain.document_loaders.UnstructuredWordDocumentLoader', 'UnstructuredWordDocumentLoader', (['filepath'], {}), '(filepath)\n', (5235, 5245), False, 'from langchain.document_loaders import UnstructuredWordDocumentLoader\n'), ((7724, 7775), 'os.environ.get', 'os.environ.get', (['"""OPENAI_EMBEDDING_API_KEY"""', 'api_key'], {}), "('OPENAI_EMBEDDING_API_KEY', api_key)\n", (7738, 7775), False, 'import os\n'), ((4547, 4578), 'src.pdf_func.parse_pdf', 'parse_pdf', (['filepath', 'two_column'], {}), '(filepath, two_column)\n', (4556, 4578), False, 'from src.pdf_func import parse_pdf\n'), ((5339, 5376), 'loguru.logger.debug', 'logger.debug', (['"""Loading PowerPoint..."""'], {}), "('Loading PowerPoint...')\n", (5351, 5376), False, 'from loguru import logger\n'), ((5486, 5524), 'langchain.document_loaders.UnstructuredPowerPointLoader', 'UnstructuredPowerPointLoader', (['filepath'], {}), '(filepath)\n', (5514, 5524), False, 'from langchain.document_loaders import UnstructuredPowerPointLoader\n'), ((4738, 4766), 'PyPDF2.PdfReader', 'PyPDF2.PdfReader', (['pdfFileObj'], {}), '(pdfFileObj)\n', (4754, 4766), False, 'import PyPDF2\n'), ((4803, 4824), 'tqdm.tqdm', 'tqdm', (['pdfReader.pages'], {}), '(pdfReader.pages)\n', (4807, 4824), False, 'from tqdm import tqdm\n'), ((5618, 5649), 'loguru.logger.debug', 'logger.debug', (['"""Loading EPUB..."""'], {}), "('Loading EPUB...')\n", (5630, 5649), False, 'from loguru import logger\n'), ((5753, 5785), 'langchain.document_loaders.UnstructuredEPubLoader', 'UnstructuredEPubLoader', (['filepath'], {}), '(filepath)\n', (5775, 5785), False, 'from langchain.document_loaders import UnstructuredEPubLoader\n'), ((5879, 5911), 'loguru.logger.debug', 'logger.debug', (['"""Loading Excel..."""'], {}), "('Loading Excel...')\n", (5891, 5911), False, 'from loguru import logger\n'), ((5940, 5965), 'src.utils.excel_to_string', 'excel_to_string', (['filepath'], {}), '(filepath)\n', (5955, 5965), False, 'from src.utils import excel_to_string, get_files_hash, load_pkl, save_pkl\n'), ((6201, 6237), 'loguru.logger.debug', 'logger.debug', (['"""Loading text file..."""'], {}), "('Loading text file...')\n", (6213, 6237), False, 'from loguru import logger\n'), ((6339, 6367), 'langchain_community.document_loaders.TextLoader', 'TextLoader', (['filepath', '"""utf8"""'], {}), "(filepath, 'utf8')\n", (6349, 6367), False, 'from langchain_community.document_loaders import TextLoader\n'), ((6065, 6123), 'langchain.schema.Document', 'Document', ([], {'page_content': 'elem', 'metadata': "{'source': filepath}"}), "(page_content=elem, metadata={'source': filepath})\n", (6073, 6123), False, 'from langchain.schema import Document\n')]