#!/usr/bin/env python3
"""
重新剪切视频片段 - 修复播放问题
"""
import os
import subprocess
import asyncio
import edge_tts
import tempfile
import json

from console_helper import safe_print, print_step, print_info


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 get_video_duration(input_video):
    """获取视频总时长(秒)"""
    args = [
        'ffprobe',
        '-v', 'quiet',
        '-print_format', 'json',
        '-show_format',
        input_video
    ]
    result = subprocess.run(args, capture_output=True, text=True)
    data = json.loads(result.stdout)
    return float(data['format']['duration'])


def time_to_seconds(time_str):
    """将时间字符串转换为秒"""
    parts = time_str.split(':')
    hours = int(parts[0])
    minutes = int(parts[1])
    seconds = float(parts[2])
    return hours * 3600 + minutes * 60 + seconds


def cut_segments_fixed(input_video, subtitles, output_dir):
    """剪切视频片段 - 使用连续时间戳逻辑"""
    segments_dir = os.path.join(output_dir, "segments")
    if not os.path.exists(segments_dir):
        os.makedirs(segments_dir)
    
    total_duration = get_video_duration(input_video)
    num_subs = len(subtitles)
    
    for i, sub in enumerate(subtitles):
        # 确定当前segment的开始和结束时间
        if i == 0:
            # 第一句从0开始
            start = 0.0
        else:
            # 其他句用上一句的结束时间
            start = time_to_seconds(subtitles[i-1]['end'])
        
        if i == num_subs - 1:
            # 最后一句用视频总时长
            end = total_duration
        else:
            # 其他句用下一句的开始时间
            end = time_to_seconds(subtitles[i+1]['start'])
        
        duration = end - start
        
        output_file = os.path.join(segments_dir, f"segment_{sub['index']:03d}.mp4")
        
        args = [
            'ffmpeg',
            '-y',
            '-i', input_video,
            '-ss', str(start),
            '-t', str(duration),
            '-c:v', 'libx264',
            '-preset', 'fast',
            '-c:a', 'aac',
            '-ar', '44100',
            '-ac', '1',
            '-pix_fmt', 'yuv420p',
            '-avoid_negative_ts', 'make_zero',
            '-movflags', '+faststart',
            output_file
        ]
        
        try:
            subprocess.run(args, check=True, capture_output=True)
        except subprocess.CalledProcessError as e:
            safe_print(f"      剪切片段 {sub['index']} 失败")
    
    return segments_dir


def generate_tts_audio(subtitles, output_dir):
    """生成TTS音频"""
    audio_dir = os.path.join(output_dir, "tts_audio")
    if not os.path.exists(audio_dir):
        os.makedirs(audio_dir)
    
    success_count = 0
    for sub in subtitles:
        text = sub['text']
        output_file = os.path.join(audio_dir, f"tts_line_{sub['index']}.wav")
        
        if os.path.exists(output_file):
            success_count += 1
            continue
        
        try:
            communicate = edge_tts.Communicate(text, rate="-25%")
            asyncio.run(communicate.save(output_file))
            if os.path.exists(output_file):
                success_count += 1
        except Exception as e:
            safe_print(f"      生成音频 {sub['index']} 异常: {str(e)}")
    
    safe_print(f"      成功生成 {success_count}/{len(subtitles)} 个音频文件")
    return audio_dir


def extract_frames(segments_dir, subtitles, output_dir):
    """提取每句字幕结束时的帧"""
    frames_dir = os.path.join(output_dir, "frames")
    if not os.path.exists(frames_dir):
        os.makedirs(frames_dir)
    
    for sub in subtitles:
        segment_file = os.path.join(segments_dir, f"segment_{sub['index']:03d}.mp4")
        output_file = os.path.join(frames_dir, f"frame_{sub['index']:03d}.jpg")
        
        if not os.path.exists(segment_file):
            continue
        
        args = [
            'ffmpeg',
            '-y',
            '-sseof', '-0.1',
            '-i', segment_file,
            '-frames:v', '1',
            '-q:v', '2',
            output_file
        ]
        
        try:
            subprocess.run(args, check=True, capture_output=True)
        except:
            safe_print(f"      提取帧 {sub['index']} 失败")
    
    return frames_dir


def estimate_text_width(text, font_size=50):
    """估算文本像素宽度"""
    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):
    """根据像素宽度换行文本"""
    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 add_subtitles_to_frames(frames_dir, subtitles, output_dir):
    """为帧图片添加字幕 - 居中显示"""
    frames_with_subtitle_dir = os.path.join(output_dir, "frames_with_subtitle")
    if not os.path.exists(frames_with_subtitle_dir):
        os.makedirs(frames_with_subtitle_dir)
    
    work_dir = os.getcwd()
    
    for sub in subtitles:
        input_frame = os.path.join(frames_dir, f"frame_{sub['index']:03d}.jpg")
        output_frame = os.path.join(frames_with_subtitle_dir, f"frame_{sub['index']:03d}_with_subtitle.jpg")
        
        if not os.path.exists(input_frame):
            continue
        
        text = sub['text']
        wrapped_text = wrap_text_by_pixel(text, max_width=650, font_size=50)
        lines = wrapped_text.split('\n')
        
        line_height = 60
        total_text_height = len(lines) * line_height
        video_height = 1280
        y_start = (video_height - total_text_height) / 2
        
        filters = []
        text_files = []
        
        for i, line in enumerate(lines):
            text_filename = f"temp_text_{sub['index']}_{i}.txt"
            text_file = os.path.join(work_dir, text_filename)
            text_files.append(text_file)
            with open(text_file, 'w', encoding='utf-8') as f:
                f.write(line)
            
            current_y = int(y_start + i * line_height)
            
            filter_str = (
                f"drawtext=textfile='{text_filename}':"
                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',
            '-y',
            '-i', input_frame,
            '-vf', vf_filter,
            '-frames:v', '1',
            output_frame
        ]
        
        try:
            subprocess.run(args, check=True, capture_output=True)
        except:
            safe_print(f"      添加字幕 {sub['index']} 失败")
        
        for text_file in text_files:
            try:
                os.remove(text_file)
            except:
                pass
    
    return frames_with_subtitle_dir


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:
        return 0


def create_silent_audio(duration, output_path):
    """创建静音音频"""
    args = [
        'ffmpeg',
        '-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:
        return False


def concatenate_audio_files(audio_files, output_path):
    """拼接音频文件"""
    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',
        '-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:
        return False


def create_static_videos(frames_with_subtitle_dir, audio_dir, subtitles, output_dir):
    """生成静态视频"""
    static_videos_dir = os.path.join(output_dir, "static_videos")
    if not os.path.exists(static_videos_dir):
        os.makedirs(static_videos_dir)
    
    temp_dir = tempfile.gettempdir()
    
    for sub in subtitles:
        index = sub['index']
        image_path = os.path.join(frames_with_subtitle_dir, f"frame_{index:03d}_with_subtitle.jpg")
        audio_path = os.path.join(audio_dir, f"tts_line_{index}.wav")
        output_path = os.path.join(static_videos_dir, f"static_video_{index:03d}.mp4")
        
        if not os.path.exists(image_path) or not os.path.exists(audio_path):
            continue
        
        original_duration = get_audio_duration(audio_path)
        if original_duration == 0:
            continue
        
        silence_duration = 0.5
        
        silent_audio = os.path.join(temp_dir, f"silent_{index}.aac")
        if not create_silent_audio(silence_duration, silent_audio):
            continue
        
        combined_audio = os.path.join(temp_dir, f"combined_{index}.aac")
        if not concatenate_audio_files([silent_audio, audio_path], combined_audio):
            try:
                os.remove(silent_audio)
            except:
                pass
            continue
        
        total_duration = silence_duration + original_duration
        
        args = [
            'ffmpeg',
            '-y',
            '-loop', '1',
            '-framerate', '25',
            '-i', image_path,
            '-i', combined_audio,
            '-c:v', 'libx264',
            '-preset', 'fast',
            '-c:a', 'aac',
            '-ar', '44100',
            '-ac', '1',
            '-t', str(total_duration),
            '-pix_fmt', 'yuv420p',
            '-movflags', '+faststart',
            output_path
        ]
        
        try:
            subprocess.run(args, check=True, capture_output=True)
        except:
            safe_print(f"      生成静态视频 {index} 失败")
        
        try:
            os.remove(silent_audio)
            os.remove(combined_audio)
        except:
            pass
    
    return static_videos_dir


def concat_final_video(segments_dir, static_videos_dir, subtitles, output_file):
    """拼接最终视频"""
    video_list = []
    
    for sub in subtitles:
        index = sub['index']
        segment_file = os.path.join(segments_dir, f"segment_{index:03d}.mp4")
        static_video_file = os.path.join(static_videos_dir, f"static_video_{index:03d}.mp4")
        
        if os.path.exists(segment_file) and os.path.exists(static_video_file):
            video_list.append(segment_file)
            video_list.append(static_video_file)
    
    if not video_list:
        return False
    
    input_args = []
    for video in video_list:
        input_args.extend(['-i', video])
    
    n = len(video_list)
    input_labels = []
    for i in range(n):
        input_labels.append(f"[{i}:v][{i}:a]")
    
    filter_complex = ''.join(input_labels) + f"concat=n={n}:v=1:a=1[outv][outa]"
    
    args = [
        'ffmpeg',
        '-y'
    ] + input_args + [
        '-filter_complex', filter_complex,
        '-map', '[outv]',
        '-map', '[outa]',
        '-c:v', 'libx264',
        '-preset', 'fast',
        '-c:a', 'aac',
        '-ar', '44100',
        '-ac', '1',
        '-pix_fmt', 'yuv420p',
        '-movflags', '+faststart',
        output_file
    ]
    
    try:
        subprocess.run(args, check=True, capture_output=True)
        return True
    except:
        return False


def process_video(video_file, srt_file, output_name):
    """处理单个视频"""
    safe_print("=" * 60)
    safe_print(f"处理视频: {output_name}")
    safe_print("=" * 60)
    
    output_dir = output_name
    
    safe_print("解析字幕文件...")
    subtitles = parse_srt_file(srt_file)
    print_info("字幕数量", str(len(subtitles)))
    safe_print("")
    
    print_step(1, 6, "剪切视频片段...")
    segments_dir = cut_segments_fixed(video_file, subtitles, output_dir)
    print_info("完成", f"片段保存在 {segments_dir}")
    
    print_step(2, 6, "生成TTS音频...")
    audio_dir = generate_tts_audio(subtitles, output_dir)
    print_info("完成", f"音频保存在 {audio_dir}")
    
    print_step(3, 6, "提取帧图片...")
    frames_dir = extract_frames(segments_dir, subtitles, output_dir)
    print_info("完成", f"帧图片保存在 {frames_dir}")
    
    print_step(4, 6, "添加字幕到帧图片...")
    frames_with_subtitle_dir = add_subtitles_to_frames(frames_dir, subtitles, output_dir)
    print_info("完成", f"带字幕帧保存在 {frames_with_subtitle_dir}")
    
    print_step(5, 6, "生成静态视频...")
    static_videos_dir = create_static_videos(frames_with_subtitle_dir, audio_dir, subtitles, output_dir)
    print_info("完成", f"静态视频保存在 {static_videos_dir}")
    
    print_step(6, 6, "拼接最终视频...")
    output_file = os.path.join(output_dir, "final_video.mp4")
    if concat_final_video(segments_dir, static_videos_dir, subtitles, output_file):
        print_info("完成", f"最终视频: {output_file}")
    else:
        safe_print("拼接失败")
        return
    
    safe_print("=" * 60)
    safe_print("完成!")
    safe_print("=" * 60)


def main():
    videos = [
        {
            'video': '功利主义#张雪峰#功利主义#实用主义_xWT111.mp4',
            'srt': '功利主义#张雪峰#功利主义#实用主义_xWT111_en.srt',
            'output': 'video1'
        },
        {
            'video': '张雪与张雪峰#张雪#张雪峰#理想主义#现实主义_xWT158.mp4',
            'srt': '张雪与张雪峰#张雪#张雪峰#理想主义#现实主义_xWT158_en.srt',
            'output': 'video2'
        }
    ]
    
    for v in videos:
        if os.path.exists(v['video']) and os.path.exists(v['srt']):
            process_video(v['video'], v['srt'], v['output'])
            safe_print("")
            safe_print("=" * 60)
            safe_print("")
        else:
            safe_print(f"跳过 {v['video']}，文件不存在")


if __name__ == '__main__':
    main()