video-gen/video-gen-api/app/services/shot_video_split_service.py

from __future__ import annotations

import asyncio
import os
import subprocess
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

from app.config import settings
from app.services.upload_video_asset_service import (
    build_upload_url_from_path,
    ensure_shot_segment_dir,
    get_ffmpeg_bin,
)


@dataclass(slots=True)
class ShotSplitResult:
    url: str
    path: str
    file_size_bytes: int
    temporary_path: str | None = None


def _date_dir_from_segment_id(segment_id: str) -> str:
    return datetime.now().strftime("%Y/%m/%d")


def _safe_unlink(path: Path) -> None:
    try:
        if path.exists():
            path.unlink()
    except Exception:
        pass


def split_video_segment(
    *,
    source_path: str | Path,
    segment_id: str,
    start_second: float,
    end_second: float,
    date_dir: str | None = None,
    attempt_key: str | None = None,
    finalize: bool = False,
) -> ShotSplitResult:
    """执行 ffmpeg 切片。

    默认只产出 attempt 专属临时文件；调用方完成 Redis/DB fencing 后再原子移动到正式路径，
    避免旧 Worker 在恢复任务接管后覆盖新结果。
    """
    source_path = Path(source_path)
    if not source_path.exists():
        raise RuntimeError(f"ffmpeg 拆镜失败：源视频不存在 {source_path}")

    start = max(float(start_second), 0.0)
    end = max(float(end_second), 0.0)
    duration = end - start
    if duration <= 0:
        raise RuntimeError(
            f"ffmpeg 拆镜失败：非法时间范围 start_second={start_second}, end_second={end_second}"
        )

    date_dir = date_dir or _date_dir_from_segment_id(segment_id)
    output_dir = ensure_shot_segment_dir(date_dir)
    output_path = output_dir / f"{segment_id}.mp4"
    suffix = str(attempt_key or "default").replace("/", "_").replace(":", "_")[:80]
    part_path = output_dir / f"{segment_id}.{suffix}.part.mp4"
    _safe_unlink(part_path)

    timeout = int(getattr(settings, "SHOT_FFMPEG_TIMEOUT_SECONDS", 120) or 120)
    cmd = [
        get_ffmpeg_bin(),
        "-y",

        # 先 seek 到起始秒，再按 duration 切割，避免 -to 在不同 ffmpeg 参数位置下语义不一致。
        "-ss",
        f"{start:.3f}",
        "-i",
        str(source_path),
        "-t",
        f"{duration:.3f}",

        # 只取主视频流，音频可选，避免 map 0 把字幕/数据流带进去导致 mp4 封装失败。
        "-map",
        "0:v:0",
        "-map",
        "0:a:0?",

        # 当前是拆镜片段，重编码更稳，避免关键帧不准导致片段首尾异常。
        "-c:v",
        "libx264",
        "-preset",
        "veryfast",
        "-crf",
        "23",
        "-pix_fmt",
        "yuv420p",
        "-c:a",
        "aac",
        "-b:a",
        "128k",
        "-movflags",
        "+faststart",

        # 即使临时文件扩展名未来被改坏，也强制指定 mp4 muxer。
        "-f",
        "mp4",
        str(part_path),
    ]

    try:
        completed = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=timeout,
            check=False,
        )
    except subprocess.TimeoutExpired as exc:
        _safe_unlink(part_path)
        raise RuntimeError(
            f"ffmpeg 拆镜超时：timeout={timeout}s, start={start:.3f}, end={end:.3f}"
        ) from exc

    if completed.returncode != 0:
        _safe_unlink(part_path)
        raise RuntimeError(
            f"ffmpeg 拆镜失败: {completed.stderr.strip() or completed.stdout.strip()}"
        )
    if not part_path.exists() or part_path.stat().st_size <= 0:
        _safe_unlink(part_path)
        raise RuntimeError("ffmpeg 拆镜失败：输出文件为空")

    if finalize:
        os.replace(part_path, output_path)
        return ShotSplitResult(
            url=build_upload_url_from_path(output_path),
            path=str(output_path),
            file_size_bytes=output_path.stat().st_size,
            temporary_path=None,
        )

    return ShotSplitResult(
        url=build_upload_url_from_path(output_path),
        path=str(output_path),
        file_size_bytes=part_path.stat().st_size,
        temporary_path=str(part_path),
    )


def finalize_split_result(result: ShotSplitResult) -> ShotSplitResult:
    if not result.temporary_path:
        return result
    temp_path = Path(result.temporary_path)
    final_path = Path(result.path)
    if not temp_path.exists() or temp_path.stat().st_size <= 0:
        raise RuntimeError("ffmpeg 拆镜临时结果不存在或为空")
    final_path.parent.mkdir(parents=True, exist_ok=True)
    os.replace(temp_path, final_path)
    return ShotSplitResult(
        url=result.url,
        path=str(final_path),
        file_size_bytes=final_path.stat().st_size,
        temporary_path=None,
    )


def cleanup_split_result(result: ShotSplitResult | None) -> None:
    if result and result.temporary_path:
        _safe_unlink(Path(result.temporary_path))


async def split_video_segment_async(
    *,
    source_path: str | Path,
    segment_id: str,
    start_second: float,
    end_second: float,
    date_dir: str | None = None,
    attempt_key: str | None = None,
    finalize: bool = False,
) -> ShotSplitResult:
    return await asyncio.to_thread(
        split_video_segment,
        source_path=source_path,
        segment_id=segment_id,
        start_second=start_second,
        end_second=end_second,
        date_dir=date_dir,
        attempt_key=attempt_key,
        finalize=finalize,
    )