FrameTour-RenderWorker/handlers/finalize_mp4.py

# -*- coding: utf-8 -*-
"""
最终 MP4 合并处理器

处理 FINALIZE_MP4 任务，将所有 TS 分片合并为最终可下载的 MP4 文件。
"""

import os
import logging
from typing import List

from handlers.base import BaseHandler
from domain.task import Task, TaskType
from domain.result import TaskResult, ErrorCode

logger = logging.getLogger(__name__)


class FinalizeMp4Handler(BaseHandler):
    """
    最终 MP4 合并处理器

    职责：
    - 下载所有 TS 分片
    - 使用 concat demuxer 合并
    - 产出最终 MP4（remux，不重编码）
    - 上传 MP4 产物

    关键约束：
    - 优先使用 remux（复制流，不重新编码）
    - 使用 aac_adtstoasc bitstream filter 处理音频
    """

    def get_supported_type(self) -> TaskType:
        return TaskType.FINALIZE_MP4

    def handle(self, task: Task) -> TaskResult:
        """处理 MP4 合并任务"""
        work_dir = self.create_work_dir(task.task_id)

        try:
            # 获取 TS 列表
            ts_list = task.get_ts_list()
            m3u8_url = task.get_m3u8_url()

            if not ts_list and not m3u8_url:
                return TaskResult.fail(
                    ErrorCode.E_SPEC_INVALID,
                    "Missing tsList or m3u8Url"
                )

            output_file = os.path.join(work_dir, 'final.mp4')

            if ts_list:
                # 方式1：使用 TS 列表
                result = self._process_ts_list(task, work_dir, ts_list, output_file)
            else:
                # 方式2：使用 m3u8 URL
                result = self._process_m3u8(task, work_dir, m3u8_url, output_file)

            if not result.success:
                return result

            # 验证输出文件
            if not self.ensure_file_exists(output_file, min_size=4096):
                return TaskResult.fail(
                    ErrorCode.E_FFMPEG_FAILED,
                    "MP4 output file is missing or too small"
                )

            # 获取文件大小
            file_size = self.get_file_size(output_file)

            # 上传产物
            mp4_url = self.upload_file(task.task_id, 'mp4', output_file)
            if not mp4_url:
                return TaskResult.fail(
                    ErrorCode.E_UPLOAD_FAILED,
                    "Failed to upload MP4"
                )

            return TaskResult.ok({
                'mp4Url': mp4_url,
                'fileSizeBytes': file_size
            })

        except Exception as e:
            logger.error(f"[task:{task.task_id}] Unexpected error: {e}", exc_info=True)
            return TaskResult.fail(ErrorCode.E_UNKNOWN, str(e))

        finally:
            self.cleanup_work_dir(work_dir)

    def _process_ts_list(
        self,
        task: Task,
        work_dir: str,
        ts_list: List[str],
        output_file: str
    ) -> TaskResult:
        """
        使用 TS 列表处理

        Args:
            task: 任务实体
            work_dir: 工作目录
            ts_list: TS URL 列表
            output_file: 输出文件路径

        Returns:
            TaskResult
        """
        # 1. 下载所有 TS 分片
        ts_files = []
        for i, ts_url in enumerate(ts_list):
            ts_file = os.path.join(work_dir, f'seg_{i}.ts')
            if not self.download_file(ts_url, ts_file):
                return TaskResult.fail(
                    ErrorCode.E_INPUT_UNAVAILABLE,
                    f"Failed to download TS segment {i}: {ts_url}"
                )
            ts_files.append(ts_file)

        logger.info(f"[task:{task.task_id}] Downloaded {len(ts_files)} TS segments")

        # 2. 创建 concat 文件列表
        concat_file = os.path.join(work_dir, 'concat.txt')
        with open(concat_file, 'w', encoding='utf-8') as f:
            for ts_file in ts_files:
                # 路径中的反斜杠需要转义或使用正斜杠
                ts_path = ts_file.replace('\\', '/')
                f.write(f"file '{ts_path}'\n")

        # 3. 构建合并命令（remux，不重编码）
        cmd = [
            'ffmpeg', '-y', '-hide_banner',
            '-f', 'concat',
            '-safe', '0',
            '-i', concat_file,
            '-c', 'copy',  # 复制流，不重编码
            '-bsf:a', 'aac_adtstoasc',  # 音频 bitstream filter
            output_file
        ]

        # 4. 执行 FFmpeg
        if not self.run_ffmpeg(cmd, task.task_id):
            return TaskResult.fail(
                ErrorCode.E_FFMPEG_FAILED,
                "MP4 concatenation failed"
            )

        return TaskResult.ok({})

    def _process_m3u8(
        self,
        task: Task,
        work_dir: str,
        m3u8_url: str,
        output_file: str
    ) -> TaskResult:
        """
        使用 m3u8 URL 处理

        Args:
            task: 任务实体
            work_dir: 工作目录
            m3u8_url: m3u8 URL
            output_file: 输出文件路径

        Returns:
            TaskResult
        """
        # 构建命令
        cmd = [
            'ffmpeg', '-y', '-hide_banner',
            '-protocol_whitelist', 'file,http,https,tcp,tls',
            '-i', m3u8_url,
            '-c', 'copy',
            '-bsf:a', 'aac_adtstoasc',
            output_file
        ]

        # 执行 FFmpeg
        if not self.run_ffmpeg(cmd, task.task_id):
            return TaskResult.fail(
                ErrorCode.E_FFMPEG_FAILED,
                "MP4 conversion from m3u8 failed"
            )

        return TaskResult.ok({})