FrameTour-RenderWorker/handlers/base.py

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

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

import os
import logging
import shutil
import tempfile
import subprocess
from abc import ABC
from typing import Optional, List, TYPE_CHECKING

from core.handler import TaskHandler
from domain.task import Task
from domain.result import TaskResult, ErrorCode
from domain.config import WorkerConfig
from util import oss
from util.ffmpeg import subprocess_args

if TYPE_CHECKING:
    from services.api_client import APIClientV2

logger = logging.getLogger(__name__)


# v2 统一视频编码参数（来自集成文档）
VIDEO_ENCODE_ARGS = [
    '-c:v', 'libx264',
    '-preset', 'medium',
    '-profile:v', 'main',
    '-level', '4.0',
    '-crf', '23',
    '-pix_fmt', 'yuv420p',
]

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


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 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 = oss.download_from_oss(url, dest)
            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 = oss.upload_to_oss(upload_url, file_path)
            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
        """
        try:
            from util.ffmpeg import probe_video_info
            _, _, 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