#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
分批拼接"我们需要知识分子吗_思考"的双语视频
"""

import os
import subprocess
import re

def get_video_files(video_dir):
    """
    获取视频文件列表，按数字顺序排序
    """
    video_files = []
    for filename in os.listdir(video_dir):
        if filename.endswith('.mp4') and filename.startswith('video_'):
            # 解析文件名获取索引
            match = re.match(r'video_(\d+)\.mp4', filename)
            if match:
                idx = int(match.group(1))
                video_files.append((idx, os.path.join(video_dir, filename)))
    
    # 按索引排序
    video_files.sort(key=lambda x: x[0])
    return video_files

def create_batch_file(batch_videos, output_file):
    """
    创建批次文件，使用filter_complex拼接视频
    """
    # 构建filter_complex参数
    filter_complex = ""
    inputs = []
    
    for i, (_, video_path) in enumerate(batch_videos):
        inputs.extend(["-i", video_path])
        # 为每个视频添加音频滤镜，统一为stereo 44100Hz
        filter_complex += f"[{i}:v]scale=1080:1920:force_original_aspect_ratio=decrease,pad=1080:1920:(ow-iw)/2:(oh-ih)/2,setsar=1[v{i}];"
        filter_complex += f"[{i}:a]aformat=sample_fmts=s16:sample_rates=44100:channel_layouts=stereo[a{i}];"
    
    # 拼接视频和音频
    video_filters = ''.join([f"[v{i}]" for i in range(len(batch_videos))])
    audio_filters = ''.join([f"[a{i}]" for i in range(len(batch_videos))])
    
    filter_complex += f"{video_filters}concat=n={len(batch_videos)}:v=1:a=0[outv];"
    filter_complex += f"{audio_filters}concat=n={len(batch_videos)}:v=0:a=1[outa]"
    
    # 构建FFmpeg命令
    cmd = [
        "ffmpeg", "-y"
    ] + inputs + [
        "-filter_complex", filter_complex,
        "-map", "[outv]",
        "-map", "[outa]",
        "-c:v", "libx264",
        "-c:a", "aac",
        "-b:a", "192k",
        "-shortest",
        output_file
    ]
    
    print(f"执行命令: {' '.join(cmd[:10])} ...")  # 只打印前10个参数
    
    result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='ignore')
    return result.returncode == 0

def main():
    # 配置路径
    chinese_video_dir = "我们需要知识分子吗_思考_videos"
    tts_video_dir = "我们需要知识分子吗_思考_tts_videos"
    output_dir = "bilingual_output"
    final_output = "我们需要知识分子吗_思考_bilingual.mp4"
    
    # 确保输出目录存在
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # 获取视频文件
    chinese_videos = get_video_files(chinese_video_dir)
    tts_videos = get_video_files(tts_video_dir)
    
    print(f"找到 {len(chinese_videos)} 个中文视频")
    print(f"找到 {len(tts_videos)} 个英文TTS视频")
    
    # 确保数量匹配
    min_count = min(len(chinese_videos), len(tts_videos))
    print(f"使用 {min_count} 对视频进行拼接")
    
    # 交替拼接，每批10对视频
    batch_size = 10
    batch_files = []
    
    for i in range(0, min_count, batch_size):
        batch_end = min(i + batch_size, min_count)
        batch_pairs = []
        
        # 构建当前批次的视频对
        for j in range(i, batch_end):
            # 交替添加中文和英文视频
            batch_pairs.append(chinese_videos[j])
            batch_pairs.append(tts_videos[j])
        
        # 生成批次文件
        batch_output = os.path.join(output_dir, f"batch_{i//batch_size:03d}.mp4")
        print(f"\n处理批次 {i//batch_size:03d} (视频 {i+1}-{batch_end})")
        print(f"  包含 {len(batch_pairs)} 个视频片段")
        
        if create_batch_file(batch_pairs, batch_output):
            print(f"  成功: {batch_output}")
            batch_files.append(batch_output)
        else:
            print(f"  失败: 生成批次文件失败")
    
    # 合并所有批次
    if batch_files:
        print(f"\n合并 {len(batch_files)} 个批次...")
        
        # 构建合并命令
        inputs = []
        filter_complex = ""
        
        for i, batch_file in enumerate(batch_files):
            inputs.extend(["-i", batch_file])
            filter_complex += f"[{i}:v]scale=1080:1920:force_original_aspect_ratio=decrease,pad=1080:1920:(ow-iw)/2:(oh-ih)/2,setsar=1[v{i}];"
            filter_complex += f"[{i}:a]aformat=sample_fmts=s16:sample_rates=44100:channel_layouts=stereo[a{i}];"
        
        # 拼接视频和音频
        video_filters = ''.join([f"[v{i}]" for i in range(len(batch_files))])
        audio_filters = ''.join([f"[a{i}]" for i in range(len(batch_files))])
        
        filter_complex += f"{video_filters}concat=n={len(batch_files)}:v=1:a=0[outv];"
        filter_complex += f"{audio_filters}concat=n={len(batch_files)}:v=0:a=1[outa]"
        
        # 构建FFmpeg命令
        cmd = [
            "ffmpeg", "-y"
        ] + inputs + [
            "-filter_complex", filter_complex,
            "-map", "[outv]",
            "-map", "[outa]",
            "-c:v", "libx264",
            "-c:a", "aac",
            "-b:a", "192k",
            "-shortest",
            final_output
        ]
        
        print(f"执行合并命令: {' '.join(cmd[:10])} ...")  # 只打印前10个参数
        
        result = subprocess.run(cmd, capture_output=True, text=True, encoding='utf-8', errors='ignore')
        
        if result.returncode == 0:
            print(f"\n成功！双语视频已生成: {final_output}")
            # 清理临时文件
            for batch_file in batch_files:
                if os.path.exists(batch_file):
                    os.remove(batch_file)
            print("清理临时文件完成")
        else:
            print(f"\n失败: 合并批次失败 - {result.stderr}")
    else:
        print("错误: 没有生成任何批次文件")


if __name__ == "__main__":
    main()
