batch_processor.py

"""
批量处理模块
支持批量处理文件夹中的所有音视频文件
提供进度跟踪、错误处理、结果统计功能
"""

import os
import time
from pathlib import Path
from typing import List, Dict, Any, Optional, Callable, Tuple
from dataclasses import dataclass, field
from datetime import datetime
import logging
import concurrent.futures
from threading import Lock

logger = logging.getLogger(__name__)


@dataclass
class BatchProgress:
    """
    批量处理进度信息
    """
    total_files: int = 0
    completed_files: int = 0
    failed_files: int = 0
    skipped_files: int = 0
    current_file_index: int = 0
    current_file_name: str = ""
    current_file_path: str = ""
    current_stage: str = ""
    current_progress: float = 0.0
    elapsed_time: float = 0.0
    estimated_time_remaining: float = 0.0
    start_time: float = field(default_factory=time.time)
    results: List[Dict[str, Any]] = field(default_factory=list)
    errors: List[Dict[str, Any]] = field(default_factory=list)
    
    @property
    def overall_progress(self) -> float:
        """
        计算整体进度百分比
        """
        if self.total_files == 0:
            return 0.0
        return (self.completed_files + self.failed_files + self.skipped_files) / self.total_files * 100.0
    
    @property
    def success_rate(self) -> float:
        """
        计算成功率
        """
        if self.completed_files + self.failed_files == 0:
            return 0.0
        return self.completed_files / (self.completed_files + self.failed_files) * 100.0


class BatchProcessor:
    """
    批量处理器类
    支持批量处理文件夹中的音视频文件
    """
    
    # 支持的文件扩展名
    SUPPORTED_EXTENSIONS = {
        # 视频格式
        '.mp4', '.mov', '.avi', '.mkv', '.flv', '.wmv', '.webm',
        # 音频格式
        '.mp3', '.wav', '.m4a', '.flac', '.aac', '.ogg', '.wma'
    }
    
    def __init__(
        self,
        max_workers: int = 1,
        skip_existing: bool = True,
        output_suffix: str = '',
        progress_callback: Optional[Callable] = None,
        error_callback: Optional[Callable] = None,
        file_callback: Optional[Callable] = None,
    ):
        """
        初始化批量处理器
        
        Args:
            max_workers: 最大并发数（建议为1，因为Whisper模型通常占满GPU）
            skip_existing: 是否跳过已存在字幕文件的视频
            output_suffix: 输出文件名后缀
            progress_callback: 进度回调函数
            error_callback: 错误回调函数
            file_callback: 文件处理完成回调函数
        """
        self.max_workers = max_workers
        self.skip_existing = skip_existing
        self.output_suffix = output_suffix
        self.progress_callback = progress_callback
        self.error_callback = error_callback
        self.file_callback = file_callback
        
        # 线程锁（用于并发安全）
        self._lock = Lock()
        
        # 处理状态
        self._is_running = False
        self._should_stop = False
    
    def _get_output_path(self, input_path: str, output_dir: str) -> str:
        """
        根据输入文件路径生成输出字幕文件路径
        
        Args:
            input_path: 输入文件路径
            output_dir: 输出目录
            
        Returns:
            输出字幕文件路径
        """
        input_path = Path(input_path).resolve()
        output_dir = Path(output_dir).resolve()
        
        # 生成输出文件名
        if self.output_suffix:
            base_name = f"{input_path.stem}{self.output_suffix}"
        else:
            base_name = input_path.stem
        
        output_file = output_dir / f"{base_name}.srt"
        return str(output_file)
    
    def scan_files(
        self,
        input_path: str,
        recursive: bool = False,
        output_dir: Optional[str] = None,
    ) -> List[Tuple[str, str]]:
        """
        扫描文件夹中的音视频文件
        
        Args:
            input_path: 输入路径（文件或文件夹）
            recursive: 是否递归扫描子文件夹
            output_dir: 输出目录（用于检查是否已存在字幕）
            
        Returns:
            文件路径列表，格式为 [(输入路径, 输出路径), ...]
        """
        input_path = Path(input_path).resolve()
        file_pairs = []
        
        if input_path.is_file():
            # 单个文件
            if input_path.suffix.lower() in self.SUPPORTED_EXTENSIONS:
                # 确定输出路径
                if output_dir:
                    out_path = self._get_output_path(str(input_path), output_dir)
                else:
                    out_path = self._get_output_path(str(input_path), str(input_path.parent))
                
                # 检查是否跳过
                if self.skip_existing and Path(out_path).exists():
                    logger.info(f"跳过已存在字幕的文件: {input_path}")
                else:
                    file_pairs.append((str(input_path), out_path))
        else:
            # 文件夹
            pattern = '**/*' if recursive else '*'
            
            for file_path in input_path.glob(pattern):
                if file_path.is_file() and file_path.suffix.lower() in self.SUPPORTED_EXTENSIONS:
                    # 确定输出路径
                    if output_dir:
                        out_path = self._get_output_path(str(file_path), output_dir)
                    else:
                        out_path = self._get_output_path(str(file_path), str(file_path.parent))
                    
                    # 检查是否跳过
                    if self.skip_existing and Path(out_path).exists():
                        logger.info(f"跳过已存在字幕的文件: {file_path}")
                        continue
                    
                    file_pairs.append((str(file_path), out_path))
        
        logger.info(f"扫描完成，共找到 {len(file_pairs)} 个待处理文件")
        return file_pairs
    
    def _update_progress(self, progress: BatchProgress, stage: str, current_progress: float = 0.0):
        """
        更新进度信息
        
        Args:
            progress: 进度对象
            stage: 当前阶段
            current_progress: 当前文件进度
        """
        with self._lock:
            progress.current_stage = stage
            progress.current_progress = current_progress
            progress.elapsed_time = time.time() - progress.start_time
            
            # 估算剩余时间
            if progress.completed_files > 0:
                avg_time_per_file = progress.elapsed_time / progress.completed_files
                remaining_files = progress.total_files - progress.completed_files - progress.failed_files
                progress.estimated_time_remaining = avg_time_per_file * remaining_files
            else:
                progress.estimated_time_remaining = 0.0
        
        # 调用回调函数
        if self.progress_callback:
            try:
                self.progress_callback(progress)
            except Exception as e:
                logger.error(f"进度回调函数执行失败: {e}")
    
    def _handle_file_result(
        self,
        progress: BatchProgress,
        input_path: str,
        output_path: str,
        success: bool,
        result: Optional[Dict[str, Any]] = None,
        error: Optional[Exception] = None,
    ):
        """
        处理文件处理结果
        
        Args:
            progress: 进度对象
            input_path: 输入文件路径
            output_path: 输出文件路径
            success: 是否成功
            result: 成功结果
            error: 错误信息
        """
        with self._lock:
            if success:
                progress.completed_files += 1
                if result:
                    result_info = {
                        'input_path': input_path,
                        'output_path': output_path,
                        'success': True,
                        'timestamp': datetime.now().isoformat(),
                        **result
                    }
                    progress.results.append(result_info)
                    logger.info(f"文件处理成功: {input_path} -> {output_path}")
            else:
                progress.failed_files += 1
                error_info = {
                    'input_path': input_path,
                    'output_path': output_path,
                    'success': False,
                    'error': str(error) if error else "未知错误",
                    'error_type': type(error).__name__ if error else "Unknown",
                    'timestamp': datetime.now().isoformat()
                }
                progress.errors.append(error_info)
                logger.error(f"文件处理失败: {input_path}, 错误: {error}")
        
        # 调用文件回调
        if self.file_callback:
            try:
                self.file_callback(input_path, output_path, success, result, error)
            except Exception as e:
                logger.error(f"文件回调函数执行失败: {e}")
        
        # 调用错误回调
        if not success and self.error_callback:
            try:
                self.error_callback(input_path, error, result)
            except Exception as e:
                logger.error(f"错误回调函数执行失败: {e}")
    
    def process_single_file(
        self,
        input_path: str,
        output_path: str,
        processor_func: Callable,
        progress: Optional[BatchProgress] = None,
    ) -> Tuple[bool, Optional[Dict[str, Any]], Optional[Exception]]:
        """
        处理单个文件
        
        Args:
            input_path: 输入文件路径
            output_path: 输出文件路径
            processor_func: 处理函数，签名为 func(input_path, output_path) -> result
            progress: 进度对象（可选）
            
        Returns:
            (是否成功, 结果数据, 错误信息)
        """
        try:
            if progress:
                with self._lock:
                    progress.current_file_name = Path(input_path).name
                    progress.current_file_path = input_path
                self._update_progress(progress, "处理中")
            
            # 调用处理函数
            result = processor_func(input_path, output_path)
            
            return (True, result, None)
            
        except Exception as e:
            logger.exception(f"处理文件失败: {input_path}")
            return (False, None, e)
    
    def process_files(
        self,
        file_pairs: List[Tuple[str, str]],
        processor_func: Callable,
    ) -> BatchProgress:
        """
        批量处理文件
        
        Args:
            file_pairs: 文件路径列表 [(输入路径, 输出路径), ...]
            processor_func: 处理函数
            
        Returns:
            批量处理进度对象
        """
        if not file_pairs:
            logger.warning("没有待处理的文件")
            return BatchProgress()
        
        # 初始化进度
        progress = BatchProgress(
            total_files=len(file_pairs),
            start_time=time.time()
        )
        
        self._is_running = True
        self._should_stop = False
        
        logger.info(f"开始批量处理，共 {len(file_pairs)} 个文件")
        logger.info(f"最大并发数: {self.max_workers}")
        
        # 串行处理（推荐，因为Whisper模型通常占满GPU）
        if self.max_workers == 1:
            for idx, (input_path, output_path) in enumerate(file_pairs):
                if self._should_stop:
                    logger.info("收到停止信号，中止处理")
                    with self._lock:
                        progress.skipped_files += (len(file_pairs) - idx)
                    break
                
                with self._lock:
                    progress.current_file_index = idx
                
                success, result, error = self.process_single_file(
                    input_path, output_path, processor_func, progress
                )
                
                self._handle_file_result(
                    progress, input_path, output_path, success, result, error
                )
                
                # 更新进度
                self._update_progress(progress, "等待中")
        else:
            # 并行处理
            with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
                future_to_file = {
                    executor.submit(
                        self.process_single_file, 
                        input_path, output_path, processor_func, None
                    ): (input_path, output_path)
                    for input_path, output_path in file_pairs
                }
                
                for future in concurrent.futures.as_completed(future_to_file):
                    if self._should_stop:
                        logger.info("收到停止信号，取消剩余任务")
                        for f in future_to_file:
                            f.cancel()
                        break
                    
                    input_path, output_path = future_to_file[future]
                    
                    try:
                        success, result, error = future.result()
                        
                        with self._lock:
                            progress.current_file_index += 1
                        
                        self._handle_file_result(
                            progress, input_path, output_path, success, result, error
                        )
                        
                    except Exception as e:
                        logger.exception(f"获取任务结果失败: {input_path}")
                        self._handle_file_result(
                            progress, input_path, output_path, False, None, e
                        )
        
        self._is_running = False
        
        # 输出统计信息
        logger.info("=" * 50)
        logger.info("批量处理完成")
        logger.info(f"总文件数: {progress.total_files}")
        logger.info(f"成功: {progress.completed_files}")
        logger.info(f"失败: {progress.failed_files}")
        logger.info(f"跳过: {progress.skipped_files}")
        logger.info(f"成功率: {progress.success_rate:.1f}%")
        logger.info(f"耗时: {progress.elapsed_time:.1f} 秒")
        logger.info("=" * 50)
        
        return progress
    
    def process_folder(
        self,
        input_folder: str,
        output_folder: Optional[str] = None,
        recursive: bool = False,
        processor_func: Optional[Callable] = None,
    ) -> BatchProgress:
        """
        处理整个文件夹
        
        Args:
            input_folder: 输入文件夹路径
            output_folder: 输出文件夹路径（如为None则与输入相同）
            recursive: 是否递归处理子文件夹
            processor_func: 处理函数（如为None则需要在子类中实现）
            
        Returns:
            批量处理进度对象
        """
        input_folder = Path(input_folder).resolve()
        
        if not input_folder.exists():
            raise FileNotFoundError(f"输入文件夹不存在: {input_folder}")
        
        if not input_folder.is_dir():
            raise NotADirectoryError(f"输入路径不是文件夹: {input_folder}")
        
        # 确定输出文件夹
        if output_folder:
            output_folder = Path(output_folder).resolve()
            output_folder.mkdir(parents=True, exist_ok=True)
        else:
            output_folder = input_folder
        
        # 扫描文件
        file_pairs = self.scan_files(
            str(input_folder), 
            recursive=recursive,
            output_dir=str(output_folder)
        )
        
        if not file_pairs:
            logger.warning("没有找到待处理的文件")
            return BatchProgress()
        
        # 如果没有指定处理函数，需要子类实现
        if processor_func is None:
            raise NotImplementedError("必须提供processor_func或在子类中实现处理逻辑")
        
        # 执行批量处理
        return self.process_files(file_pairs, processor_func)
    
    def stop(self):
        """
        停止处理
        """
        self._should_stop = True
        logger.info("正在停止批量处理...")
    
    def is_running(self) -> bool:
        """
        检查是否正在运行
        
        Returns:
            是否正在运行
        """
        return self._is_running
    
    def generate_report(self, progress: BatchProgress, output_path: Optional[str] = None) -> str:
        """
        生成处理报告
        
        Args:
            progress: 进度对象
            output_path: 报告输出路径（可选）
            
        Returns:
            报告内容
        """
        lines = []
        lines.append("=" * 60)
        lines.append("          VideoSubtitleAI 批量处理报告")
        lines.append("=" * 60)
        lines.append("")
        
        # 基本信息
        lines.append("【基本信息】")
        lines.append(f"  处理时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
        lines.append(f"  总文件数: {progress.total_files}")
        lines.append(f"  成功: {progress.completed_files}")
        lines.append(f"  失败: {progress.failed_files}")
        lines.append(f"  跳过: {progress.skipped_files}")
        lines.append(f"  成功率: {progress.success_rate:.1f}%")
        lines.append(f"  总耗时: {progress.elapsed_time:.1f} 秒 ({progress.elapsed_time/60:.1f} 分钟)")
        lines.append("")
        
        # 成功文件列表
        if progress.results:
            lines.append("【成功处理的文件】")
            for idx, result in enumerate(progress.results, 1):
                lines.append(f"  {idx}. {Path(result['input_path']).name}")
                lines.append(f"     -> {result['output_path']}")
            lines.append("")
        
        # 失败文件列表
        if progress.errors:
            lines.append("【处理失败的文件】")
            for idx, error in enumerate(progress.errors, 1):
                lines.append(f"  {idx}. {Path(error['input_path']).name}")
                lines.append(f"     错误: {error['error']}")
                lines.append(f"     类型: {error['error_type']}")
            lines.append("")
        
        lines.append("=" * 60)
        lines.append("报告结束")
        lines.append("=" * 60)
        
        report_content = '\n'.join(lines)
        
        # 写入文件
        if output_path:
            try:
                with open(output_path, 'w', encoding='utf-8') as f:
                    f.write(report_content)
                logger.info(f"处理报告已保存: {output_path}")
            except Exception as e:
                logger.error(f"保存报告失败: {e}")
        
        return report_content