# -*- coding: utf-8 -*-
"""
术语格式化器 — 正则后处理层（零 token 消耗）

在 ASR 结果出来后、AI 校对之前执行。
从 A 稿中提取正确的术语写法，构建映射表，对 ASR 文本做确定性替换。

解决的问题:
- 讯飞 ASR 丢失英文型号中的短横线（F-15J→F15J, V-22→V22）
- 武器昵称引号丢失（A稿有引号但ASR没带出来）
- 中文数字被转成阿拉伯数字（数十→数10）
- 数字范围符号（~→到）
- 顿号分隔词加空格
- 小数点丢失修复（09马赫→0.9马赫）
- 军事领域高频同音字修正（建制→舰只等）
"""

import re
from typing import List, Tuple, Dict, Set


# ========================================================================
# 型号短横线修复
# ========================================================================

MODEL_PATTERN = re.compile(r'[A-Z]{1,4}-\d{1,4}[A-Z]?(?:/[A-Z])?')


def _build_model_mapping(script_text: str) -> Dict[str, str]:
    mapping = {}
    models = set(MODEL_PATTERN.findall(script_text))
    for model in models:
        no_hyphen = model.replace("-", "")
        if no_hyphen != model:
            mapping[no_hyphen] = model
    return mapping


def _fix_model_hyphens(text: str, mapping: Dict[str, str]) -> str:
    if not mapping:
        return text
    for no_hyphen in sorted(mapping.keys(), key=len, reverse=True):
        correct = mapping[no_hyphen]
        pattern = re.compile(re.escape(no_hyphen) + r'(?![A-Za-z0-9])')
        text = pattern.sub(correct, text)
    return text


# ========================================================================
# 武器昵称引号修复（上下文感知版）
# ========================================================================

# 匹配 A 稿中 "xxx"号 / "xxx"级 / "xxx"型 / 单独 "xxx" 的模式
QUOTED_WITH_SUFFIX = re.compile(r'“([^“”„‟""]{1,8})”([号级型式舰]?)')


def _build_quote_mapping(script_text: str) -> Dict[str, Set[str]]:
    """
    从 A 稿提取引号词及其后缀上下文。
    返回 {词: {出现过的后缀集合}}，后缀为空字符串表示单独使用。
    例: {"日向": {"号"}, "鱼鹰": {""}} 表示 A 稿有"日向"号但没有"日向"级，有单独的"鱼鹰"
    """
    mapping: Dict[str, Set[str]] = {}
    for match in QUOTED_WITH_SUFFIX.finditer(script_text):
        word = match.group(1).strip()
        suffix = match.group(2)
        if 2 <= len(word) <= 6:
            if word not in mapping:
                mapping[word] = set()
            mapping[word].add(suffix)
    return mapping


def _check_bare_occurrences(script_text: str, word: str, suffixes: Set[str]) -> Set[str]:
    """
    检查 A 稿中该词的无引号出现，看哪些后缀组合是不加引号的。
    例如 A 稿有 "日向级"（无引号），说明"日向级"不该加引号。
    """
    bare_suffixes = set()
    for suffix in ["号", "级", "型", "式", "舰", ""]:
        bare_pattern = word + suffix if suffix else word
        quoted_pattern = f"“{word}”{suffix}"
        # 在 A 稿中出现了无引号版本 且 没有对应的有引号版本
        if bare_pattern in script_text and quoted_pattern not in script_text:
            bare_suffixes.add(suffix)
    return bare_suffixes


def _fix_weapon_quotes(text: str, quote_mapping: Dict[str, Set[str]], script_text: str) -> str:
    """对文本中无引号的武器昵称补上引号（上下文感知）"""
    if not quote_mapping:
        return text
    for word in sorted(quote_mapping.keys(), key=len, reverse=True):
        quoted_suffixes = quote_mapping[word]
        bare_suffixes = _check_bare_occurrences(script_text, word, quoted_suffixes)

        # 对每个在 A 稿中确实带引号的后缀组合，在 ASR 文本中补引号
        for suffix in quoted_suffixes:
            if suffix and suffix not in bare_suffixes:
                # 匹配 "word+suffix"（无引号），替换为 "word"+suffix
                target = word + suffix
                replacement = f"“{word}”{suffix}"
                pattern = re.compile(
                    r'(?<!“)' + re.escape(target) + r'(?!”)'
                )
                text = pattern.sub(replacement, text)
            elif not suffix:
                # 单独出现（无后缀），但要避免替换那些在 A 稿中不带引号的后缀组合
                # 用负向前瞻排除不该加引号的后缀
                exclude_chars = "".join(bare_suffixes - {""}) if bare_suffixes else ""
                if exclude_chars:
                    lookahead = f'(?![{re.escape(exclude_chars)}])'
                else:
                    lookahead = ''
                pattern = re.compile(
                    r'(?<!“)(?<!《)' + re.escape(word) + lookahead + r'(?!”)(?!》)'
                )
                text = pattern.sub(f'“{word}”', text)
    return text


# ========================================================================
# 中文数字修复
# ========================================================================

CHINESE_NUM_FIXES = [
    (re.compile(r'数10([年架艘枚门辆台套件个发种类])'), r'数十\1'),
    (re.compile(r'数100([年架艘枚门辆台套件个发种类])'), r'数百\1'),
    (re.compile(r'数1000([年架艘枚门辆台套件个发种类])'), r'数千\1'),
    (re.compile(r'几10([年架艘枚门辆台套件个发种类])'), r'几十\1'),
    (re.compile(r'几100([年架艘枚门辆台套件个发种类])'), r'几百\1'),
]


def _fix_chinese_numbers(text: str) -> str:
    for pattern, replacement in CHINESE_NUM_FIXES:
        text = pattern.sub(replacement, text)
    return text


# ========================================================================
# 数字范围符号修复：~ ～ → 到
# ========================================================================

# 匹配 数字~数字 或 数字～数字 的模式
RANGE_TILDE = re.compile(r'(\d)[~～](\d)')


def _fix_range_symbol(text: str) -> str:
    return RANGE_TILDE.sub(r'\1到\2', text)


# ========================================================================
# 顿号→空格（唱词中并列词用空格分隔）
# ========================================================================

def _fix_enumeration_pause(text: str) -> str:
    return text.replace("、", " ")


# ========================================================================
# 节目名称书名号补全
# ========================================================================

# 需要带书名号的固定名称（节目名等）
# 格式: (裸名称, 带书名号版本)
BOOK_TITLE_NAMES = [
    ("军事科技", "《军事科技》"),
    ("军事报道", "《军事报道》"),
]


def _fix_book_titles(text: str) -> str:
    for bare, titled in BOOK_TITLE_NAMES:
        # 只替换没有被书名号包围的裸名称
        pattern = re.compile(r'(?<!《)' + re.escape(bare) + r'(?!》)')
        text = pattern.sub(titled, text)
    return text


# ========================================================================
# 小数点丢失修复（09马赫→0.9马赫 等）
# ========================================================================

# 匹配丢失小数点的情况：
# 1. "09马赫" → "0.9马赫"（数字在单位前）
# 2. "马赫数09" → "马赫数0.9"（数字在单位后）
# 3. 通用：非正常的 0+单个数字 紧跟/紧接单位
LOST_DECIMAL_BEFORE_UNIT = re.compile(r'(?<!\d)0(\d)(\s*(?:马赫|倍|秒|米|千米|公里))')
LOST_DECIMAL_AFTER_UNIT = re.compile(r'(马赫数|倍数|速度约)0(\d)(?!\d)')


def _fix_lost_decimal(text: str) -> str:
    text = LOST_DECIMAL_BEFORE_UNIT.sub(r'0.\1\2', text)
    text = LOST_DECIMAL_AFTER_UNIT.sub(r'\g<1>0.\2', text)
    return text


# ========================================================================
# 军事领域高频同音字修正
# ========================================================================

# 格式: (错误写法正则, 正确写法, A稿中应有的验证词)
# 只有当 A 稿中存在正确写法时才替换，避免误改
HOMOPHONE_PAIRS = [
    # 海军
    ("建制", "舰只", "舰只"),
    ("舰手", "舰艏", "舰艏"),
    ("舰位", "舰尾", "舰尾"),
    ("继承", "击沉", "击沉"),
    ("沉默", "沉没", "沉没"),
    ("空花弹", "滑翔弹", "滑翔弹"),
    ("建支", "舰只", "舰只"),
    ("坚支", "舰只", "舰只"),
    # 其他
    ("符和", "符合", "符合"),
    ("决意", "决议", "决议"),
]


def _build_homophone_mapping(script_text: str) -> Dict[str, str]:
    mapping = {}
    for wrong, correct, verify_word in HOMOPHONE_PAIRS:
        if verify_word in script_text:
            mapping[wrong] = correct
    return mapping


def _fix_homophones(text: str, mapping: Dict[str, str]) -> str:
    if not mapping:
        return text
    for wrong, correct in mapping.items():
        text = text.replace(wrong, correct)
    return text


# ========================================================================
# 主入口
# ========================================================================

def normalize_terms(
    sentences: List[Tuple[int, int, str, int]],
    script_text: str,
) -> List[Tuple[int, int, str, int]]:
    """
    对 ASR 句子列表做术语格式化（确定性正则替换，不调 AI）。
    在 ASR 结果出来后、AI 校对之前调用。
    """
    if not sentences:
        return []
    if not script_text:
        return list(sentences)

    model_mapping = _build_model_mapping(script_text)
    quote_mapping = _build_quote_mapping(script_text)
    homophone_mapping = _build_homophone_mapping(script_text)

    if model_mapping:
        print(f"[术语格式化] 型号映射 {len(model_mapping)} 条: {list(model_mapping.items())[:5]}")
    if quote_mapping:
        print(f"[术语格式化] 引号昵称 {len(quote_mapping)} 个: {dict((k, list(v)) for k, v in list(quote_mapping.items())[:5])}")
    if homophone_mapping:
        print(f"[术语格式化] 同音字映射 {len(homophone_mapping)} 条: {list(homophone_mapping.items())[:5]}")

    result = []
    fix_count = 0
    for bg, ed, text, spk in sentences:
        original = text
        text = _fix_model_hyphens(text, model_mapping)
        text = _fix_weapon_quotes(text, quote_mapping, script_text)
        text = _fix_chinese_numbers(text)
        text = _fix_range_symbol(text)
        text = _fix_enumeration_pause(text)
        text = _fix_lost_decimal(text)
        text = _fix_homophones(text, homophone_mapping)
        text = _fix_book_titles(text)
        if text != original:
            fix_count += 1
        result.append((bg, ed, text, spk))

    print(f"[术语格式化] 完成，修正 {fix_count} 句")
    return result