"""
Whisper视频转录模块

使用OpenAI Whisper对视频进行语音识别,生成文本转录和SRT字幕文件
"""
import os
from pathlib import Path
from typing import Dict, Any, Optional, List
import whisper
from tqdm import tqdm


class BilibiliTranscriber:
    """Whisper视频转录器
    
    使用OpenAI Whisper模型对视频进行语音识别,生成文本转录和SRT字幕
    """
    
    # 支持的模型列表
    AVAILABLE_MODELS = ['tiny', 'base', 'small', 'medium', 'large']
    
    def __init__(self, model_name: str = "base", language: str = "zh"):
        """初始化转录器
        
        Args:
            model_name: Whisper模型名称 (tiny/base/small/medium/large)
                       默认base,平衡速度和准确度
            language: 语言代码 (zh/en/ja等), 默认zh(中文)
                     设置语言可提升识别准确度
        
        Raises:
            ValueError: 模型名称不在支持列表中
        """
        if model_name not in self.AVAILABLE_MODELS:
            raise ValueError(
                f"不支持的模型: {model_name}\n"
                f"可用模型: {', '.join(self.AVAILABLE_MODELS)}"
            )
        
        self.model_name = model_name
        self.language = language
        self.model = None  # 延迟加载,首次调用时加载
    
    def load_model(self):
        """加载Whisper模型(延迟加载)
        
        首次调用时下载模型到~/.cache/whisper/
        后续调用直接使用缓存
        
        支持GPU自动检测:
        - CUDA (NVIDIA GPU)
        - MPS (Apple Silicon)
        - CPU (fallback)
        """
        if self.model is None:
            print(f"正在加载 Whisper {self.model_name} 模型...")
            print("(首次使用会自动下载模型,请耐心等待)")
            self.model = whisper.load_model(self.model_name)
            print("✅ 模型加载完成")
    
    def transcribe_video(
        self, 
        video_path: Path, 
        output_path: Optional[Path] = None,
        verbose: bool = True
    ) -> Dict[str, Any]:
        """转录视频为文本
        
        Args:
            video_path: 视频文件路径(.mp4/.mkv/.avi等)
            output_path: 输出文本路径(可选)
                        默认为: {video_path}_transcript.txt
            verbose: 是否显示详细进度信息
            
        Returns:
            转录结果字典:
            {
                'text': str,           # 完整转录文本
                'segments': List[Dict], # 分段信息,每段包含{start,end,text}
                'language': str        # 检测到的语言
            }
            
        Raises:
            FileNotFoundError: 视频文件不存在
            Exception: 转录过程失败
        """
        video_path = Path(video_path)
        
        if not video_path.exists():
            raise FileNotFoundError(f"视频文件不存在: {video_path}")
        
        # 加载模型
        self.load_model()
        
        # 执行转录
        if verbose:
            print(f"\n开始转录: {video_path.name}")
        
        result = self.model.transcribe(
            str(video_path),
            language=self.language,
            verbose=verbose
        )
        
        # 保存文本
        if output_path is None:
            output_path = video_path.parent / f"{video_path.stem}_transcript.txt"
        
        output_path = Path(output_path)
        output_path.write_text(result['text'], encoding='utf-8')
        
        if verbose:
            print(f"✅ 文本已保存: {output_path.name}")
        
        return result
    
    def transcribe_to_srt(
        self,
        video_path: Path,
        output_path: Optional[Path] = None
    ) -> Path:
        """转录视频并生成SRT字幕文件
        
        Args:
            video_path: 视频文件路径
            output_path: SRT输出路径(可选)
                        默认为: {video_path.stem}.srt
            
        Returns:
            SRT文件路径
            
        生成两个文件:
        1. {video_name}.srt - SRT字幕文件
        2. {video_name}_transcript.txt - 纯文本转录
        """
        video_path = Path(video_path)
        
        # 执行转录
        result = self.transcribe_video(video_path, verbose=True)
        
        # 生成SRT内容
        srt_content = self.generate_srt_content(result['segments'])
        
        # 保存SRT文件
        if output_path is None:
            output_path = video_path.parent / f"{video_path.stem}.srt"
        
        output_path = Path(output_path)
        output_path.write_text(srt_content, encoding='utf-8')
        
        print(f"✅ SRT字幕已保存: {output_path.name}")
        
        return output_path
    
    def batch_transcribe(
        self,
        video_dir: Path,
        output_dir: Optional[Path] = None,
        recursive: bool = False,
        generate_srt: bool = True
    ) -> Dict[str, Any]:
        """批量转录目录下的所有视频
        
        Args:
            video_dir: 视频目录路径
            output_dir: 输出目录(可选)
                       默认为各视频所在目录
            recursive: 是否递归处理子目录
            generate_srt: 是否生成SRT字幕(默认True)
            
        Returns:
            统计信息字典:
            {
                'total': int,      # 总文件数
                'success': int,    # 成功数量
                'failed': int,     # 失败数量
                'failed_files': List[str]  # 失败文件列表
            }
            
        处理策略:
        - 只处理.mp4文件
        - 跳过已存在的转录文件(已有.srt或_transcript.txt)
        - 单个文件失败不中断整体处理
        - 显示总体进度条
        """
        video_dir = Path(video_dir)
        
        if not video_dir.exists():
            raise FileNotFoundError(f"目录不存在: {video_dir}")
        
        # 扫描视频文件
        if recursive:
            video_files = list(video_dir.rglob("*.mp4"))
        else:
            video_files = list(video_dir.glob("*.mp4"))
        
        if not video_files:
            print("⚠️  未找到任何MP4视频文件")
            return {
                'total': 0,
                'success': 0,
                'failed': 0,
                'failed_files': []
            }
        
        print(f"\n找到 {len(video_files)} 个视频文件")
        
        # 预加载模型(避免每个文件重复加载)
        self.load_model()
        
        # 统计信息
        stats = {
            'total': len(video_files),
            'success': 0,
            'failed': 0,
            'failed_files': []
        }
        
        # 批量处理
        for video_path in tqdm(video_files, desc="批量转录", unit="个"):
            try:
                # 检查是否已存在转录文件
                transcript_path = video_path.parent / f"{video_path.stem}_transcript.txt"
                srt_path = video_path.parent / f"{video_path.stem}.srt"
                
                if transcript_path.exists() or srt_path.exists():
                    print(f"⏭️  跳过(已存在): {video_path.name}")
                    stats['success'] += 1
                    continue
                
                # 执行转录
                if generate_srt:
                    self.transcribe_to_srt(video_path)
                else:
                    self.transcribe_video(video_path, verbose=False)
                
                stats['success'] += 1
                
            except Exception as e:
                print(f"❌ 失败: {video_path.name} - {e}")
                stats['failed'] += 1
                stats['failed_files'].append(str(video_path))
        
        return stats
    
    @staticmethod
    def format_timestamp(seconds: float) -> str:
        """格式化时间戳为SRT格式
        
        Args:
            seconds: 秒数 (如 65.5)
            
        Returns:
            格式化字符串 "HH:MM:SS,mmm" (如 "00:01:05,500")
            
        示例:
            0.0 -> "00:00:00,000"
            65.5 -> "00:01:05,500"
            3661.123 -> "01:01:01,123"
        """
        hours = int(seconds // 3600)
        minutes = int((seconds % 3600) // 60)
        secs = int(seconds % 60)
        millis = int((seconds % 1) * 1000)
        return f"{hours:02d}:{minutes:02d}:{secs:02d},{millis:03d}"
    
    def generate_srt_content(self, segments: List[Dict]) -> str:
        """根据Whisper segments生成SRT内容
        
        Args:
            segments: Whisper返回的分段列表
                     每个segment包含: {start, end, text}
            
        Returns:
            SRT格式字符串
            
        格式:
            序号
            开始时间 --> 结束时间
            文本内容
            (空行)
        """
        srt_lines = []
        
        for i, segment in enumerate(segments, start=1):
            # 序号
            srt_lines.append(str(i))
            
            # 时间范围
            start_time = self.format_timestamp(segment['start'])
            end_time = self.format_timestamp(segment['end'])
            srt_lines.append(f"{start_time} --> {end_time}")
            
            # 文本内容
            text = segment['text'].strip()
            srt_lines.append(text)
            
            # 空行分隔
            srt_lines.append("")
        
        return "\n".join(srt_lines)
