import whisper_timestamped
from textgrid import TextGrid, IntervalTier
import os
import soundfile as sf

def transcribe_audio(audio_path, language=None):
    """
    使用 whisper_timestamped 对音频文件进行转录
    
    参数:
    audio_path (str): 音频文件的路径
    
    返回:
    dict: 包含转录结果的字典
    """
    model = whisper_timestamped.load_model("large-v3-turbo")
    if language is None:
        result = whisper_timestamped.transcribe(model, audio_path)
    else:
        result = whisper_timestamped.transcribe(model, audio_path, language=language)

    # 获取音频文件的最大时长
    with sf.SoundFile(audio_path) as f:
        max_duration = len(f) / f.samplerate
        
    # 创建一个新的 TextGrid 对象
    tg = TextGrid()

    # 创建一个 IntervalTier 用于存储转录文本
    tier = IntervalTier('Transcription', 0, max_duration)

    # 合并
    merged_segments = []
    prev_segment = result["segments"][0]
    for segment in result["segments"][1:]:
        if segment["start"] >= max_duration:
            continue
        if prev_segment["end"] >= segment["start"]:  # 检查前后两个seg是否相邻
            # 合并相邻词段
            merged_segment = {
                "start": prev_segment["start"],
                "end": segment["end"],
                "text": prev_segment["text"] + segment["text"]
            }
            prev_segment = merged_segment
        else:
            merged_segments.append(prev_segment)
            prev_segment = segment
    merged_segments.append(prev_segment)




    for segment in merged_segments:
        if segment["start"] >= max_duration:
            continue
        if segment["end"] > max_duration:
            segment["end"] = max_duration
        if not segment["text"].strip():
            continue
        tier.add(segment['start'], segment['end'], segment['text'])
    # 将 IntervalTier 添加到 TextGrid 中
    tg.append(tier)

    # 保存 TextGrid 文件，文件路径与音频文件同目录，名称添加 _transcription 后缀
    base_name = os.path.splitext(audio_path)[0]
    textgrid_path = base_name + '_transcription.TextGrid'
    # 获取音频文件所在目录的父目录
    parent_dir = os.path.dirname(os.path.dirname(audio_path))
    # 拼接 test_results 目录路径
    test_results_dir = os.path.join(parent_dir, 'test_results')
    # 确保 test_results 目录存在
    os.makedirs(test_results_dir, exist_ok=True)
    # 从音频文件名获取基础名称
    base_name = os.path.basename(os.path.splitext(audio_path)[0])
    # 拼接完整的 TextGrid 文件路径，保存到 test_results 目录
    textgrid_path = os.path.join(test_results_dir, base_name + '_whisper_timestamped.TextGrid')
    tg.write(textgrid_path)
    # print(f"TextGrid 文件已保存至: {textgrid_path}")



    return result


if __name__ == "__main__":
    audio_path = r'C:\Users\User\Desktop\Praasper\data\test_audio.wav'  # 替换为你的音频文件路径
    result = transcribe_audio(audio_path)
    # print(result)
