import json
from pypinyin import lazy_pinyin, load_phrases_dict, Style
import re
import logging

# -------------------------- 新增：日志配置（核心调试用） --------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(funcName)s - %(message)s",
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger(__name__)

# -------------------------- 新增：拼音去声调工具函数（保留声母+韵母） --------------------------
def remove_tone_from_pinyin(pinyin_char: str) -> str:
    """
    单个拼音字符去声调（保留完整声母+韵母）
    例如：yī → yi，zhòng → zhong，āi → ai
    """
    tone_map = {
        'ā': 'a', 'á': 'a', 'ǎ': 'a', 'à': 'a',
        'ē': 'e', 'é': 'e', 'ě': 'e', 'è': 'e',
        'ī': 'i', 'í': 'i', 'ǐ': 'i', 'ì': 'i',
        'ō': 'o', 'ó': 'o', 'ǒ': 'o', 'ò': 'o',
        'ū': 'u', 'ú': 'u', 'ǔ': 'u', 'ù': 'u',
        'ǖ': 'v', 'ǘ': 'v', 'ǚ': 'v', 'ǜ': 'v',
        'ń': 'n', 'ň': 'n', 'ǹ': 'n'
    }
    return ''.join([tone_map.get(char, char) for char in pinyin_char])

def process_idiom_pinyin_from_xinhua(original_pinyin: str) -> tuple:
    """
    处理新华成语库的带声调拼音：去声调+拆分单字拼音列表
    输入："yī xīn yī yì"（新华库的带声调拼音）
    输出：("yixinyiyi", ["yi", "xin", "yi", "yi"])（完整拼音串+单字拼音列表）
    """
    if not original_pinyin or not isinstance(original_pinyin, str):
        logger.warning("无效的拼音输入：%s", original_pinyin)
        return "", []
    # 按空格拆分单字拼音（新华库的拼音是空格分隔的）
    single_pinyins = original_pinyin.strip().split()
    if not single_pinyins:
        logger.warning("拆分后无单字拼音：%s", original_pinyin)
        return "", []
    # 逐个单字拼音去声调
    clean_single_pinyins = []
    for p in single_pinyins:
        clean_p = remove_tone_from_pinyin(p)
        clean_single_pinyins.append(clean_p)
        logger.debug("单字拼音去声调：%s → %s", p, clean_p)
    # 拼接成连续拼音串（用于构建前缀树）
    full_pinyin_str = ''.join(clean_single_pinyins)
    logger.info("成语拼音处理完成：%s → %s（单字列表：%s）", original_pinyin, full_pinyin_str, clean_single_pinyins)
    return full_pinyin_str, clean_single_pinyins

# -------------------------- 1. 加载新华成语库JSON（原有逻辑保留，新增日志） --------------------------
def load_xinhua_idioms_json(json_path: str) -> list:
    try:
        with open(json_path, "r", encoding="utf-8") as f:
            data = json.load(f)
        logger.info("成功加载新华成语库，共%s条数据", len(data))
        return data
    except Exception as e:
        logger.error("加载新华成语库失败：%s", str(e), exc_info=True)
        raise

# -------------------------- 2. 构建拼音前缀树（Trie）（原有逻辑优化，新增拼音入参） --------------------------

# 基类：通用的Trie节点
class BaseTrieNode:
    """前缀树节点基类"""
    def __init__(self):
        self.children = {}  # 子节点：{字符: BaseTrieNode}
        self.idioms = []    # 该节点对应的成语（叶子节点存储）

# 基类：通用的Trie树
class BaseTrie:
    """前缀树基类"""
    def __init__(self):
        self.root = BaseTrieNode()
        self.added_idiom_count = 0  # 统计添加的成语数

    def to_dict(self, node=None):
        """将前缀树转换为字典（用于序列化JSON）"""
        if node is None:
            node = self.root
        node_dict = {
            "children": {},
            "idioms": node.idioms
        }
        for char, child in node.children.items():
            node_dict["children"][char] = self.to_dict(child)
        return node_dict

# 子类：拼音前缀树（保持原有拼音逻辑不变）
class PinyinTrie(BaseTrie):
    """拼音前缀树（保持原有拼音逻辑不变）"""
    def add_idiom(self, word: str, pinyin_str: str = None):
        """
        支持传入新华库处理后的拼音串，避免pypinyin转换
        若未传入拼音串，仍用pypinyin兜底（兼容原有逻辑）
        """
        if not word:  # 空成语直接返回
            return
        # 优先使用新华库的拼音串，否则用pypinyin兜底
        if pinyin_str:
            final_pinyin_str = pinyin_str
        else:
            pinyin_list = lazy_pinyin(word, style=Style.NORMAL)
            final_pinyin_str = ''.join(pinyin_list)
            logger.debug("使用pypinyin兜底生成拼音：%s → %s", word, final_pinyin_str)
        
        if not final_pinyin_str:
            logger.warning("成语%s无有效拼音，跳过", word)
            return
        
        node = self.root
        for char in final_pinyin_str:
            if char not in node.children:
                node.children[char] = BaseTrieNode()
            node = node.children[char]
        # 避免重复添加同一成语
        if word not in node.idioms:
            node.idioms.append(word)
            self.added_idiom_count += 1
        logger.debug("成语%s已加入拼音前缀树，拼音串：%s", word, final_pinyin_str)

# 子类：汉字前缀树（新增）
class ChineseTrie(BaseTrie):
    """汉字前缀树（新增）"""
    def add_idiom(self, word: str):
        """
        将成语按汉字字符添加到前缀树中
        """
        if not word:  # 空成语直接返回
            return
        
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = BaseTrieNode()
            node = node.children[char]
        # 避免重复添加同一成语
        if word not in node.idioms:
            node.idioms.append(word)
            self.added_idiom_count += 1
        logger.debug("成语%s已加入汉字前缀树", word)

# -------------------------- 3. 处理拼音：去声调+构建自定义词典（原有逻辑重构） --------------------------
def process_pinyin(idioms: list) -> tuple:
    """
    重构：基于新华库自带拼音处理，构建自定义多音字词典
    返回：(成语拼音映射字典, 成语集合, 自定义词典)
    成语拼音映射：{成语: 处理后的拼音串}
    自定义词典格式：{成语: [字1拼音, 字2拼音, ...]}（符合pypinyin的加载要求）
    """
    idiom_set = set()
    custom_dict = {}
    idiom_pinyin_map = {}  # 新增：存储成语与处理后的拼音串的映射
    for idx, idiom in enumerate(idioms):
        # 兼容成语数据为字典的情况（确保取到"word"和"pinyin"键）
        if not isinstance(idiom, dict) or "word" not in idiom or "pinyin" not in idiom:
            logger.warning("第%s条数据格式错误，缺少word/pinyin键：%s", idx+1, idiom)
            continue
        word = idiom["word"].strip()
        original_pinyin = idiom["pinyin"].strip()
        if not word:  # 空成语跳过
            continue
        # 处理新华库的带声调拼音
        pinyin_str, clean_single_pinyins = process_idiom_pinyin_from_xinhua(original_pinyin)
        if not pinyin_str:
            continue
        # 构建自定义词典：成语作为键，字拼音列表作为值（匹配pypinyin要求）
        custom_dict[word] = clean_single_pinyins
        # 存储成语与拼音串的映射
        idiom_pinyin_map[word] = pinyin_str
        idiom_set.add(word)
    # 加载自定义词典到pypinyin（兜底用，优先用新华库拼音）
    load_phrases_dict(custom_dict)
    logger.info("拼音处理完成，有效成语数：%s，自定义词典条目数：%s", len(idiom_set), len(custom_dict))
    return idiom_pinyin_map, idiom_set, custom_dict
    
# -------------------------- 4. 生成双树（拼音树+汉字树）并输出JSON（新增双树逻辑） --------------------------
if __name__ == "__main__":
    # 1. 加载并处理新华成语库
    idioms = load_xinhua_idioms_json("idiom.json")  # 你的新华成语库路径
    idiom_pinyin_map, idiom_set, custom_dict = process_pinyin(idioms)

    # 2. 构建拼音前缀树（优先使用新华库处理后的拼音串）
    pinyin_trie = PinyinTrie()
    for word in idiom_set:
        pinyin_str = idiom_pinyin_map.get(word)
        pinyin_trie.add_idiom(word, pinyin_str)  # 传入新华库的拼音串

    # 3. 构建汉字前缀树（新增）
    chinese_trie = ChineseTrie()
    for word in idiom_set:
        chinese_trie.add_idiom(word)

    # 4. 生成包含双树结构的JSON（拼音树+汉字树）
    dual_tree_json = {
        "pinyin_trie": pinyin_trie.to_dict(),
        "chinese_trie": chinese_trie.to_dict()
    }
    
    # 5. 保存JSON文件
    with open("idiom_trie.json", 'w', encoding='utf-8') as f:
        json.dump(dual_tree_json, f, ensure_ascii=False)
    
    # 新增：打印统计信息
    logger.info("双树结构JSON已生成！拼音树成语数：%s，汉字树成语数：%s", 
                pinyin_trie.added_idiom_count, chinese_trie.added_idiom_count)
    print("\n双树结构JSON已生成，包含拼音树和汉字树！")
    # 可选：打印根节点的子节点（调试用）
    logger.debug("拼音树根节点子节点：%s", list(pinyin_trie.root.children.keys())[:10])
    logger.debug("汉字树根节点子节点：%s", list(chinese_trie.root.children.keys())[:10])