#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
合并Artemis_II_十天绕月的字幕图片和TTS音频为视频片段
"""

import os
import subprocess
import re


def get_tts_audio_files(tts_dir):
    """
    获取TTS音频文件列表
    """
    audio_files = []
    for filename in os.listdir(tts_dir):
        if filename.endswith('.mp3') and filename.startswith('tts_'):
            # 解析文件名获取索引
            match = re.match(r'tts_(\d+)\.mp3', filename)
            if match:
                idx = int(match.group(1))
                audio_files.append((idx, os.path.join(tts_dir, filename)))
    
    # 按索引排序
    audio_files.sort(key=lambda x: x[0])
    return audio_files


def create_video_segment(image_path, audio_path, output_path):
    """
    使用FFmpeg将图片和音频合并为视频片段
    """
    cmd = [
        'ffmpeg', '-y',
        '-loop', '1',
        '-i', image_path,
        '-i', audio_path,
        '-c:v', 'libx264',
        '-tune', 'stillimage',
        '-c:a', 'aac',
        '-b:a', '192k',
        '-pix_fmt', 'yuv420p',
        '-shortest',
        output_path
    ]
    
    result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='ignore')
    return result.returncode == 0


def main():
    # 配置路径
    image_dir = "Artemis_II_十天绕月探测远非你所想象的那么简单_那它都要完成哪些任务呢_images"
    tts_dir = "Artemis_II_十天绕月_tts"
    output_dir = "Artemis_II_十天绕月_tts_videos"
    
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 获取TTS音频文件
    audio_files = get_tts_audio_files(tts_dir)
    print(f"找到 {len(audio_files)} 个TTS音频文件")
    
    # 合并每个片段
    for idx, audio_path in audio_files:
        # 构建图片路径
        image_path = os.path.join(image_dir, f"vertical_subtitle_{idx:03d}.png")
        
        # 检查图片是否存在
        if not os.path.exists(image_path):
            print(f"  警告: 图片不存在 - {image_path}")
            continue
        
        # 构建输出路径
        output_path = os.path.join(output_dir, f"video_{idx:03d}.mp4")
        
        # 如果视频已存在，跳过
        if os.path.exists(output_path):
            print(f"视频片段 {idx:03d} 已存在，跳过...")
            continue
        
        # 创建视频片段
        print(f"正在生成视频片段 {idx:03d}...")
        if create_video_segment(image_path, audio_path, output_path):
            print(f"  成功: {output_path}")
        else:
            print(f"  失败: video_{idx:03d}.mp4")
    
    print(f"\n完成！视频片段保存到: {output_dir}")


if __name__ == "__main__":
    main()