"""
中文分词器模块
用于BM25检索的自定义中文分词功能
支持中英文混合文本的分词处理
"""
import re
from typing import List
import jieba

# 尝试导入Stemmer，如果失败则使用备用方案
try:
    import Stemmer
    HAS_STEMMER = True
except ImportError:
    HAS_STEMMER = False
    Stemmer = None


class JiebaTokenizer:
    """基于jieba的中文分词器，支持停用词过滤和词干提取"""
    
    def __init__(self, language: str = "chinese", use_stemmer: bool = True):
        """
        初始化分词器
        
        Args:
            language: 语言类型，支持 "chinese" 和 "english"
            use_stemmer: 是否使用词干提取（仅对英文有效）
        """
        self.language = language.lower()
        self.use_stemmer = use_stemmer and HAS_STEMMER
        
        # 初始化英文词干提取器
        if self.use_stemmer and self.language == "english":
            self.stemmer = Stemmer.Stemmer("english")
        else:
            self.stemmer = None
        
        # 加载中文停用词
        self.stopwords = self._load_stopwords()
    
    def _load_stopwords(self) -> set:
        """
        加载停用词表
        
        Returns:
            set: 停用词集合
        """
        stopwords = set()
        
        # 中文常见停用词
        chinese_stopwords = {
            "的", "了", "在", "是", "我", "有", "和", "就", "不", "人", "都", "一", "一个",
            "上", "也", "很", "到", "说", "要", "去", "你", "会", "着", "没有", "看", "好",
            "自己", "这", "那", "里", "为", "把", "被", "给", "让", "从", "向", "以", "可以",
            "能", "但", "而", "与", "或", "因为", "所以", "如果", "为了", "对", "于", "由",
            "及", "对于", "关于", "根据", "按照", "通过", "之", "得", "地", "个", "等", "并",
        }
        
        # 英文常见停用词
        english_stopwords = {
            "a", "an", "and", "are", "as", "at", "be", "but", "by", "for", "if", "in",
            "into", "is", "it", "no", "not", "of", "on", "or", "such", "that", "the",
            "their", "then", "there", "these", "they", "this", "to", "was", "will", "with",
        }
        
        if self.language == "chinese":
            stopwords.update(chinese_stopwords)
            stopwords.update(english_stopwords)  # 中文也可能包含英文
        elif self.language == "english":
            stopwords.update(english_stopwords)
        
        return stopwords
    
    def tokenize(self, text: str) -> List[str]:
        """
        对文本进行分词
        
        Args:
            text: 输入文本
            
        Returns:
            List[str]: 分词结果列表
        """
        if not text or not isinstance(text, str):
            return []
        
        # 转换为小写
        text = text.lower().strip()
        
        if self.language == "chinese":
            return self._tokenize_chinese(text)
        elif self.language == "english":
            return self._tokenize_english(text)
        else:
            # 默认使用中文分词
            return self._tokenize_chinese(text)
    
    def _tokenize_chinese(self, text: str) -> List[str]:
        """
        中文分词（支持中英文混合）
        
        Args:
            text: 输入文本
            
        Returns:
            List[str]: 分词结果
        """
        # 使用jieba进行中文分词
        words = jieba.cut(text)
        
        # 过滤停用词和无效词
        tokens = []
        for word in words:
            word = word.strip()
            
            # 跳过空白和单字符词（除非是数字或英文单词）
            if not word:
                continue
            
            # 跳过停用词
            if word in self.stopwords:
                continue
            
            # 跳过纯标点符号
            if re.match(r'^[^\w\s]+$', word):
                continue
            
            # 如果是英文单词且启用了词干提取
            if self.stemmer and re.match(r'^[a-z]+$', word):
                word = self.stemmer.stemWord(word)
            
            tokens.append(word)
        
        return tokens
    
    def _tokenize_english(self, text: str) -> List[str]:
        """
        英文分词
        
        Args:
            text: 输入文本
            
        Returns:
            List[str]: 分词结果
        """
        # 使用正则表达式分割英文单词
        words = re.findall(r'\b[a-z]+\b', text)
        
        # 过滤停用词
        tokens = []
        for word in words:
            if word not in self.stopwords:
                # 词干提取
                if self.stemmer:
                    word = self.stemmer.stemWord(word)
                tokens.append(word)
        
        return tokens
    
    def __call__(self, text: str) -> List[str]:
        """
        使分词器可直接调用
        
        Args:
            text: 输入文本
            
        Returns:
            List[str]: 分词结果
        """
        return self.tokenize(text)


def create_tokenizer(language: str = "chinese", use_stemmer: bool = True) -> JiebaTokenizer:
    """
    创建分词器的便捷函数
    
    Args:
        language: 语言类型
        use_stemmer: 是否使用词干提取
        
    Returns:
        JiebaTokenizer: 分词器实例
    """
    return JiebaTokenizer(language=language, use_stemmer=use_stemmer)
