#!/usr/bin/env python3
"""
重新生成所有静态视频，修复像素格式问题
使用yuv420p而不是yuvj420p
"""
import os
import subprocess
import tempfile

# 导入编码处理模块
from console_helper import (
    safe_print, print_header, print_step, print_info,
    print_success, print_fail, print_done
)


def get_audio_duration(audio_path):
    """获取音频时长"""
    args = [
        'ffprobe',
        '-v', 'quiet',
        '-print_format', 'json',
        '-show_format',
        audio_path
    ]
    
    try:
        result = subprocess.run(args, check=True, capture_output=True, text=True)
        import json
        data = json.loads(result.stdout)
        return float(data['format']['duration'])
    except Exception as e:
        safe_print(f"      获取音频时长失败: {e}")
        return 0


def create_silent_audio(duration, output_path, ffmpeg_path="ffmpeg"):
    """创建静音音频"""
    args = [
        ffmpeg_path,
        '-y',
        '-f', 'lavfi',
        '-i', 'anullsrc=r=44100:cl=mono',
        '-c:a', 'aac',
        '-ar', '44100',
        '-ac', '1',
        '-t', str(duration),
        output_path
    ]
    
    try:
        subprocess.run(args, check=True, capture_output=True)
        return True
    except subprocess.CalledProcessError as e:
        safe_print(f"      创建静音音频失败: {e}")
        return False


def concatenate_audio_files(audio_files, output_path, ffmpeg_path="ffmpeg"):
    """使用filter_complex拼接多个音频文件"""
    inputs = []
    filters = []
    
    for i, audio_file in enumerate(audio_files):
        inputs.extend(['-i', audio_file])
        filters.append(f"[{i}:a]")
    
    filters.append(f"concat=n={len(audio_files)}:v=0:a=1[out]")
    
    args = [
        ffmpeg_path,
        '-y',
    ] + inputs + [
        '-filter_complex', ''.join(filters),
        '-map', '[out]',
        '-c:a', 'aac',
        '-ar', '44100',
        '-ac', '1',
        output_path
    ]
    
    try:
        subprocess.run(args, check=True, capture_output=True)
        return True
    except subprocess.CalledProcessError as e:
        safe_print(f"      拼接音频失败: {e}")
        return False


def create_static_video(image_path, audio_path, output_path, silence_duration=0.5, ffmpeg_path="ffmpeg"):
    """使用图片和音频生成静态视频，音频前加空白音"""
    temp_dir = tempfile.gettempdir()
    
    original_duration = get_audio_duration(audio_path)
    if original_duration == 0:
        safe_print(f"      无法获取音频时长")
        return False
    
    # 创建静音音频
    silent_audio = os.path.join(temp_dir, f"silent_{os.path.basename(audio_path)}.aac")
    if not create_silent_audio(silence_duration, silent_audio):
        return False
    
    # 拼接静音音频和原始音频
    combined_audio = os.path.join(temp_dir, f"combined_{os.path.basename(audio_path)}.aac")
    if not concatenate_audio_files([silent_audio, audio_path], combined_audio):
        try:
            os.remove(silent_audio)
        except:
            pass
        return False
    
    total_duration = silence_duration + original_duration
    
    # 生成视频 - 使用yuv420p像素格式（与原始片段一致）
    args = [
        ffmpeg_path,
        '-y',
        '-loop', '1',
        '-i', image_path,
        '-i', combined_audio,
        '-c:v', 'libx264',
        '-c:a', 'aac',
        '-ar', '44100',
        '-ac', '1',
        '-t', str(total_duration),
        '-pix_fmt', 'yuv420p',  # 使用标准yuv420p，与原始片段一致
        '-color_range', 'tv',   # 使用TV范围（16-235）
        output_path
    ]
    
    try:
        subprocess.run(args, check=True, capture_output=True)
        try:
            os.remove(silent_audio)
            os.remove(combined_audio)
        except:
            pass
        return True
    except subprocess.CalledProcessError as e:
        safe_print(f"      生成视频失败: {e}")
        try:
            os.remove(silent_audio)
            os.remove(combined_audio)
        except:
            pass
        return False


def main():
    # 配置
    frames_dir = "all_frames_with_subtitle"
    audio_dir = "tts_audio_all"
    output_dir = "all_static_videos_with_silence"
    silence_duration = 0.5
    
    print_header("重新生成所有静态视频（修复像素格式）")
    safe_print(f"图片目录: {frames_dir}")
    safe_print(f"音频目录: {audio_dir}")
    safe_print(f"输出目录: {output_dir}")
    safe_print(f"空白音时长: {silence_duration*1000}毫秒")
    safe_print(f"像素格式: yuv420p（与原始片段一致）")
    safe_print("")
    
    # 检查目录是否存在
    if not os.path.exists(frames_dir):
        safe_print(f"错误: 图片目录 {frames_dir} 不存在")
        return
    
    if not os.path.exists(audio_dir):
        safe_print(f"错误: 音频目录 {audio_dir} 不存在")
        return
    
    # 创建输出目录
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 获取所有帧图片
    frame_files = sorted([f for f in os.listdir(frames_dir) if f.endswith('.jpg')])
    
    print_step(1, 1, f"生成静态视频...")
    generated_count = 0
    failed_count = 0
    
    for frame_file in frame_files:
        # 从文件名提取行号
        import re
        match = re.match(r'frame_(\d+)_with_subtitle\.jpg', frame_file)
        if match:
            line_number = int(match.group(1))
            
            image_path = os.path.join(frames_dir, frame_file)
            audio_path = os.path.join(audio_dir, f"tts_line_{line_number}.wav")
            output_path = os.path.join(output_dir, f"static_video_{line_number:03d}.mp4")
            
            # 检查音频文件是否存在
            if not os.path.exists(audio_path):
                safe_print(f"      警告: 音频文件 {audio_path} 不存在，跳过第{line_number}句")
                failed_count += 1
                continue
            
            safe_print(f"      生成第{line_number}句静态视频...")
            
            if create_static_video(image_path, audio_path, output_path, silence_duration):
                duration = get_audio_duration(output_path)
                print_info(f"      已生成", f"static_video_{line_number:03d}.mp4 ({duration:.2f}秒)")
                generated_count += 1
            else:
                safe_print(f"      生成第{line_number}句静态视频失败")
                failed_count += 1
    
    safe_print("")
    print_info("成功生成", f"{generated_count} 个静态视频")
    if failed_count > 0:
        print_info("失败", f"{failed_count} 个")
    print_info("输出目录", output_dir)
    
    print_done()


if __name__ == '__main__':
    main()