from pypinyin import lazy_pinyin
import re
import json
import logging
from collections import defaultdict
import os

# 日志配置（保留但降低调试级别的输出频率）
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(funcName)s - %(message)s",
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

# -------------------------- 1. 嵌入拼音前缀树的JSON（关键！） --------------------------
# 替换为你生成的3万+成语的拼音树JSON字符串
IDIOM_TRIE_JSON = ""  # 请粘贴你的JSON内容

# -------------------------- 2. 工具函数优化（减少重复计算） --------------------------
def chinese_to_pinyin(text: str) -> str:
    """将纯汉字转为无声调拼音字符串，非纯汉字返回空（添加缓存）"""
    if not isinstance(text, str) or not re.match(r'^[\u4e00-\u9fa5]+$', text):
        return ""
    # 单次调用完成拼音转换，避免循环内多次调用
    return ''.join(lazy_pinyin(text, style=0))

def calculate_match_rate(pinyin1: str, pinyin2: str) -> float:
    """优化匹配度计算：使用编辑距离算法，更好地处理字符插入、删除和替换的情况"""
    if not pinyin1 or not pinyin2:
        return 0.0
    
    len1, len2 = len(pinyin1), len(pinyin2)
    
    # 创建编辑距离矩阵
    dp = [[0] * (len2 + 1) for _ in range(len1 + 1)]
    
    # 初始化矩阵边界
    for i in range(len1 + 1):
        dp[i][0] = i
    for j in range(len2 + 1):
        dp[0][j] = j
    
    # 计算编辑距离
    for i in range(1, len1 + 1):
        for j in range(1, len2 + 1):
            if pinyin1[i-1] == pinyin2[j-1]:
                dp[i][j] = dp[i-1][j-1]  # 字符相同，无需操作
            else:
                # 选择插入、删除或替换中代价最小的操作
                dp[i][j] = min(
                    dp[i][j-1] + 1,      # 插入
                    dp[i-1][j] + 1,      # 删除
                    dp[i-1][j-1] + 1     # 替换
                )
    
    # 计算最大可能的编辑距离（两个字符串长度之和）
    max_distance = len1 + len2
    
    # 将编辑距离转换为匹配率（0.0-1.0）
    # 匹配率 = 1 - (编辑距离 / 最大可能距离)
    match_rate = 1 - (dp[len1][len2] / max_distance)
    
    return match_rate

# -------------------------- 3. 前缀树节点类（轻量化改造） --------------------------
class TrieNode:
    def __init__(self, node_dict: dict):
        self.children = node_dict["children"]  # 直接存储字典，避免递归创建对象
        self.idioms = node_dict["idioms"]

# -------------------------- 4. 检索器核心优化（预缓存+重写模糊匹配） --------------------------
class IdiomRetriever:
    def __init__(self, trie_json: str):
        """初始化：加载双树结构+预缓存所有成语的拼音映射（核心优化）"""
        # 加载包含双树的JSON
        trie_data = json.loads(trie_json)
        # 分别获取汉字树和拼音树
        self.chinese_trie_dict = trie_data.get("chinese_trie", {})
        self.pinyin_trie_dict = trie_data.get("pinyin_trie", {})
        # 初始化拼音树的节点和缓存
        self.pinyin_root = TrieNode(self.pinyin_trie_dict)
        # 预缓存：{拼音长度: {拼音串: 成语列表}} —— 避免模糊检索遍历树
        self.idiom_pinyin_cache = defaultdict(dict)
        # 预缓存：所有成语的集合（快速去重）
        self.all_idioms = set()
        # 从前缀树中提取所有成语和拼音（仅初始化时执行一次）
        self._build_pinyin_cache(self.pinyin_trie_dict, current_pinyin="")
        logger.info(f"初始化完成，缓存成语数：{len(self.all_idioms)}，拼音长度分组：{list(self.idiom_pinyin_cache.keys())}")

    def _build_pinyin_cache(self, node_dict: dict, current_pinyin: str):
        """深度优先遍历树，预缓存成语拼音映射（仅初始化执行）"""
        # 叶子节点有成语，加入缓存
        if node_dict["idioms"]:
            pinyin_len = len(current_pinyin)
            self.idiom_pinyin_cache[pinyin_len][current_pinyin] = node_dict["idioms"]
            self.all_idioms.update(node_dict["idioms"])
        # 遍历子节点
        for char, child_dict in node_dict["children"].items():
            self._build_pinyin_cache(child_dict, current_pinyin + char)

    def retrieve_exact_chinese(self, chinese: str) -> list:
        """汉字精准检索：使用汉字树查找"""
        if not chinese:
            return []
        node_dict = self.chinese_trie_dict
        for char in chinese:
            if char not in node_dict["children"]:
                return []
            node_dict = node_dict["children"][char]
        return node_dict["idioms"]

    def retrieve_exact(self, pinyin: str) -> list:
        """拼音精准检索：优化为字典遍历，替代对象属性访问"""
        if not pinyin:
            return []
        node_dict = self.pinyin_trie_dict
        for char in pinyin:
            if char not in node_dict["children"]:
                return []
            node_dict = node_dict["children"][char]
        return node_dict["idioms"]

    def retrieve_fuzzy(self, pinyin: str, threshold: float = 0.8, max_count: int = 0) -> list:
        """模糊检索：基于预缓存的拼音映射，替代DFS遍历（核心优化）
        
        Args:
            pinyin: 要搜索的拼音字符串
            threshold: 匹配度阈值
            max_count: 最大返回结果数量，0表示不限制
            
        Returns:
            匹配的成语列表
        """
        if not pinyin or threshold <= 0:
            return []
        
        # 存储(匹配度, 成语)元组的列表，用于排序
        matched_with_scores = []
        
        # 遍历所有可能长度的拼音串（不仅限于相同长度）
        # 为了提高效率，我们只考虑与目标拼音长度相近的拼音串
        target_len = len(pinyin)
        
        # 检查长度相近的拼音串
        for length_diff in [-1, 0, 1]:
            length = target_len + length_diff
            if length not in self.idiom_pinyin_cache:
                continue
                
            # 对于每个候选拼音，使用简化的过滤条件
            for cache_pinyin, idioms in self.idiom_pinyin_cache[length].items():
                # 1. 简化过滤：只检查第一个字符是否匹配（提高效率同时保证基本相关性）
                if pinyin[0] != cache_pinyin[0]:
                    continue
                
                # 2. 精确计算匹配率（去掉了过于严格的预过滤）
                match_rate = calculate_match_rate(pinyin, cache_pinyin)
                if match_rate >= threshold:
                    # 为每个成语记录匹配度
                    for idiom in idioms:
                        matched_with_scores.append((-match_rate, idiom))  # 使用负号进行降序排序
                        
                        # 如果已经找到足够数量的结果，可以提前结束搜索
                        if max_count > 0 and len(matched_with_scores) >= max_count * 2:  # 预留一些缓冲区用于去重
                            break
                
                # 如果已经找到足够数量的结果，可以提前结束搜索
                if max_count > 0 and len(matched_with_scores) >= max_count * 2:
                    break
            
            # 如果已经找到足够数量的结果，可以提前结束搜索
            if max_count > 0 and len(matched_with_scores) >= max_count * 2:
                break
        
        # 先排序（按匹配度降序），再去重
        # 使用负号排序是为了避免自定义排序函数，提高性能
        matched_with_scores.sort()
        
        # 去重并提取成语列表
        seen = set()
        result = []
        for _, idiom in matched_with_scores:
            if idiom not in seen:
                seen.add(idiom)
                result.append(idiom)
                
                # 应用max_count限制，找到足够数量就停止
                if max_count > 0 and len(result) >= max_count:
                    break
        
        return result

    def find_idioms_in_sentence(self, sentence: str, fuzzy_threshold: float = 0.8, max_count: int = 10) -> list:
        """
        find_idioms_in_sentence 
        Args:
            sentence: 要搜索的中文句子
            fuzzy_threshold: 模糊匹配的阈值
            max_count: 最大返回结果数量
            
        Returns:
            包含精准和模糊匹配结果的成语列表（精准在前，模糊在后）
        """
        # 清洗句子：移除非汉字字符（单次处理）
        clean_sentence = re.sub(r'[^\u4e00-\u9fa5]', '', sentence)
        if len(clean_sentence) < 4:
            return []  # 无结果返回空列表
        
        # 三级匹配流程：
        # 1. 先使用汉字树匹配
        chinese_idioms = self._match_chinese_exact_in_sentence(clean_sentence)
        
        # 2. 再进行拼音精准匹配
        pinyin_idioms = self._match_pinyin_exact_in_sentence(clean_sentence)
        
        # 3. 合并所有精准匹配结果（去重）
        all_exact_idioms = chinese_idioms.union(pinyin_idioms)
        
        if all_exact_idioms:  # 如果有任何精准匹配结果
            # 排序并应用max_count限制
            result = list(all_exact_idioms)
            result.sort(key=lambda x: len(x), reverse=True)
            if max_count > 0:
                result = result[:max_count]
            logger.info(f"精准匹配成功（汉字{len(chinese_idioms)}个，拼音{len(pinyin_idioms)}个），返回{len(result)}条结果")
            return result
        
        # 2. 如果汉字树匹配失败，将汉字翻译成拼音，使用拼音树匹配
        sentence_pinyin = chinese_to_pinyin(clean_sentence)
        if sentence_pinyin:
            pinyin_idioms = self._match_pinyin_exact_in_sentence(clean_sentence)
            if pinyin_idioms:  # 如果拼音树匹配成功
                # 排序并应用max_count限制
                result = list(pinyin_idioms)
                result.sort(key=lambda x: len(x), reverse=True)
                if max_count > 0:
                    result = result[:max_count]
                logger.info(f"拼音精准匹配成功，返回{len(result)}条结果")
                return result
        
        # 3. 如果拼音树匹配失败，最后进行拼音模糊匹配
        fuzzy_results = self._match_fuzzy_in_sentence(clean_sentence, fuzzy_threshold, max_count, set())
        logger.info(f"拼音模糊匹配成功，返回{len(fuzzy_results)}条结果")
        return fuzzy_results
        
    def _match_chinese_exact_in_sentence(self, clean_sentence: str) -> set:
        """
        使用汉字Trie树多模式匹配句子中的成语
        """
        exact_idioms = set()
        sentence_len = len(clean_sentence)
        
        # 对句子中的每个字符作为起点，尝试匹配任意长度的成语
        for i in range(sentence_len):
            current_dict = self.chinese_trie_dict
            
            # 从当前位置开始，尝试匹配任意长度的成语
            for j in range(i, sentence_len):
                char = clean_sentence[j]
                
                # 检查当前汉字是否在汉字Trie树中
                if char not in current_dict["children"]:
                    break  # 如果当前汉字不在Trie树中，中断当前匹配
                
                # 移动到汉字Trie树的下一个节点
                current_dict = current_dict["children"][char]
                
                # 如果当前节点有成语，说明找到了匹配
                if current_dict["idioms"]:
                    exact_idioms.update(current_dict["idioms"])
        
        return exact_idioms
        
    def _match_pinyin_exact_in_sentence(self, clean_sentence: str) -> set:
        """
        将汉字翻译成拼音，然后使用拼音Trie树多模式匹配
        """
        exact_idioms = set()
        
        # 将句子翻译成拼音
        sentence_pinyin = chinese_to_pinyin(clean_sentence)
        if not sentence_pinyin:
            return exact_idioms
        
        pinyin_len = len(sentence_pinyin)
        
        # 对拼音串中的每个字符作为起点，尝试匹配任意长度的成语
        for i in range(pinyin_len):
            current_dict = self.pinyin_trie_dict
            
            # 从当前位置开始，尝试匹配任意长度的成语
            for j in range(i, pinyin_len):
                char = sentence_pinyin[j]
                
                # 检查当前拼音字符是否在拼音Trie树中
                if char not in current_dict["children"]:
                    break  # 如果当前拼音字符不在Trie树中，中断当前匹配
                
                # 移动到拼音Trie树的下一个节点
                current_dict = current_dict["children"][char]
                
                # 如果当前节点有成语，说明找到了匹配
                if current_dict["idioms"]:
                    exact_idioms.update(current_dict["idioms"])
        
        return exact_idioms
        
    def _match_fuzzy_in_sentence(self, clean_sentence: str, threshold: float, max_count: int, exclude_set: set) -> list:
        """
        _match_fuzzy_in_sentence对句子中的片段进行模糊搜索
        
        Args:
            clean_sentence: 清洗后的中文句子
            threshold: 模糊匹配阈值
            max_count: 最大返回结果数量
            exclude_set: 需要排除的成语集合（已精准匹配的成语）
            
        Returns:
            排序后的模糊匹配成语列表
        """
        fuzzy_idiom_set = set()  # 用于去重的模糊成语集合
        fuzzy_list = []  # 存储排序后的模糊成语列表
        sentence_len = len(clean_sentence)
        
        # 支持的成语长度范围（4-15字）
        # 优化策略：
        # 1. 从长到短匹配（15字到4字），确保长成语不会被短成语覆盖
        # 2. 这样可以保证如果输入包含多字成语，能够正确匹配到
        for length in range(15, 3, -1):  # 从15字到4字，倒序检查
            # 如果句子长度不够，跳过当前长度
            if sentence_len < length:
                continue
                
            # 检查是否已经收集了足够的结果
            if max_count > 0 and len(fuzzy_idiom_set) >= max_count:
                break
                
            # 遍历当前长度的所有可能片段
            for i in range(sentence_len - length + 1):
                # 检查是否已经收集了足够的结果
                if max_count > 0 and len(fuzzy_idiom_set) >= max_count:
                    break
                    
                fragment = clean_sentence[i:i+length]
                fragment_pinyin = chinese_to_pinyin(fragment)
                if not fragment_pinyin:
                    continue

                # 模糊检索：提取成语名称，排除已精准匹配的，加入模糊集合
                remaining_needed = max_count - len(fuzzy_idiom_set) if max_count > 0 else 0
                fuzzy_idioms = self.retrieve_fuzzy(fragment_pinyin, threshold, remaining_needed)
                
                # 过滤掉已在精准匹配中的成语，并去重
                for idiom in fuzzy_idioms:
                    if idiom not in exclude_set and idiom not in fuzzy_idiom_set:
                        fuzzy_idiom_set.add(idiom)
                        fuzzy_list.append(idiom)
                        
                        # 检查是否已经收集了足够的结果
                        if max_count > 0 and len(fuzzy_idiom_set) >= max_count:
                            break
        
        return fuzzy_list


# 模块级别初始化检索器（预缓存仅执行一次）
current_dir = os.path.abspath(os.path.dirname(__file__))
json_name = os.path.join(current_dir, "idiom_trie.json")
with open(json_name, "r", encoding="utf-8") as f:
    content = f.read()
# 创建全局检索器实例，所有search_idiom调用共享
_retriever_instance = IdiomRetriever(content)


# -------------------------- 4.暴露函数供外部调用 --------------------------
def search_idiom(sentence: str, fuzzy_threshold: float = 0.8, max_count: int = 5) -> list:
    """搜索句子中的成语，使用预初始化的检索器实例"""
    return _retriever_instance.find_idioms_in_sentence(sentence, fuzzy_threshold, max_count)
    
# -------------------------- 5. 测试检索（适配纯成语列表的输出格式） --------------------------
if __name__ == "__main__":
    # 测试句子：包含不同长度的成语
    test_sentence1 = "戎马劻勷"
    test_sentence2 = "他做事情总是三欣二意，从来没有一心一意的时候。"
    test_sentence3 = "天下兴亡匹夫有责，我们每个人都应该有这样的责任感。"
    test_sentence4 = "千里之行始于足下，不积跬步无以至千里。"

    # 识别成语（计时查看优化效果）
    import time
    start = time.time()
    result1 = search_idiom(test_sentence1)
    result2 = search_idiom(test_sentence2, fuzzy_threshold=0.8)
    result3 = search_idiom(test_sentence3, fuzzy_threshold=0.8)
    result4 = search_idiom(test_sentence4, fuzzy_threshold=0.8)
    end = time.time()

    # 输出结果：纯成语列表格式
    print("===== 测试句子1结果 =====")
    print(f"匹配成语：{result1}")
    print("\n===== 测试句子2结果 =====")
    print(f"匹配成语：{result2}")
    print("\n===== 测试句子3结果 =====")
    print(f"匹配成语：{result3}")
    print("\n===== 测试句子4结果 =====")
    print(f"匹配成语：{result4}")
    print(f"\n总检索耗时：{end - start:.4f} 秒")