# ------------
# 语义切割模块 - 基于标点符号分割ASR结果
# ------------
import json                                                                     # 导入JSON处理模块
import os                                                                       # 导入操作系统模块
import re                                                                       # 导入正则表达式模块
import pandas as pd                                                             # 导入pandas处理Excel文件
from core.utils import *                                                        # 导入核心工具函数
from core.utils.models import _2_CLEANED_CHUNKS, _SPLIT_RESULT                  # 导入文件路径常量



@check_file_exists(_SPLIT_RESULT)                                              # 检查切割结果文件是否存在的装饰器
def split_semantic():
    """
    语义切割主函数：读取ASR结果，进行标点符号切割，输出切割结果
    
    处理流程：
    1. 读取ASR步骤6保存的Excel文件（cleaned_chunks.xlsx）
    2. 组合文本进行标点符号切割
    3. 通过文本匹配获取时间戳
    4. 将切割结果重新组织为JSON格式并保存
    """
    
    # ------------
    # 步骤1: 读取ASR步骤6的Excel结果
    # ------------
    rprint("[cyan]📖 读取ASR转录结果...[/cyan]")
    
    if not os.path.exists(_2_CLEANED_CHUNKS):
        raise FileNotFoundError(f"ASR结果文件不存在: {_2_CLEANED_CHUNKS}")
    
    # 读取Excel文件，包含text, start, end, speaker_id列
    chunks_df = pd.read_excel(_2_CLEANED_CHUNKS)
    chunks_df['text'] = chunks_df['text'].apply(lambda x: x.strip('"').strip(""))  # 移除引号
    
    if chunks_df.empty:
        raise ValueError("ASR结果为空")
    
    rprint(f"[green]✅ 成功读取 {len(chunks_df)} 个ASR词汇片段[/green]")
    
    # ------------
    # 步骤2: 组合文本进行切割
    # ------------
    rprint("[cyan]📝 组合文本内容进行切割...[/cyan]")
    
    # 将所有ASR词汇组合成一个长文本
    combined_text = ' '.join(chunks_df['text'].tolist())
    
    # ------------
    # 步骤3: 执行简化的标点符号切割
    # ------------
    rprint("[cyan]✂️  执行标点符号切割...[/cyan]")
    
    # 使用简化的切割逻辑
    split_sentences = simple_punctuation_split(combined_text)
    
    rprint(f"[green]✅ 文本切割完成，共得到 {len(split_sentences)} 个句子片段[/green]")
    
    # ------------
    # 步骤4: 通过文本匹配获取时间戳
    # ------------
    rprint("[cyan]🕐 通过文本匹配获取时间戳...[/cyan]")
    
    split_result = match_timestamps_from_words(chunks_df, split_sentences)
    
    # ------------
    # 步骤6: 保存切割结果
    # ------------
    rprint("[cyan]💾 保存语义切割结果...[/cyan]")
    save_split_result(split_result)
    
    rprint(f"[green]🎉 语义切割完成！结果已保存到 {_SPLIT_RESULT}[/green]")

def simple_punctuation_split(text):
    """
    简化的标点符号切割函数
    
    Args:
        text: 要切割的文本
    
    Returns:
        list: 切割后的句子列表
    """
    
    # 第一步：按句号、问号、感叹号切割
    sentences = re.split(r'[。！？.!?]', text)
    
    # 第二步：对较长的句子按中文逗号切割
    final_sentences = []
    for sentence in sentences:
        sentence = sentence.strip()
        if not sentence:
            continue
            
        # 如果句子较长（超过20个字符），尝试按逗号切割
        if len(sentence) > 20 and '，' in sentence:
            # 按中文逗号切割
            comma_parts = sentence.split('，')
            for i, part in enumerate(comma_parts):
                part = part.strip()
                if part:
                    # 除了最后一部分，其他部分都加回逗号
                    if i < len(comma_parts) - 1:
                        part += '，'
                    final_sentences.append(part)
        else:
            final_sentences.append(sentence)
    
    # 过滤掉空句子和过短的句子
    final_sentences = [s for s in final_sentences if len(s.strip()) > 2]
    
    return final_sentences

def match_timestamps_from_words(chunks_df, split_sentences):
    """
    通过文本匹配从原始词汇级时间戳中获取句子的时间戳
    
    Args:
        chunks_df: 包含词汇级时间戳的DataFrame (text, start, end, speaker_id)
        split_sentences: 切割后的句子列表
    
    Returns:
        list: 包含时间戳的切割结果列表
    """
    
    split_result = []
    words_list = chunks_df['text'].tolist()
    
    # 为每个切割后的句子找到对应的时间戳
    word_index = 0  # 当前处理到的词汇索引
    
    for sentence_id, sentence in enumerate(split_sentences):
        sentence = sentence.strip()
        if not sentence:
            continue
        
        # 在句子中找到包含的词汇
        sentence_words = []
        sentence_start_word_idx = word_index
        
        # 简单的文本匹配逻辑：逐个匹配句子中的词汇
        remaining_sentence = sentence
        temp_word_idx = word_index
        
        while temp_word_idx < len(words_list) and remaining_sentence:
            current_word = words_list[temp_word_idx]
            
            # 检查当前词汇是否在剩余句子的开头
            if remaining_sentence.startswith(current_word):
                sentence_words.append(temp_word_idx)
                # 移除已匹配的词汇部分
                remaining_sentence = remaining_sentence[len(current_word):].lstrip()
                temp_word_idx += 1
            elif current_word in remaining_sentence:
                # 如果词汇在句子中但不在开头，可能是因为标点符号的影响
                # 尝试找到词汇在句子中的位置
                word_pos = remaining_sentence.find(current_word)
                if word_pos != -1:
                    sentence_words.append(temp_word_idx)
                    remaining_sentence = remaining_sentence[word_pos + len(current_word):].lstrip()
                    temp_word_idx += 1
                else:
                    temp_word_idx += 1
            else:
                temp_word_idx += 1
            
            # 防止无限循环
            if temp_word_idx - sentence_start_word_idx > 50:  # 最多匹配50个词
                break
        
        # 如果找到了匹配的词汇，计算时间戳
        if sentence_words:
            start_time = chunks_df.iloc[sentence_words[0]]['start']
            end_time = chunks_df.iloc[sentence_words[-1]]['end']
            word_index = sentence_words[-1] + 1  # 更新词汇索引
        else:
            # 如果没有找到匹配的词汇，使用前一个句子的结束时间或默认时间
            if split_result:
                start_time = split_result[-1]['end']
                end_time = start_time + 1.0  # 默认1秒长度
            else:
                start_time = 0.0
                end_time = 1.0
        
        split_result.append({
            "id": sentence_id,
            "text": sentence,
            "start": round(float(start_time), 2),
            "end": round(float(end_time), 2)
        })
        
        rprint(f"[blue]📍 句子 {sentence_id}: '{sentence[:30]}...' -> {start_time:.2f}s - {end_time:.2f}s[/blue]")
    
    return split_result

def save_split_result(split_result):
    """
    保存语义切割结果到JSON文件
    
    Args:
        split_result: 切割结果列表
    """
    
    # 确保输出目录存在
    os.makedirs('output', exist_ok=True)
    
    # 保存到JSON文件
    with open(_SPLIT_RESULT, 'w', encoding='utf-8') as f:
        json.dump(split_result, f, ensure_ascii=False, indent=2)
    
    rprint(f"[green]📄 语义切割结果已保存到 {_SPLIT_RESULT}，共 {len(split_result)} 个片段[/green]")

if __name__ == "__main__":                                                      # 主程序入口
    split_semantic()                                                            # 执行语义切割函数
