"""
基于spaCy句法分析的英文字幕切分模块
使用句法依存分析找到合适的切分点，保证意群完整性
"""
import warnings
from typing import List, Tuple
from core.spacy_utils.load_nlp_model import init_nlp
from core.utils import rprint
from core.utils.config_utils import load_key

warnings.filterwarnings("ignore", category=FutureWarning)


def find_syntactic_split_points(text: str, max_length: int, nlp) -> List[str]:
    """
    基于句法分析找到最佳切分点
    
    Args:
        text: 要切分的文本
        max_length: 最大长度限制
        nlp: spaCy模型实例
        
    Returns:
        切分后的句子列表
    """
    # 如果文本长度已经符合要求，直接返回
    if len(text) <= max_length:
        return [text]
    
    doc = nlp(text)
    
    # 如果解析失败或只有一个token，使用简单切分
    if len(doc) <= 1:
        return simple_split_by_length(text, max_length)
    
    # 寻找句法切分点
    split_points = []
    
    for token in doc:
        # 跳过开头和结尾的token
        if token.i <= 2 or token.i >= len(doc) - 3:
            continue
            
        # 检查当前位置是否适合切分
        if is_good_split_point(token, max_length):
            split_points.append(token.i)
    
    # 如果没找到合适的切分点，使用标点符号作为备选
    if not split_points:
        split_points = find_punctuation_split_points(doc, max_length)
    
    # 如果还是没找到，使用简单的长度切分
    if not split_points:
        return simple_split_by_length(text, max_length)
    
    # 根据找到的切分点进行切分
    return split_by_points(doc, split_points, max_length)


def is_good_split_point(token, max_length: int) -> bool:
    """
    判断一个token是否是好的切分点
    
    优先考虑的切分点：
    1. 从句连接词之前 (SCONJ)
    2. 并列连词之前 (CCONJ) 
    3. 关系代词之前 (who, which, that等)
    4. 介词短语开头 (IN)
    5. 副词从句开头 (ADVMOD)
    6. 分词短语开头 (VBG, VBN)
    """
    
    # 从句连接词 (subordinating conjunctions)
    if token.pos_ == 'SCONJ' or token.dep_ == 'mark':
        return True
    
    # 并列连词，但要检查是否连接两个独立的部分
    if token.pos_ == 'CCONJ' and token.text.lower() in ['and', 'but', 'or', 'yet', 'so']:
        # 检查左右两边是否都有足够的内容
        left_length = token.i
        right_length = len(token.doc) - token.i - 1
        if left_length >= 5 and right_length >= 5:
            return True
    
    # 关系代词和关系副词
    if token.text.lower() in ['who', 'which', 'that', 'where', 'when', 'why', 'how'] and token.dep_ in ['nsubj', 'dobj', 'advmod']:
        return True
    
    # 介词短语开头，特别是表示时间、地点、原因的
    if token.pos_ == 'ADP' and token.text.lower() in ['in', 'on', 'at', 'by', 'with', 'for', 'during', 'after', 'before', 'because', 'since']:
        # 确保不是句子开头的介词
        if token.i > 3:
            return True
    
    # 分词短语开头 (-ing, -ed)
    if token.pos_ in ['VERB'] and token.tag_ in ['VBG', 'VBN'] and token.dep_ in ['advcl', 'acl']:
        return True
    
    # 副词从句开头
    if token.pos_ == 'ADV' and token.dep_ == 'advmod' and token.head.pos_ == 'VERB':
        return True
    
    return False


def find_punctuation_split_points(doc, max_length: int) -> List[int]:
    """
    寻找标点符号切分点作为备选方案
    """
    split_points = []
    
    for token in doc:
        if token.i <= 2 or token.i >= len(doc) - 3:
            continue
            
        # 逗号切分点
        if token.text == ',' and token.i > len(doc) // 3:
            # 确保左右两边都有足够内容
            left_text = doc[:token.i].text
            right_text = doc[token.i+1:].text
            if len(left_text.strip()) >= 10 and len(right_text.strip()) >= 10:
                split_points.append(token.i + 1)  # 在逗号后切分
        
        # 分号切分点
        elif token.text == ';':
            split_points.append(token.i + 1)
    
    return split_points


def split_by_points(doc, split_points: List[int], max_length: int) -> List[str]:
    """
    根据切分点列表进行文本切分
    """
    # 选择最佳切分点（尽量让每部分长度均匀）
    text_length = len(doc.text)
    target_parts = (text_length + max_length - 1) // max_length  # 向上取整
    
    if target_parts <= 1:
        return [doc.text]
    
    # 如果只需要切成两部分，选择最接近中间的切分点
    if target_parts == 2:
        middle = len(doc) // 2
        best_point = min(split_points, key=lambda x: abs(x - middle))
        
        part1 = doc[:best_point].text.strip()
        part2 = doc[best_point:].text.strip()
        
        return [part1, part2] if part1 and part2 else [doc.text]
    
    # 如果需要切成更多部分，选择多个切分点
    else:
        # 选择分布较均匀的切分点
        selected_points = select_best_split_points(split_points, len(doc), target_parts - 1)
        
        if not selected_points:
            return simple_split_by_length(doc.text, max_length)
        
        parts = []
        start = 0
        
        for point in selected_points:
            part = doc[start:point].text.strip()
            if part:
                parts.append(part)
            start = point
        
        # 添加最后一部分
        final_part = doc[start:].text.strip()
        if final_part:
            parts.append(final_part)
        
        return parts if len(parts) > 1 else [doc.text]


def select_best_split_points(split_points: List[int], total_length: int, num_splits: int) -> List[int]:
    """
    从候选切分点中选择最佳的几个点
    """
    if len(split_points) <= num_splits:
        return sorted(split_points)
    
    # 计算理想的切分位置
    ideal_positions = []
    for i in range(1, num_splits + 1):
        ideal_pos = (total_length * i) // (num_splits + 1)
        ideal_positions.append(ideal_pos)
    
    # 为每个理想位置选择最接近的切分点
    selected = []
    used_points = set()
    
    for ideal_pos in ideal_positions:
        # 找到最接近理想位置且未被使用的切分点
        available_points = [p for p in split_points if p not in used_points]
        if available_points:
            best_point = min(available_points, key=lambda x: abs(x - ideal_pos))
            selected.append(best_point)
            used_points.add(best_point)
    
    return sorted(selected)


def simple_split_by_length(text: str, max_length: int) -> List[str]:
    """
    简单的按长度切分，作为最后的备选方案
    在单词边界处切分
    """
    if len(text) <= max_length:
        return [text]
    
    parts = []
    words = text.split()
    current_part = []
    current_length = 0
    
    for word in words:
        # 如果单个词本身就超长，执行硬切分
        if len(word) > max_length:
            if current_part:
                parts.append(' '.join(current_part))
                current_part = []
                current_length = 0
            # 硬切分这个超长的单词
            for i in range(0, len(word), max_length):
                parts.append(word[i:i + max_length])
            continue

        word_length = len(word) + (1 if current_part else 0)  # +1 for space
        
        if current_length + word_length <= max_length:
            current_part.append(word)
            current_length += word_length
        else:
            if current_part:
                parts.append(' '.join(current_part))
                current_part = [word]
                current_length = len(word)
            else:
                # 这种情况不应该发生，因为我们已经在前面处理了超长的单个词
                parts.append(word)
                current_length = 0
    
    if current_part:
        parts.append(' '.join(current_part))
    
    return parts


def split_english_subtitle_by_syntax(text: str, max_length: int = 50) -> List[str]:
    """
    主函数：使用spaCy句法分析切分英文字幕
    
    Args:
        text: 要切分的英文文本
        max_length: 最大长度限制
        
    Returns:
        切分后的文本列表
    """
    if not text or not text.strip():
        return []
    
    text = text.strip()
    
    # 如果文本已经符合长度要求，直接返回
    if len(text) <= max_length:
        return [text]
    
    try:
        # 初始化spaCy模型
        nlp = init_nlp()
        
        # 执行句法切分
        result = find_syntactic_split_points(text, max_length, nlp)
        
        # 验证和清理结果，移除空字符串
        final_result = [part.strip() for part in result if part and part.strip()]

        # 检查是否有部分仍然过长，如果存在，则使用更简单的分割方法
        validated_result = []
        for part in final_result:
            if len(part) > max_length:
                # 对于仍然过长的部分，直接使用基于长度的简单分割
                validated_result.extend(simple_split_by_length(part, max_length))
            else:
                validated_result.append(part)

        return validated_result if validated_result else [text]
        
    except Exception as e:
        rprint(f"[red]❌ Syntax-based splitting failed: {e}[/red]")
        rprint("[yellow]⚠️  Falling back to simple word-boundary splitting...[/yellow]")
        return simple_split_by_length(text, max_length)


if __name__ == "__main__":
    # 测试代码
    test_sentences = [
        "This is a very long sentence that needs to be split because it exceeds the maximum length limit that we have set for subtitles.",
        "Although the weather was terrible and it was raining heavily, we decided to go out and enjoy the beautiful scenery that the park had to offer.",
        "The quick brown fox jumps over the lazy dog while the sun is shining brightly in the clear blue sky.",
    ]
    
    for i, sentence in enumerate(test_sentences):
        rprint(f"\n[cyan]Test {i+1}:[/cyan] {sentence}")
        rprint(f"[yellow]Length:[/yellow] {len(sentence)}")
        
        result = split_english_subtitle_by_syntax(sentence, max_length=50)
        
        rprint(f"[green]Split result ({len(result)} parts):[/green]")
        for j, part in enumerate(result):
            rprint(f"  {j+1}. [{len(part)}] {part}")

