# /*
#  * Copyright 2024 KylinSoft Co., Ltd.
#  *
#  * This program is free software: you can redistribute it and/or modify it under
#  * the terms of the GNU General Public License as published by the Free Software
#  * Foundation, either version 3 of the License, or (at your option) any later
#  * version.
#  *
#  * This program is distributed in the hope that it will be useful, but WITHOUT
#  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
#  * FOR A PARTICULAR PURPOSE. See the GNU General Public License for more
#  * details.
#  *
#  * You should have received a copy of the GNU General Public License along with
#  * this program. If not, see <https://www.gnu.org/licenses/>.
#  */

import re
from base_chunker import BaseChunker

class TextChunker(BaseChunker):
    def is_english(self,text):
        return all(ord(c) < 128 for c in text)
    def apply_regex(self, text) -> str:
        newline_regex = r"\n{3,}"
        space_regex = r"\s+"
        double_newline_regex = r"\n\n"
        single_punctuation_regex = r"([;；.!?。！？\?])([^”’'])"
        ellipsis_regex = r"(\\.\\.\\.)([^\"’”」』])"
        chinese_ellipsis_regex = r"(…{2})([^\"’”」』])"
        end_quote_punctuation_regex = r"([;；!?。！？\?][\"’”」』]{0,2})([^;；!?，。！？\?])"

        # 处理多余的换行符
        text = re.sub(newline_regex, "\n", text)

        # 标准化空格
        text = re.sub(space_regex, " ", text)

        # 移除双换行
        text = re.sub(double_newline_regex, "", text)

        # 处理标点符号后的换行
        text = re.sub(single_punctuation_regex, r"\1\n\2", text)

        # 处理省略号后的换行
        text = re.sub(ellipsis_regex, r"\1\n\2", text)

        # 处理中文省略号后的换行
        text = re.sub(chinese_ellipsis_regex, r"\1\n\2", text)

        # 处理结尾引号后的标点
        text = re.sub(end_quote_punctuation_regex, r"\1\n\2", text)

        # 删除所有空白字符
        if self.is_english(text):
            text = re.sub(r"\s+", " ", text)
        else:
            text = re.sub(r"\s+", "", text)

        return text

    def split_sentence_into_chunks(self,sentence, max_chunk_size=100):
        punctuation = r'.+?[.?!,;，。！？：]+'

        chunks = []
        current_chunk = ""
        for part in re.split(f"({punctuation})", sentence):
            part = part.strip()
            if not part:
                continue
            if len(current_chunk) + len(part) <= max_chunk_size:
                current_chunk += part
            else:
                if current_chunk:
                    chunks.append(current_chunk.strip())
                current_chunk = part
        if current_chunk:
            chunks.append(current_chunk.strip())

        return chunks
    
    def split_sentence_length(self, sentence, max_chunk_size=100):
        sentence_list = []
        for sen in sentence:
            sentence_list.extend([sen[i:i+max_chunk_size] for i in range(0, len(sen), max_chunk_size)])
        return sentence_list
    
    def segment_tokenize(self, texts, max_length) -> list:
        if isinstance(texts, str):
            texts = [texts]  # 将字符串转换为列表
        result_list = []
        for text in texts:
            sentence = self.split_sentence_into_chunks(sentence=text,max_chunk_size=max_length)
            sentence_list = self.split_sentence_length(sentence=sentence,max_chunk_size=max_length)
            result_list.extend(sentence_list)
        return result_list

# if __name__ == "__main__":
#     text = ""
#     textchunker = TextChunker()
#     textchunker.segment_tokenize(text, 500)