#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
字幕视频拼接器
根据字幕文件逐句切片视频，并为每句字幕生成TTS音频，然后拼接成完整视频
"""

import os
import re
import tempfile
import subprocess
import sqlite3
from pathlib import Path
from dataclasses import dataclass
from typing import List, Tuple, Optional
import argparse


@dataclass
class SubtitleEntry:
    """字幕条目"""
    index: int
    start_time: float  # 秒
    end_time: float    # 秒
    text: str


class SubtitleParser:
    """字幕解析器"""
    
    @staticmethod
    def parse_srt(content: str) -> List[SubtitleEntry]:
        """解析SRT格式的字幕"""
        entries = []
        
        # 按空行分割字幕块
        blocks = re.split(r'\n\s*\n', content.strip())
        
        for block in blocks:
            lines = block.strip().split('\n')
            if len(lines) < 3:
                continue
            
            try:
                # 第一行是序号
                index = int(lines[0].strip())
                
                # 第二行是时间码
                time_line = lines[1].strip()
                time_match = re.match(
                    r'(\d{2}):(\d{2}):(\d{2}),(\d{3})\s*-->\s*(\d{2}):(\d{2}):(\d{2}),(\d{3})',
                    time_line
                )
                
                if time_match:
                    start_time = (
                        int(time_match.group(1)) * 3600 +
                        int(time_match.group(2)) * 60 +
                        int(time_match.group(3)) +
                        int(time_match.group(4)) / 1000.0
                    )
                    end_time = (
                        int(time_match.group(5)) * 3600 +
                        int(time_match.group(6)) * 60 +
                        int(time_match.group(7)) +
                        int(time_match.group(8)) / 1000.0
                    )
                    
                    # 剩余行是字幕文本
                    text = ' '.join(lines[2:]).strip()
                    # 去除HTML标签
                    text = re.sub(r'<[^>]+>', '', text)
                    
                    entries.append(SubtitleEntry(
                        index=index,
                        start_time=start_time,
                        end_time=end_time,
                        text=text
                    ))
            except (ValueError, IndexError) as e:
                print(f"警告: 解析字幕块失败: {e}")
                continue
        
        return entries
    
    @classmethod
    def load_file(cls, filepath: str) -> List[SubtitleEntry]:
        """从文件加载字幕"""
        with open(filepath, 'r', encoding='utf-8') as f:
            content = f.read()
        return cls.parse_srt(content)
    
    @classmethod
    def load_from_db(cls, dbfile: str, dbtab: str, adid: int) -> List[SubtitleEntry]:
        """从数据库加载字幕"""
        entries = []
        
        try:
            conn = sqlite3.connect(dbfile)
            cursor = conn.cursor()
            
            # 查询该视频的字幕，使用 adname 字段（存储视频文件名），包含idx字段
            cursor.execute(f"SELECT idx, start, end, en FROM {dbtab} WHERE adname LIKE ? ORDER BY idx", ("%睡觉是人生头等大事%",))
            rows = cursor.fetchall()
            
            # 按照idx分组，选择每个idx对应的正确记录
            idx_entries = {}
            
            for i, row in enumerate(rows):
                idx = row[0]  # 获取idx字段
                start_time = row[1] / 1000.0  # 转换为秒
                end_time = row[2] / 1000.0    # 转换为秒
                text = row[3]
                
                # 过滤空字幕
                if not text or text.strip() == '':
                    continue
                
                # 只添加新的idx条目（避免重复）
                if idx not in idx_entries:
                    idx_entries[idx] = SubtitleEntry(
                        index=idx + 1,  # 从1开始编号
                        start_time=start_time,
                        end_time=end_time,
                        text=text
                    )
                    print(f"  添加字幕 {idx+1}: {text[:50]}...")
            
            # 转换为列表
            entries = list(idx_entries.values())
            
            conn.close()
        except Exception as e:
            print(f"数据库错误: {e}")
        
        return entries


class VideoProcessor:
    """视频处理器"""
    
    def __init__(self, ffmpeg_path: str = "ffmpeg"):
        self.ffmpeg_path = ffmpeg_path
    
    def _run_ffmpeg(self, args: List[str]) -> bool:
        """运行ffmpeg命令"""
        cmd = [self.ffmpeg_path] + args
        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                check=True
            )
            return True
        except subprocess.CalledProcessError as e:
            print(f"FFmpeg错误: {e.stderr}")
            return False
    
    def cut_video_segment(
        self,
        input_video: str,
        start_time: float,
        end_time: float,
        output_path: str
    ) -> bool:
        """剪切视频片段"""
        duration = end_time - start_time
        args = [
            '-y',  # 覆盖输出文件
            '-i', input_video,
            '-ss', str(start_time),
            '-t', str(duration),
            '-c:v', 'libx264',
            '-c:a', 'copy',
            '-pix_fmt', 'yuv420p',
            '-avoid_negative_ts', 'make_zero',
            output_path
        ]
        return self._run_ffmpeg(args)
    
    def create_silent_video(
        self,
        duration: float,
        output_path: str,
        width: int = 1280,
        height: int = 720
    ) -> bool:
        """创建静音视频（用于填充）"""
        args = [
            '-y',
            '-f', 'lavfi',
            '-i', f'color=c=black:s={width}x{height}:d={duration}',
            '-f', 'lavfi',
            '-i', 'anullsrc=r=44100:cl=mono',
            '-shortest',
            '-c:v', 'libx264',
            '-c:a', 'aac',
            '-t', str(duration),
            output_path
        ]
        return self._run_ffmpeg(args)


class TTSGenerator:
    """TTS音频生成器"""
    
    def __init__(self):
        self.temp_dir = tempfile.mkdtemp()
    
    def generate_audio(self, text: str, output_path: str, lang: str = 'zh') -> bool:
        """
        生成TTS音频
        使用edge-tts (Microsoft Edge TTS)
        """
        try:
            import edge_tts
            import asyncio
            
            async def _generate():
                # 根据语言选择语音
                if lang == 'zh':
                    voice = "zh-CN-XiaoxiaoNeural"
                elif lang == 'en':
                    voice = "en-US-JennyNeural"
                else:
                    voice = "zh-CN-XiaoxiaoNeural"
                
                # 控制语速 (0.8倍速)
                rate = "-20%"  # 0.8倍速
                
                communicate = edge_tts.Communicate(text, voice, rate=rate)
                await communicate.save(output_path)
            
            asyncio.run(_generate())
            return True
            
        except ImportError:
            print("错误: 未安装edge-tts模块。请运行: pip install edge-tts")
            return False
        except Exception as e:
            print(f"TTS生成错误: {e}")
            return False
    
    def get_audio_duration(self, audio_path: str) -> float:
        """获取音频时长"""
        try:
            from pydub import AudioSegment
            audio = AudioSegment.from_file(audio_path)
            return len(audio) / 1000.0
        except ImportError:
            # 使用ffprobe
            try:
                result = subprocess.run(
                    ['ffprobe', '-v', 'error', '-show_entries', 
                     'format=duration', '-of', 
                     'default=noprint_wrappers=1:nokey=1', audio_path],
                    capture_output=True,
                    text=True,
                    check=True
                )
                return float(result.stdout.strip())
            except Exception as e:
                print(f"获取音频时长失败: {e}")
                return 0.0


class VideoAudioMerger:
    """视频音频合并器"""
    
    def __init__(self, ffmpeg_path: str = "ffmpeg"):
        self.ffmpeg_path = ffmpeg_path
        self.video_processor = VideoProcessor(ffmpeg_path)
        self.tts_generator = TTSGenerator()
    
    def _run_ffmpeg(self, args: List[str]) -> bool:
        """运行ffmpeg命令"""
        cmd = [self.ffmpeg_path] + args
        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                check=True
            )
            return True
        except subprocess.CalledProcessError as e:
            print(f"FFmpeg错误: {e.stderr}")
            return False
    
    def merge_video_audio(
        self,
        video_path: str,
        audio_path: str,
        output_path: str,
        audio_duration: float = None
    ) -> bool:
        """
        将视频和音频合并
        如果音频比视频长，则循环播放视频
        """
        if audio_duration is None:
            audio_duration = self.tts_generator.get_audio_duration(audio_path)
        
        video_duration = self._get_video_duration(video_path)
        
        # 构建ffmpeg命令
        if audio_duration > video_duration:
            # 音频比视频长，需要循环视频
            args = [
                '-y',
                '-stream_loop', '-1',  # 无限循环视频
                '-i', video_path,
                '-i', audio_path,
                '-t', str(audio_duration),  # 以音频时长为准
                '-c:v', 'libx264',
                '-c:a', 'aac',
                '-shortest',
                '-pix_fmt', 'yuv420p',
                output_path
            ]
        else:
            # 音频比视频短或相等
            args = [
                '-y',
                '-i', video_path,
                '-i', audio_path,
                '-c:v', 'copy',
                '-c:a', 'aac',
                '-t', str(audio_duration),
                output_path
            ]
        
        return self._run_ffmpeg(args)
    
    def _get_video_duration(self, video_path: str) -> float:
        """获取视频时长"""
        try:
            result = subprocess.run(
                ['ffprobe', '-v', 'error', '-show_entries',
                 'format=duration', '-of',
                 'default=noprint_wrappers=1:nokey=1', video_path],
                capture_output=True,
                text=True,
                check=True
            )
            return float(result.stdout.strip())
        except Exception as e:
            print(f"获取视频时长失败: {e}")
            return 0.0
    
    def concatenate_videos(self, video_list: List[str], output_path: str) -> bool:
        """拼接多个视频"""
        # 创建临时文件列表
        list_file = os.path.join(tempfile.gettempdir(), 'video_list.txt')
        with open(list_file, 'w', encoding='utf-8') as f:
            for video in video_list:
                # 处理Windows路径中的反斜杠
                escaped_path = video.replace('\\', '/')
                f.write(f"file '{escaped_path}'\n")
        
        args = [
            '-y',
            '-f', 'concat',
            '-safe', '0',
            '-i', list_file,
            '-c:v', 'libx264',
            '-c:a', 'copy',
            '-pix_fmt', 'yuv420p',
            output_path
        ]
        
        result = self._run_ffmpeg(args)
        
        # 清理临时文件
        try:
            os.remove(list_file)
        except:
            pass
        
        return result


class SubtitleVideoCreator:
    """字幕视频创建器 - 主类"""
    
    def __init__(
        self,
        ffmpeg_path: str = "ffmpeg",
        temp_dir: Optional[str] = None
    ):
        self.ffmpeg_path = ffmpeg_path
        self.temp_dir = temp_dir or tempfile.mkdtemp()
        self.video_processor = VideoProcessor(ffmpeg_path)
        self.tts_generator = TTSGenerator()
        self.merger = VideoAudioMerger(ffmpeg_path)
        
        os.makedirs(self.temp_dir, exist_ok=True)
    
    def process(
        self,
        video_path: str,
        output_path: str,
        lang: str = 'zh',
        pause_between: float = 0.5,
        subtitle_path: Optional[str] = None,
        dbfile: Optional[str] = None,
        dbtab: Optional[str] = None,
        adid: Optional[int] = None
    ) -> bool:
        """
        处理视频和字幕，生成最终视频
        
        Args:
            video_path: 输入视频路径
            output_path: 输出视频路径
            lang: TTS语言 ('zh'为中文, 'en'为英文)
            pause_between: 句与句之间的停顿时间(秒)
            subtitle_path: 字幕文件路径(SRT格式)，如果指定了数据库参数则忽略
            dbfile: 数据库文件路径
            dbtab: 数据库表名
            adid: 视频ID
        """
        print(f"开始处理...")
        print(f"视频: {video_path}")
        
        # 1. 解析字幕
        print("\n[1/4] 解析字幕...")
        subtitles = []
        
        if dbfile and dbtab and adid:
            print(f"从数据库加载字幕: adid={adid}")
            subtitles = SubtitleParser.load_from_db(dbfile, dbtab, adid)
        elif subtitle_path:
            print(f"从文件加载字幕: {subtitle_path}")
            subtitles = SubtitleParser.load_file(subtitle_path)
        
        print(f"找到 {len(subtitles)} 条字幕")
        
        if not subtitles:
            print("错误: 没有解析到字幕")
            return False
        
        # 2. 处理每个字幕片段
        print("\n[2/4] 处理字幕片段...")
        segment_videos = []
        
        # 只处理前10句
        subtitles = subtitles[:10]

        
        for i, sub in enumerate(subtitles):
            print(f"  处理第 {i+1}/{len(subtitles)} 条: {sub.text[:30]}...")
            
            # 生成临时文件名
            video_segment = os.path.join(self.temp_dir, f"segment_{i:04d}_video.mp4")
            audio_file = os.path.join(self.temp_dir, f"segment_{i:04d}_audio.mp3")
            last_frame = os.path.join(self.temp_dir, f"segment_{i:04d}_last_frame.png")
            paused_video = os.path.join(self.temp_dir, f"segment_{i:04d}_paused.mp4")
            
            # 2.1 剪切视频片段
            if not self.video_processor.cut_video_segment(
                video_path, sub.start_time, sub.end_time, video_segment
            ):
                print(f"  警告: 剪切视频片段失败，跳过此段")
                continue
            
            # 2.2 生成TTS音频
            if not self.tts_generator.generate_audio(sub.text, audio_file, lang):
                print(f"  警告: 生成TTS音频失败，跳过此段")
                continue
            
            # 2.3 提取视频片段的最后一帧
            if not self._extract_last_frame(video_segment, last_frame):
                print(f"  警告: 提取最后一帧失败，跳过此段")
                continue
            
            # 2.4 将最后一帧延长到TTS音频的时长，并添加TTS音频
            audio_duration = self.tts_generator.get_audio_duration(audio_file)
            print(f"  TTS音频时长: {audio_duration:.2f}秒")
            if audio_duration > 0:
                # 提取视频片段的最后一帧
                last_frame = os.path.join(self.temp_dir, f"segment_{i:04d}_last_frame.png")
                if not self._extract_last_frame(video_segment, last_frame):
                    print(f"  警告: 提取最后一帧失败，使用黑色背景")
                    # 简化版本：直接将TTS音频转换为视频（黑色背景）
                    audio_video = os.path.join(self.temp_dir, f"segment_{i:04d}_audio_video.mp4")
                    if self._create_audio_video(audio_file, audio_video, audio_duration):
                        print(f"  音频视频创建成功: {audio_video}")
                        # 添加视频片段
                        segment_videos.append(video_segment)
                        # 添加音频视频片段
                        segment_videos.append(audio_video)
                        print(f"  已添加片段: {len(segment_videos)}个")
                    else:
                        print(f"  警告: 创建音频视频失败，跳过此段")
                        continue
                else:
                    # 创建带字幕的暂停视频
                    paused_video = os.path.join(self.temp_dir, f"segment_{i:04d}_paused.mp4")
                    # 显示字幕
                    if self._create_simple_paused_video(last_frame, audio_file, paused_video, audio_duration, sub.text):
                        print(f"  暂停视频创建成功: {paused_video}")
                        # 添加视频片段
                        segment_videos.append(video_segment)
                        # 添加暂停的视频片段（带TTS音频）
                        segment_videos.append(paused_video)
                        print(f"  已添加片段: {len(segment_videos)}个")
                    else:
                        print(f"  警告: 创建暂停视频失败，跳过此段")
                        continue
            else:
                print(f"  警告: TTS音频时长为0，跳过此段")
                continue
        
        if not segment_videos:
            print("错误: 没有成功生成任何视频片段")
            return False
        
        # 3. 拼接所有片段
        print("\n[3/4] 拼接所有片段...")
        if not self.merger.concatenate_videos(segment_videos, output_path):
            print("错误: 拼接视频失败")
            return False
        
        # 4. 清理临时文件
        print("\n[4/4] 清理临时文件...")
        self._cleanup()
        
        print(f"\n完成! 输出文件: {output_path}")
        return True
    
    def _extract_last_frame(self, video_path: str, output_path: str) -> bool:
        """提取视频的最后一帧"""
        args = [
            self.ffmpeg_path,
            '-y',
            '-sseof', '-1',  # 从视频末尾开始
            '-i', video_path,
            '-frames:v', '1',  # 只提取1帧
            '-q:v', '2',
            output_path
        ]
        try:
            subprocess.run(args, capture_output=True, check=True)
            return True
        except:
            return False
    
    def _create_simple_paused_video(self, image_path: str, audio_path: str, output_path: str, duration: float, subtitle_text: str = "") -> bool:
        """创建简单的暂停视频（静态画面 + TTS音频 + 字幕）"""
        # 基本命令
        args = [
            self.ffmpeg_path,
            '-y',
            '-loop', '1',
            '-framerate', '30',
            '-i', image_path,
            '-i', audio_path
        ]
        
        # 如果有字幕，添加字幕滤镜
        if subtitle_text:
            # 按单词换行，确保单词完整性
            words = subtitle_text.split()
            lines = []
            current_line = []
            current_length = 0
            max_length = 30  # 每行最大字符数
            
            for word in words:
                # 计算加上当前单词后的长度（包括空格）
                word_length = len(word)
                space_length = 1 if current_line else 0
                total_length = current_length + space_length + word_length
                
                if total_length > max_length:
                    # 行满了，添加当前行并开始新行
                    lines.append(' '.join(current_line))
                    current_line = [word]
                    current_length = word_length
                else:
                    # 单词可以加入当前行
                    if current_line:
                        current_line.append(word)
                        current_length += space_length + word_length
                    else:
                        current_line = [word]
                        current_length = word_length
            
            # 添加最后一行
            if current_line:
                lines.append(' '.join(current_line))
            
            # 更彻底地清理文本，避免特殊字符
            safe_lines = []
            for line in lines:
                # 只保留基本的ASCII字符
                safe_line = ''.join(c for c in line if ord(c) < 128)
                # 移除可能导致问题的字符
                safe_line = safe_line.replace('"', '').replace("'", '').replace('\\', '').replace('|', '').replace('&', '').replace(';', '')
                # 移除其他可能导致FFmpeg命令失败的字符
                safe_line = safe_line.replace('%', '').replace('#', '').replace('$', '').replace('`', '')
                if safe_line:
                    safe_lines.append(safe_line)
            
            # 使用多行drawtext滤镜，每行单独显示
            # 从底部向上排列
            vf_filters = []
            base_y = 50  # 最上面一行的Y坐标
            line_height = 45  # 每行高度
            
            for i, line in enumerate(safe_lines):
                y_pos = base_y + i * line_height
                vf_filters.append(f"drawtext=text='{line}':fontsize=50:fontcolor=red:box=0:bordercolor=green:borderw=2:x=(w-text_w)/2:y={y_pos}")
            
            # 组合所有滤镜
            if vf_filters:
                args.extend(['-vf', ','.join(vf_filters)])
        
        # 完成命令
        args.extend([
            '-c:v', 'libx264',
            '-c:a', 'aac',
            '-t', str(duration),
            '-pix_fmt', 'yuv420p',
            output_path
        ])
        
        try:
            # 使用capture_output=False，直接显示输出
            subprocess.run(args, check=True)
            return True
        except Exception as e:
            print(f"  简单暂停视频错误: {e}")
            # 如果字幕导致错误，尝试不带字幕的版本
            try:
                simple_args = [
                    self.ffmpeg_path,
                    '-y',
                    '-loop', '1',
                    '-i', image_path,
                    '-i', audio_path,
                    '-c:v', 'libx264',
                    '-c:a', 'aac',
                    '-t', str(duration),
                    '-pix_fmt', 'yuv420p',
                    output_path
                ]
                subprocess.run(simple_args, check=True)
                print(f"  不带字幕的版本创建成功")
                return True
            except:
                return False
    
    def _create_paused_video(self, image_path: str, audio_path: str, output_path: str, duration: float, subtitle_text: str) -> bool:
        """创建暂停视频（静态画面 + TTS音频 + 字幕）"""
        # 转义字幕文本中的特殊字符，避免ffmpeg命令出错
        safe_text = subtitle_text.replace('\\', '\\\\').replace("'", "\\'").replace('"', '\\"')
        
        # 构建字幕滤镜
        # 字幕会显示在视频上方，带半透明背景，红色50号字体，自动换行
        text_filter = f"drawtext=text='{safe_text}':fontsize=50:fontcolor=red:box=1:boxcolor=black@0.5:boxborderw=5:x=(w-text_w)/2:y=50:enable='between(t,0,{duration})':line_spacing=10:max_width=1100"
        
        args = [
            self.ffmpeg_path,
            '-y',
            '-loop', '1',
            '-i', image_path,
            '-i', audio_path,
            '-vf', text_filter,
            '-c:v', 'libx264',
            '-c:a', 'aac',
            '-t', str(duration),
            '-pix_fmt', 'yuv420p',
            '-shortest',
            output_path
        ]
        try:
            subprocess.run(args, capture_output=True, check=True)
            return True
        except Exception as e:
            print(f"  错误详情: {e}")
            return False
    
    def _create_audio_video(self, audio_path: str, output_path: str, duration: float) -> bool:
        """将TTS音频转换为视频（黑色背景，只有音频）"""
        args = [
            self.ffmpeg_path,
            '-y',
            '-f', 'lavfi',
            '-i', 'color=c=black:s=1280x720:d={}'.format(duration),
            '-i', audio_path,
            '-c:v', 'libx264',
            '-c:a', 'aac',
            '-t', str(duration),
            '-pix_fmt', 'yuv420p',
            output_path
        ]
        try:
            subprocess.run(args, capture_output=True, check=True)
            return True
        except:
            return False
    
    def _create_pause_frame(self, duration: float, output_path: str) -> bool:
        """创建停顿帧（黑色静音视频）"""
        args = [
            self.ffmpeg_path,
            '-y',
            '-f', 'lavfi',
            '-i', 'color=c=black:s=1280x720:d={}'.format(duration),
            '-f', 'lavfi',
            '-i', 'anullsrc=r=44100:cl=mono',
            '-shortest',
            '-c:v', 'libx264',
            '-c:a', 'aac',
            '-t', str(duration),
            output_path
        ]
        try:
            subprocess.run(args, capture_output=True, check=True)
            return True
        except:
            return False
    
    def _cleanup(self):
        """清理临时文件"""
        import shutil
        try:
            shutil.rmtree(self.temp_dir)
        except:
            pass


def main():
    # 直接设置参数
    video_path = "睡觉是人生头等大事#健康#张雪峰_xWT111.mp4"  # 输入视频文件路径
    
    # 选择数据源：SRT文件或数据库
    use_srt_file = True  # 设置为True使用SRT文件，False使用数据库
    
    if use_srt_file:
        subtitle_path = "output.srt"  # 字幕文件路径
        dbfile = None
        dbtab = None
        adid = None
        lang = "zh"  # TTS语言 (zh中文, en英文)
        output_path = "output_srt_version.mp4"  # 输出视频文件路径
    else:
        subtitle_path = None
        dbfile = r'\\ll\d\xtrssvjj\dpd.db'  # 数据库文件路径
        dbtab = 'srts'  # 数据库表名
        adid = 8243  # 视频ID
        lang = "en"  # TTS语言 (zh中文, en英文)
        output_path = "output_db_version.mp4"  # 输出视频文件路径
    
    pause = 0.5  # 句间停顿时间(秒)
    ffmpeg_path = "ffmpeg"  # FFmpeg可执行文件路径
    temp_dir = None  # 临时文件目录 (None表示使用系统临时目录)
    
    # 检查输入文件
    if not os.path.exists(video_path):
        print(f"错误: 视频文件不存在: {video_path}")
        return 1
    
    # 检查数据源
    if use_srt_file:
        if not os.path.exists(subtitle_path):
            print(f"错误: 字幕文件不存在: {subtitle_path}")
            return 1
    else:
        if not os.path.exists(dbfile):
            print(f"错误: 数据库文件不存在: {dbfile}")
            return 1
    
    # 创建处理器并运行
    creator = SubtitleVideoCreator(
        ffmpeg_path=ffmpeg_path,
        temp_dir=temp_dir
    )
    
    success = creator.process(
        video_path=video_path,
        output_path=output_path,
        lang=lang,
        pause_between=pause,
        subtitle_path=subtitle_path,
        dbfile=dbfile,
        dbtab=dbtab,
        adid=adid
    )
    
    return 0 if success else 1


if __name__ == '__main__':
    exit(main())