#!/usr/bin/env python3
"""
重新生成所有带字幕的帧图片和静态视频
使用基于像素宽度的智能换行
"""
import os
import subprocess
import tempfile
import textwrap
import json

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


def parse_srt_file(srt_file):
    """解析SRT字幕文件"""
    subtitles = []
    
    with open(srt_file, 'r', encoding='utf-8') as f:
        content = f.read()
    
    blocks = content.strip().split('\n\n')
    
    for block in blocks:
        lines = block.strip().split('\n')
        if len(lines) >= 3:
            index = int(lines[0])
            time_line = lines[1]
            text = '\n'.join(lines[2:])
            
            # 解析时间
            times = time_line.split(' --> ')
            start_time = times[0].replace(',', '.')
            end_time = times[1].replace(',', '.')
            
            subtitles.append({
                'index': index,
                'start': start_time,
                'end': end_time,
                'text': text
            })
    
    return subtitles


def estimate_text_width(text, font_size=50):
    """
    估算文本像素宽度
    英文字符约为字体大小的0.5倍
    中文字符约为字体大小
    """
    width = 0
    for char in text:
        if '\u4e00' <= char <= '\u9fff':  # 中文字符
            width += font_size
        elif char.isalpha() or char.isdigit():  # 英文字母和数字
            width += font_size * 0.5
        elif char in '.,!?;:\'"()-':  # 标点符号
            width += font_size * 0.3
        else:  # 空格等
            width += font_size * 0.3
    return width


def wrap_text_by_pixel(text, max_width=650, font_size=50):
    """
    根据像素宽度换行文本
    视频宽度720，左右边距各35像素，可用宽度650
    """
    words = text.split()
    lines = []
    current_line = ""
    current_width = 0
    
    for word in words:
        word_width = estimate_text_width(word + " ", font_size)
        
        if current_width + word_width <= max_width:
            current_line += word + " "
            current_width += word_width
        else:
            if current_line:
                lines.append(current_line.strip())
            current_line = word + " "
            current_width = word_width
    
    if current_line:
        lines.append(current_line.strip())
    
    return '\n'.join(lines)


def escape_text_for_ffmpeg(text):
    """为FFmpeg的drawtext滤镜转义文本"""
    escaped = text.replace("'", "'\\''")
    return escaped


def add_subtitle_to_image(input_image, output_image, subtitle_text, ffmpeg_path="ffmpeg"):
    """为图片添加字幕，自动换行"""
    # 根据像素宽度换行
    wrapped_text = wrap_text_by_pixel(subtitle_text, max_width=650, font_size=50)
    
    # 将文本按行分割
    lines = wrapped_text.split('\n')
    
    # 构建多个drawtext滤镜，每行一个
    filters = []
    y_position = 50  # 起始Y位置
    line_height = 60  # 每行高度（字体50 + 间距10）
    
    for i, line in enumerate(lines):
        # 转义特殊字符
        safe_line = escape_text_for_ffmpeg(line)
        current_y = y_position + i * line_height
        
        filter_str = (
            f"drawtext=text='{safe_line}':"
            f"fontsize=50:"
            f"fontcolor=yellow:"
            f"borderw=3:"
            f"bordercolor=red:"
            f"x=(w-text_w)/2:"
            f"y={current_y}"
        )
        filters.append(filter_str)
    
    # 组合所有滤镜
    vf_filter = ','.join(filters)
    
    args = [
        ffmpeg_path,
        '-y',
        '-i', input_image,
        '-vf', vf_filter,
        '-frames:v', '1',
        output_image
    ]
    
    try:
        subprocess.run(args, check=True, capture_output=True)
        return True
    except subprocess.CalledProcessError as e:
        safe_print(f"      添加字幕失败: {e}")
        return False


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)
        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_temp.aac")
    if not create_silent_audio(silence_duration, silent_audio):
        return False
    
    # 拼接静音音频和原始音频
    combined_audio = os.path.join(temp_dir, f"combined_temp.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',
        '-framerate', '25',
        '-i', image_path,
        '-i', combined_audio,
        '-c:v', 'libx264',
        '-c:a', 'aac',
        '-ar', '44100',
        '-ac', '1',
        '-t', str(total_duration),
        '-pix_fmt', 'yuv420p',
        '-color_range', 'tv',
        '-vf', 'format=yuv420p',
        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():
    print_header("重新生成所有带字幕的帧图片和静态视频")
    
    subtitle_file = "subtitles.srt"
    frames_dir = "all_frames"
    frames_with_subtitle_dir = "all_frames_with_subtitle"
    audio_dir = "tts_audio_all"
    static_videos_dir = "all_static_videos_with_silence"
    
    # 检查输入文件
    if not os.path.exists(subtitle_file):
        safe_print(f"错误: 字幕文件 {subtitle_file} 不存在")
        return
    
    # 创建输出目录
    for dir_path in [frames_with_subtitle_dir, static_videos_dir]:
        if not os.path.exists(dir_path):
            os.makedirs(dir_path)
    
    # 解析字幕文件
    safe_print("解析字幕文件...")
    subtitles = parse_srt_file(subtitle_file)
    if not subtitles:
        safe_print("解析字幕失败，退出")
        return
    
    print_info("字幕数量", str(len(subtitles)))
    safe_print("")
    
    # 处理所有字幕句
    max_lines = len(subtitles)
    success_count = 0
    
    for i, sub in enumerate(subtitles):
        index = sub['index']
        text = sub['text']
        
        print_step(i + 1, max_lines, f"处理第{index}句: {text[:30]}...")
        
        # 输入文件
        input_frame = f"{frames_dir}/frame_{index:03d}.jpg"
        input_audio = f"{audio_dir}/tts_line_{index}.wav"
        
        # 输出文件
        output_frame = f"{frames_with_subtitle_dir}/frame_{index:03d}_with_subtitle.jpg"
        output_video = f"{static_videos_dir}/static_video_{index:03d}.mp4"
        
        # 检查输入文件
        if not os.path.exists(input_frame):
            safe_print(f"      警告: 帧图片 {input_frame} 不存在")
            continue
        
        if not os.path.exists(input_audio):
            safe_print(f"      警告: 音频文件 {input_audio} 不存在")
            continue
        
        # 生成带字幕的帧图片
        safe_print(f"      生成带字幕帧图片...")
        if not add_subtitle_to_image(input_frame, output_frame, text):
            continue
        
        # 生成静态视频
        safe_print(f"      生成静态视频...")
        if not create_static_video(output_frame, input_audio, output_video):
            continue
        
        duration = get_audio_duration(output_video)
        print_info(f"      完成", f"{output_video} ({duration:.2f}秒)")
        success_count += 1
    
    safe_print("")
    print_info("成功处理", f"{success_count}/{max_lines} 句")
    print_done()


if __name__ == '__main__':
    main()