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)
    # 仅匹配长度完全相同的拼音（解决一心一计/力的误匹配）
    if len1 != len2:
        return 0.0
    match_count = 0
    # 用普通循环替代生成器，减少解释器开销
    for a, b in zip(pinyin1, pinyin2):
        if a == b:
            match_count += 1
    return match_count / len1

# -------------------------- 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):
        """初始化：加载树+预缓存所有成语的拼音映射（核心优化）"""
        self.trie_dict = json.loads(trie_json)
        self.root = TrieNode(self.trie_dict)
        # 预缓存：{拼音长度: {拼音串: 成语列表}} —— 避免模糊检索遍历树
        self.idiom_pinyin_cache = defaultdict(dict)
        # 预缓存：所有成语的集合（快速去重）
        self.all_idioms = set()
        # 从前缀树中提取所有成语和拼音（仅初始化时执行一次）
        self._build_pinyin_cache(self.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(self, pinyin: str) -> list:
        """精准检索：优化为字典遍历，替代对象属性访问"""
        if not pinyin:
            return []
        node_dict = self.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) -> list:
        """模糊检索：基于预缓存的拼音映射，替代DFS遍历（核心优化）"""
        if not pinyin or threshold <= 0:
            return []
        target_len = len(pinyin)
        matched_idioms = []
        # 仅匹配长度完全相同的拼音
        if target_len not in self.idiom_pinyin_cache:
            return []
        # 遍历同长度的拼音串
        for cache_pinyin, idioms in self.idiom_pinyin_cache[target_len].items():
            if calculate_match_rate(pinyin, cache_pinyin) >= threshold:
                matched_idioms.extend(idioms)
        # 去重：用集合快速去重
        return list(set(matched_idioms))

    def find_idioms_in_sentence(self, sentence: str, fuzzy_threshold: float = 0.8) -> list:
        """
        核心修改：返回仅包含成语名称的一维列表，精准在前，模糊在后
        1. 提取元组中的成语字符串，去掉句子片段
        2. 全局去重，避免精准和模糊中出现重复成语
        """
        # 清洗句子：移除非汉字字符（单次处理）
        clean_sentence = re.sub(r'[^\u4e00-\u9fa5]', '', sentence)
        if len(clean_sentence) < 4:
            return []  # 无结果返回空列表
        
        exact_idiom_set = set()  # 仅存成语名称的精准集合
        fuzzy_idiom_set = set()  # 仅存成语名称的模糊集合
        sentence_len = len(clean_sentence)

        # 仅遍历4字片段（核心优化：减少2/3的循环次数）
        for i in range(sentence_len - 3):
            fragment = clean_sentence[i:i+4]
            fragment_pinyin = chinese_to_pinyin(fragment)
            if not fragment_pinyin:
                continue

            # 精准检索：提取成语名称，加入精准集合
            exact_idioms = self.retrieve_exact(fragment_pinyin)
            exact_idiom_set.update(exact_idioms)

            # 模糊检索：提取成语名称，排除已精准匹配的，加入模糊集合
            fuzzy_idioms = self.retrieve_fuzzy(fragment_pinyin, fuzzy_threshold)
            fuzzy_idiom_set.update([idiom for idiom in fuzzy_idioms if idiom not in exact_idiom_set])

        # 核心：拼接为纯成语列表，精准在前，模糊在后，最终去重
        exact_list = list(exact_idiom_set)
        fuzzy_list = list(fuzzy_idiom_set)
        final_idiom_list = exact_list + fuzzy_list  # 精准成语 + 模糊成语

        logger.info(f"检索完成，精准匹配{len(exact_list)}条，模糊匹配{len(fuzzy_list)}条，总计{len(final_idiom_list)}条")
        return final_idiom_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) -> list:
    """搜索句子中的成语，使用预初始化的检索器实例"""
    return _retriever_instance.find_idioms_in_sentence(sentence, fuzzy_threshold)
    
# -------------------------- 5. 测试检索（适配纯成语列表的输出格式） --------------------------
if __name__ == "__main__":
    # 测试句子
    test_sentence1 = "他做事情总是三心二意，从来没有一心一意的时候。"
    test_sentence2 = "他一欣一意地完成工作，结果却画蛇添足了。"

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

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