FrameTour-RenderWorker/handlers/base.py

# -*- coding: utf-8 -*-
"""
任务处理器基类

提供所有处理器共用的基础功能。
"""

import os
import json
import logging
import shutil
import tempfile
import subprocess
from abc import ABC
from typing import Optional, List, Dict, Any, Tuple, TYPE_CHECKING

from core.handler import TaskHandler
from domain.task import Task
from domain.result import TaskResult, ErrorCode
from domain.config import WorkerConfig
from services import storage
from constant import (
    HW_ACCEL_NONE, HW_ACCEL_QSV, HW_ACCEL_CUDA,
    VIDEO_ENCODE_PARAMS, VIDEO_ENCODE_PARAMS_QSV, VIDEO_ENCODE_PARAMS_CUDA
)

if TYPE_CHECKING:
    from services.api_client import APIClientV2

logger = logging.getLogger(__name__)


def get_video_encode_args(hw_accel: str = HW_ACCEL_NONE) -> List[str]:
    """
    根据硬件加速配置获取视频编码参数

    Args:
        hw_accel: 硬件加速类型 (none, qsv, cuda)

    Returns:
        FFmpeg 视频编码参数列表
    """
    if hw_accel == HW_ACCEL_QSV:
        params = VIDEO_ENCODE_PARAMS_QSV
        return [
            '-c:v', params['codec'],
            '-preset', params['preset'],
            '-profile:v', params['profile'],
            '-level', params['level'],
            '-global_quality', params['global_quality'],
            '-look_ahead', params['look_ahead'],
        ]
    elif hw_accel == HW_ACCEL_CUDA:
        params = VIDEO_ENCODE_PARAMS_CUDA
        return [
            '-c:v', params['codec'],
            '-preset', params['preset'],
            '-profile:v', params['profile'],
            '-level', params['level'],
            '-rc', params['rc'],
            '-cq', params['cq'],
            '-b:v', '0',  # 配合 vbr 模式使用 cq
        ]
    else:
        # 软件编码（默认）
        params = VIDEO_ENCODE_PARAMS
        return [
            '-c:v', params['codec'],
            '-preset', params['preset'],
            '-profile:v', params['profile'],
            '-level', params['level'],
            '-crf', params['crf'],
            '-pix_fmt', params['pix_fmt'],
        ]


def get_hwaccel_decode_args(hw_accel: str = HW_ACCEL_NONE) -> List[str]:
    """
    获取硬件加速解码参数（输入文件之前使用）

    Args:
        hw_accel: 硬件加速类型 (none, qsv, cuda)

    Returns:
        FFmpeg 硬件加速解码参数列表
    """
    if hw_accel == HW_ACCEL_CUDA:
        # CUDA 硬件加速解码
        # 注意：使用 cuda 作为 hwaccel，但输出到系统内存以便 CPU 滤镜处理
        return ['-hwaccel', 'cuda', '-hwaccel_output_format', 'cuda']
    elif hw_accel == HW_ACCEL_QSV:
        # QSV 硬件加速解码
        return ['-hwaccel', 'qsv', '-hwaccel_output_format', 'qsv']
    else:
        return []


def get_hwaccel_filter_prefix(hw_accel: str = HW_ACCEL_NONE) -> str:
    """
    获取硬件加速滤镜前缀（用于 hwdownload 从 GPU 到 CPU）

    注意：由于大多数复杂滤镜（如 lut3d, overlay, crop 等）不支持硬件表面，
    我们需要在滤镜链开始时将硬件表面下载到系统内存。

    Args:
        hw_accel: 硬件加速类型

    Returns:
        需要添加到滤镜链开头的 hwdownload 滤镜字符串
    """
    if hw_accel == HW_ACCEL_CUDA:
        return 'hwdownload,format=nv12,'
    elif hw_accel == HW_ACCEL_QSV:
        return 'hwdownload,format=nv12,'
    else:
        return ''


# v2 统一视频编码参数（兼容旧代码，使用软件编码）
VIDEO_ENCODE_ARGS = get_video_encode_args(HW_ACCEL_NONE)

# v2 统一音频编码参数
AUDIO_ENCODE_ARGS = [
    '-c:a', 'aac',
    '-b:a', '128k',
    '-ar', '48000',
    '-ac', '2',
]


def subprocess_args(include_stdout: bool = True) -> Dict[str, Any]:
    """
    创建跨平台的 subprocess 参数

    在 Windows 上使用 Pyinstaller --noconsole 打包时，需要特殊处理以避免弹出命令行窗口。

    Args:
        include_stdout: 是否包含 stdout 捕获

    Returns:
        subprocess.run 使用的参数字典
    """
    ret: Dict[str, Any] = {}

    # Windows 特殊处理
    if hasattr(subprocess, 'STARTUPINFO'):
        si = subprocess.STARTUPINFO()
        si.dwFlags |= subprocess.STARTF_USESHOWWINDOW
        ret['startupinfo'] = si
        ret['env'] = os.environ

    # 重定向 stdin 避免 "handle is invalid" 错误
    ret['stdin'] = subprocess.PIPE

    if include_stdout:
        ret['stdout'] = subprocess.PIPE

    return ret


def probe_video_info(video_file: str) -> Tuple[int, int, float]:
    """
    探测视频信息（宽度、高度、时长）

    Args:
        video_file: 视频文件路径

    Returns:
        (width, height, duration) 元组，失败返回 (0, 0, 0)
    """
    try:
        result = subprocess.run(
            [
                'ffprobe', '-v', 'error',
                '-select_streams', 'v:0',
                '-show_entries', 'stream=width,height:format=duration',
                '-of', 'csv=s=x:p=0',
                video_file
            ],
            capture_output=True,
            timeout=30,
            **subprocess_args(False)
        )

        if result.returncode != 0:
            logger.warning(f"ffprobe failed for {video_file}")
            return 0, 0, 0

        output = result.stdout.decode('utf-8').strip()
        if not output:
            return 0, 0, 0

        lines = output.split('\n')
        if len(lines) >= 2:
            wh = lines[0].strip()
            duration_str = lines[1].strip()
            width, height = wh.split('x')
            return int(width), int(height), float(duration_str)

        return 0, 0, 0

    except Exception as e:
        logger.warning(f"probe_video_info error: {e}")
        return 0, 0, 0


def probe_duration_json(file_path: str) -> Optional[float]:
    """
    使用 ffprobe JSON 输出探测媒体时长

    Args:
        file_path: 媒体文件路径

    Returns:
        时长（秒），失败返回 None
    """
    try:
        result = subprocess.run(
            [
                'ffprobe', '-v', 'error',
                '-show_entries', 'format=duration',
                '-of', 'json',
                file_path
            ],
            capture_output=True,
            timeout=30,
            **subprocess_args(False)
        )

        if result.returncode != 0:
            return None

        data = json.loads(result.stdout.decode('utf-8'))
        duration = data.get('format', {}).get('duration')
        return float(duration) if duration else None

    except Exception as e:
        logger.warning(f"probe_duration_json error: {e}")
        return None


class BaseHandler(TaskHandler, ABC):
    """
    任务处理器基类

    提供所有处理器共用的基础功能，包括：
    - 临时目录管理
    - 文件下载/上传
    - FFmpeg 命令执行
    - 日志记录
    """

    def __init__(self, config: WorkerConfig, api_client: 'APIClientV2'):
        """
        初始化处理器

        Args:
            config: Worker 配置
            api_client: API 客户端
        """
        self.config = config
        self.api_client = api_client

    def get_video_encode_args(self) -> List[str]:
        """
        获取当前配置的视频编码参数

        Returns:
            FFmpeg 视频编码参数列表
        """
        return get_video_encode_args(self.config.hw_accel)

    def get_hwaccel_decode_args(self) -> List[str]:
        """
        获取硬件加速解码参数（在输入文件之前使用）

        Returns:
            FFmpeg 硬件加速解码参数列表
        """
        return get_hwaccel_decode_args(self.config.hw_accel)

    def get_hwaccel_filter_prefix(self) -> str:
        """
        获取硬件加速滤镜前缀

        Returns:
            需要添加到滤镜链开头的 hwdownload 滤镜字符串
        """
        return get_hwaccel_filter_prefix(self.config.hw_accel)

    def before_handle(self, task: Task) -> None:
        """处理前钩子"""
        logger.debug(f"[task:{task.task_id}] Before handle: {task.task_type.value}")

    def after_handle(self, task: Task, result: TaskResult) -> None:
        """处理后钩子"""
        status = "success" if result.success else "failed"
        logger.debug(f"[task:{task.task_id}] After handle: {status}")

    def create_work_dir(self, task_id: str = None) -> str:
        """
        创建临时工作目录

        Args:
            task_id: 任务 ID（用于目录命名）

        Returns:
            工作目录路径
        """
        # 确保临时根目录存在
        os.makedirs(self.config.temp_dir, exist_ok=True)

        # 创建唯一的工作目录
        prefix = f"task_{task_id}_" if task_id else "task_"
        work_dir = tempfile.mkdtemp(dir=self.config.temp_dir, prefix=prefix)

        logger.debug(f"Created work directory: {work_dir}")
        return work_dir

    def cleanup_work_dir(self, work_dir: str) -> None:
        """
        清理临时工作目录

        Args:
            work_dir: 工作目录路径
        """
        if not work_dir or not os.path.exists(work_dir):
            return

        try:
            shutil.rmtree(work_dir)
            logger.debug(f"Cleaned up work directory: {work_dir}")
        except Exception as e:
            logger.warning(f"Failed to cleanup work directory {work_dir}: {e}")

    def download_file(self, url: str, dest: str, timeout: int = None) -> bool:
        """
        下载文件

        Args:
            url: 文件 URL
            dest: 目标路径
            timeout: 超时时间（秒）

        Returns:
            是否成功
        """
        if timeout is None:
            timeout = self.config.download_timeout

        try:
            result = storage.download_file(url, dest, timeout=timeout)
            if result:
                file_size = os.path.getsize(dest) if os.path.exists(dest) else 0
                logger.debug(f"Downloaded: {url} -> {dest} ({file_size} bytes)")
            return result
        except Exception as e:
            logger.error(f"Download failed: {url} -> {e}")
            return False

    def upload_file(
        self,
        task_id: str,
        file_type: str,
        file_path: str,
        file_name: str = None
    ) -> Optional[str]:
        """
        上传文件并返回访问 URL

        Args:
            task_id: 任务 ID
            file_type: 文件类型（video/audio/ts/mp4）
            file_path: 本地文件路径
            file_name: 文件名（可选）

        Returns:
            访问 URL，失败返回 None
        """
        # 获取上传 URL
        upload_info = self.api_client.get_upload_url(task_id, file_type, file_name)
        if not upload_info:
            logger.error(f"[task:{task_id}] Failed to get upload URL")
            return None

        upload_url = upload_info.get('uploadUrl')
        access_url = upload_info.get('accessUrl')

        if not upload_url:
            logger.error(f"[task:{task_id}] Invalid upload URL response")
            return None

        # 上传文件
        try:
            result = storage.upload_file(upload_url, file_path, timeout=self.config.upload_timeout)
            if result:
                file_size = os.path.getsize(file_path)
                logger.info(f"[task:{task_id}] Uploaded: {file_path} ({file_size} bytes)")
                return access_url
            else:
                logger.error(f"[task:{task_id}] Upload failed: {file_path}")
                return None
        except Exception as e:
            logger.error(f"[task:{task_id}] Upload error: {e}")
            return None

    def run_ffmpeg(
        self,
        cmd: List[str],
        task_id: str,
        timeout: int = None
    ) -> bool:
        """
        执行 FFmpeg 命令

        Args:
            cmd: FFmpeg 命令参数列表
            task_id: 任务 ID（用于日志）
            timeout: 超时时间（秒）

        Returns:
            是否成功
        """
        if timeout is None:
            timeout = self.config.ffmpeg_timeout

        # 日志记录命令（限制长度）
        cmd_str = ' '.join(cmd)
        if len(cmd_str) > 500:
            cmd_str = cmd_str[:500] + '...'
        logger.info(f"[task:{task_id}] FFmpeg: {cmd_str}")

        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                timeout=timeout,
                **subprocess_args(False)
            )

            if result.returncode != 0:
                stderr = result.stderr.decode('utf-8', errors='replace')[:1000]
                logger.error(f"[task:{task_id}] FFmpeg failed (code={result.returncode}): {stderr}")
                return False

            return True

        except subprocess.TimeoutExpired:
            logger.error(f"[task:{task_id}] FFmpeg timeout after {timeout}s")
            return False
        except Exception as e:
            logger.error(f"[task:{task_id}] FFmpeg error: {e}")
            return False

    def probe_duration(self, file_path: str) -> Optional[float]:
        """
        探测媒体文件时长

        Args:
            file_path: 文件路径

        Returns:
            时长（秒），失败返回 None
        """
        # 首先尝试 JSON 输出方式
        duration = probe_duration_json(file_path)
        if duration is not None:
            return duration

        # 回退到旧方式
        try:
            _, _, duration = probe_video_info(file_path)
            return float(duration) if duration else None
        except Exception as e:
            logger.warning(f"Failed to probe duration: {file_path} -> {e}")
            return None

    def get_file_size(self, file_path: str) -> int:
        """
        获取文件大小

        Args:
            file_path: 文件路径

        Returns:
            文件大小（字节）
        """
        try:
            return os.path.getsize(file_path)
        except Exception:
            return 0

    def ensure_file_exists(self, file_path: str, min_size: int = 0) -> bool:
        """
        确保文件存在且大小满足要求

        Args:
            file_path: 文件路径
            min_size: 最小大小（字节）

        Returns:
            是否满足要求
        """
        if not os.path.exists(file_path):
            return False
        return os.path.getsize(file_path) >= min_size