from moviepy.editor import *
from gtts import gTTS
from pydub import AudioSegment
import tempfile
import os

os.environ["IMAGEIO_FFMPEG_EXE"] = "/usr/local/bin/ffmpeg"
# 如果需要，设置 ImageMagick 的路径
# os.environ["IMAGEMAGICK_BINARY"] = "/usr/local/bin/convert"

def gen_audio(subtitles, tmpdir):
# 字幕列表，表示在指定时间段显示对应的字幕
# subtitles = [
#    [[0, 3], "這裡請注意右臂"],
#     [[4, 8], "這裡請注意左臂"],
#     [[9, 12], "這裡請注意左手細節"]
# ]
    subtitles = merge_subtitles(subtitles)
    print("===> Subtitles:")
    for subtitle in subtitles:
        print(subtitle)
    combined_clip = VideoFileClip(tmpdir + "/output.mp4")
    # 获取视频总时长
    total_duration = combined_clip.duration

    # 添加字幕的函数
    def add_subtitles(clip, subtitles):
        txt_clips = []
        for start_end, text in subtitles:
            start_time, end_time = start_end
            duration = end_time - start_time
            txt_clip = (TextClip(text, fontsize=clip.w//20, color='white', font='bold.ttf', method='caption', size=(clip.w * 0.9, None))  # 添加 method='caption'
                        .set_duration(duration)  # 设置字幕持续时间
                        .set_position(('center', 4 * clip.h // 5))
                        .set_start(start_time))  # 字幕开始时间
            txt_clips.append(txt_clip)
        
        # 将字幕叠加到视频上
        return CompositeVideoClip([clip, *txt_clips])

    # 将字幕添加到视频中
    video_with_subtitles = add_subtitles(combined_clip, subtitles)

    # 生成语音文件并返回音频剪辑的函数
    def generate_audio(subtitles, total_duration):
        temp_audio_path = tempfile.mkdtemp()  # 创建临时目录存储音频文件
        clips = []

        for i, (start_end, text) in enumerate(subtitles):
            start_time, end_time = start_end
            duration = end_time - start_time

            # 使用 gtts 生成语音
            tts = gTTS(text=text, lang='zh')
            audio_path = os.path.join(temp_audio_path, f'subtitle_{i}.mp3')
            tts.save(audio_path)
            
            # 使用 pydub 加载音频并将其转换为 wav 格式
            audio_segment = AudioSegment.from_mp3(audio_path)

            # 将 pydub 音频转换为 moviepy 可用的 AudioFileClip
            new_audio_path = os.path.join(temp_audio_path, f'subtitle_{i}_adjusted.wav')
            audio_segment.export(new_audio_path, format="wav")

            audio_clip = AudioFileClip(new_audio_path).set_start(start_time).set_duration(min(duration, audio_segment.duration_seconds))
            
            clips.append(audio_clip)
        
        # 合并所有音频剪辑
        final_audio = CompositeAudioClip(clips)
        return final_audio

    # 生成语音
    audio_clip = generate_audio(subtitles, total_duration)

    # 将语音添加到视频
    video_with_audio = video_with_subtitles.set_audio(audio_clip)

    # 导出带语音和字幕的视频
    video_with_audio.write_videofile(tmpdir + '/output_with_subtitles.mp4', fps=24)
    return tmpdir + '/output_with_subtitles.mp4'


def merge_subtitles(subtitles):
    merged_subtitles = []
    i = 0
    n = len(subtitles)
    
    while i < n:
        # 初始化当前字幕的起始时间、结束时间和内容
        current_start, current_end = subtitles[i][0]
        current_content = subtitles[i][1]
        duration = current_end - current_start
        i += 1  # 移动到下一个字幕索引，供后续合并使用

        # 尝试合并后续字幕，直到持续时间至少为3秒
        while duration < 3:
            if i >= n:
                # 已经处理完所有字幕，延长结束时间
                current_end = current_start + 3
                duration = current_end - current_start
                break
            next_start, next_end = subtitles[i][0]
            next_content = subtitles[i][1]
            if next_start == current_end:
                # 与下一个字幕相邻，合并
                current_end = next_end
                current_content += "、" + next_content
                duration = current_end - current_start
                i += 1  # 继续检查下一个字幕
            else:
                # 不相邻，无法合并，延长结束时间
                if duration < 3:
                    current_end = current_start + 3
                    duration = current_end - current_start
                break  # 退出合并循环

        # 添加合并后的字幕到新列表
        current_content = '這裡請注意' + '、'.join(list(set(current_content.replace("細節", "").split('、'))))
        merged_subtitles.append([[current_start, current_end], current_content])
    
    return merged_subtitles