import stanza
import re
import gc
from langid.langid import LanguageIdentifier, model
from concurrent.futures import ThreadPoolExecutor, as_completed


class TextProcessor:
    def __init__(self, model_dir=None, use_gpu=False, num_threads=4):
        """
        初始化 TextProcessor 实例。

        参数：
        - model_dir: 模型存储和加载的路径。如果为 None，则使用默认路径。
        - use_gpu: 是否使用 GPU。默认为 False。
        - num_threads: 并行处理的线程数量。默认为 4。
        """
        self.model_dir = model_dir
        self.use_gpu = use_gpu
        self.num_threads = num_threads
        self.nlp_models = {}
        self.supported_languages = ["zh", "ja", "en"]
        self.language_mapping = {"zh": "zh", "ja": "ja", "en": "en"}
        self._initialize_models()
        # 初始化 langid
        self.identifier = LanguageIdentifier.from_modelstring(model, norm_probs=True)
        # 语言检测缓存
        self.lang_cache = {}

    def download_models(self):
        """
        下载所需的 Stanza 语言模型。
        """
        for lang in self.supported_languages:
            try:
                stanza.download(lang, model_dir=self.model_dir, verbose=False)
            except Exception as e:
                print(f"下载语言模型 '{lang}' 失败: {e}")

    def process_text(self, text):
        """
        处理输入的文本，提取其中的名词和动词。

        参数：
        - text: 要处理的字符串文本。

        返回：
        - filtered_tokens: 包含名词和动词的词语列表。
        """
        if not isinstance(text, str):
            raise ValueError("输入的文本必须是字符串类型。")

        # 根据语言变化分割文本
        segments = self._split_text_by_language(text)

        filtered_tokens = []

        # 并行处理每个片段
        with ThreadPoolExecutor(max_workers=self.num_threads) as executor:
            futures = {
                executor.submit(self._process_segment, lang, segment): (lang, segment)
                for lang, segment in segments
            }
            for future in as_completed(futures):
                try:
                    tokens = future.result()
                    filtered_tokens.extend(tokens)
                except Exception as e:
                    lang, segment = futures[future]
                    print(f"处理语言 '{lang}' 的文本片段时出错: {e}")

        return filtered_tokens

    def _process_segment(self, lang, segment):
        """
        处理单个语言片段，提取名词和动词。

        参数：
        - lang: 语言代码。
        - segment: 文本片段。

        返回：
        - tokens: 提取的名词和动词列表。
        """
        tokens = []
        # 检查模型是否已加载
        if lang not in self.nlp_models or self.nlp_models[lang] is None:
            print(f"语言模型 '{lang}' 未加载，跳过处理。")
            return tokens
        try:
            doc = self.nlp_models[lang](segment)
            # 提取名词和动词
            for sent in doc.sentences:
                for word in sent.words:
                    if word.upos in ["NOUN", "VERB"]:
                        tokens.append(word.text)
        except Exception as e:
            print(f"处理语言 '{lang}' 的文本时出错: {e}")
        return tokens

    def close(self):
        """
        释放模型资源，节省内存。
        """
        self.nlp_models.clear()
        # 清空缓存
        self.lang_cache.clear()
        # 显式调用垃圾回收
        gc.collect()
        print("已释放模型资源。")

    def _detect_language(self, text):
        """
        检测给定文本的语言。

        参数：
        - text: 要检测的文本。

        返回：
        - lang: 检测到的语言代码（'zh', 'ja', 'en'），未检测到则返回 None。
        """
        # 检查缓存
        if text in self.lang_cache:
            return self.lang_cache[text]

        try:
            lang_code, prob = self.identifier.classify(text)
            # 映射到支持的语言
            lang = self.language_mapping.get(lang_code)
            if lang:
                self.lang_cache[text] = lang
                return lang
            else:
                self.lang_cache[text] = None
                return None
        except Exception:
            self.lang_cache[text] = None
            return None

    def _split_text_by_language(self, text):
        """
        根据语言变化分割文本。

        参数：
        - text: 要分割的文本。

        返回：
        - segments: 列表，每个元素为 (lang, segment) 格式的元组。
        """
        segments = []
        current_lang = None
        current_segment = ""
        idx = 0
        text_length = len(text)
        window_size = 10  # 窗口大小

        while idx < text_length:
            # 取窗口内的文本
            end_idx = min(idx + window_size, text_length)
            sample_text = text[idx:end_idx]
            lang = self._detect_language(sample_text)
            if not lang:
                # 无法检测语言，假设与当前语言相同
                lang = current_lang

            if lang != current_lang:
                if current_segment:
                    segments.append((current_lang, current_segment))
                current_segment = text[idx]
                current_lang = lang
            else:
                current_segment += text[idx]
            idx += 1

        if current_segment:
            segments.append((current_lang, current_segment))

        # 合并相邻相同语言的片段
        merged_segments = []
        for lang, segment in segments:
            if not lang:
                continue
            if not merged_segments or merged_segments[-1][0] != lang:
                merged_segments.append((lang, segment))
            else:
                merged_segments[-1] = (lang, merged_segments[-1][1] + segment)

        return merged_segments

    def _initialize_models(self):
        """
        初始化各语言的 Stanza 处理管线。
        """
        for lang in self.supported_languages:
            try:
                self.nlp_models[lang] = stanza.Pipeline(
                    lang,
                    processors="tokenize,pos",
                    dir=self.model_dir,
                    verbose=False,
                    use_gpu=self.use_gpu,
                )
            except Exception as e:
                self.nlp_models[lang] = None
                print(f"加载语言模型 '{lang}' 时出错: {e}")
-----
分析上买的代码，语种检测能不能快，更准一些