xiaoxia-saas/apps/worker/worker_app/tasks/generation.py

"""
视频生成任务
支持四种剪辑模式：一镜到底、画中画、口播、口播+画中画
"""

import logging
import os
import shutil
import subprocess  # nosec B404
import tempfile
from pathlib import Path
from urllib.parse import urlparse
from typing import Optional

import oss2

from worker_app.celery_app import celery_app

OUTPUT_WIDTH = 1280
OUTPUT_HEIGHT = 720
OUTPUT_FPS = 25.0
OUTPUT_DURATION_SECONDS = 5.0
FFMPEG_BIN = shutil.which("ffmpeg") or "ffmpeg"
FFPROBE_BIN = shutil.which("ffprobe") or "ffprobe"
GENERATED_FILES_DIR = Path(os.getenv("GENERATED_FILES_DIR", "/app/generated"))
GENERATED_FILES_URL_PREFIX = os.getenv("GENERATED_FILES_URL_PREFIX", "/generated-files")
PUBLIC_API_BASE_URL = os.getenv("PUBLIC_API_BASE_URL", "https://api.xiaoxiajianji.com").rstrip("/")

logger = logging.getLogger(__name__)


def _run_ffmpeg(command: list[str]) -> None:
    """执行 FFmpeg 命令"""
    subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)  # nosec B603


def _oss_settings() -> tuple[str, str, str, str] | None:
    """获取 OSS 配置"""
    access_key_id = os.getenv("OSS_ACCESS_KEY_ID")
    access_key_secret = os.getenv("OSS_ACCESS_KEY_SECRET")
    endpoint = os.getenv("OSS_ENDPOINT")
    bucket_name = os.getenv("OSS_BUCKET_NAME")
    if not all([access_key_id, access_key_secret, endpoint, bucket_name]):
        return None
    return access_key_id, access_key_secret, endpoint, bucket_name


def _oss_bucket() -> oss2.Bucket | None:
    """获取 OSS Bucket"""
    settings = _oss_settings()
    if settings is None:
        return None
    access_key_id, access_key_secret, endpoint, bucket_name = settings
    return oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)


def _normalize_storage_key(storage_key_or_url: str) -> str:
    """标准化存储键"""
    if storage_key_or_url.startswith(("http://", "https://")):
        return urlparse(storage_key_or_url).path.lstrip("/")
    return storage_key_or_url.lstrip("/")


def _download_asset(asset_storage_key: str, local_path: Path) -> bool:
    """下载素材文件"""
    bucket = _oss_bucket()
    if bucket is None:
        return False
    try:
        bucket.get_object_to_file(_normalize_storage_key(asset_storage_key), str(local_path))
        return local_path.exists() and local_path.stat().st_size > 0
    except Exception:
        return False


def _probe_duration(local_path: Path) -> float:
    """获取视频时长"""
    try:
        result = subprocess.run(
            [
                FFPROBE_BIN, "-v", "error",
                "-show_entries", "format=duration",
                "-of", "default=noprint_wrappers=1:nokey=1",
                str(local_path),
            ],
            check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True,
        )  # nosec B603
        return round(float(result.stdout.strip()), 3)
    except Exception:
        return OUTPUT_DURATION_SECONDS


def _create_fallback_clip(output_path: Path, title: str) -> None:
    """创建 fallback 视频（无素材时）"""
    safe_title = title.replace(":", "\\:").replace("'", "\\'")[:80]
    _run_ffmpeg(
        [
            FFMPEG_BIN, "-y", "-f", "lavfi",
            "-i", f"color=c=#111827:s={OUTPUT_WIDTH}x{OUTPUT_HEIGHT}:d={OUTPUT_DURATION_SECONDS}:r={int(OUTPUT_FPS)}",
            "-vf", f"drawtext=text='{safe_title}':fontcolor=white:fontsize=48:x=(w-text_w)/2:y=(h-text_h)/2",
            "-c:v", "libx264", "-pix_fmt", "yuv420p", "-movflags", "+faststart",
            str(output_path),
        ]
    )


def _download_voice_asset(voice_library_id: str, local_path: Path) -> bool:
    """下载配音文件"""
    if not voice_library_id:
        return False
    bucket = _oss_bucket()
    if bucket is None:
        return False
    storage_key = f"voice/{voice_library_id}.mp3"
    try:
        bucket.get_object_to_file(_normalize_storage_key(storage_key), str(local_path))
        return local_path.exists() and local_path.stat().st_size > 0
    except Exception:
        return False


def _download_library_assets(
    asset_library_id: str,
    temp_path: Path,
    video_extensions: tuple = (".mp4", ".mov", ".avi", ".mkv", ".webm"),
) -> list[str]:
    """
    从素材库下载所有视频素材

    Args:
        asset_library_id: 素材库 ID
        temp_path: 临时目录路径
        video_extensions: 支持的视频扩展名

    Returns:
        下载成功的视频文件路径列表
    """
    # 导入模型和会话
    try:
        from packages.adapters.sqlalchemy_impl.models import AssetModel
        from worker_app.db import SessionLocal

        session = SessionLocal()

        try:
            # 查询素材库中的视频素材
            assets = session.query(AssetModel).filter(
                AssetModel.asset_library_id == asset_library_id,
                AssetModel.status == "ready",
                AssetModel.file_type.in_(["video", "video/mp4", "video/quicktime"]),
            ).order_by(AssetModel.created_at).all()

            if not assets:
                logger.info(f"No video assets found in library {asset_library_id}")
                return []

            downloaded_videos = []
            for i, asset in enumerate(assets):
                # 获取文件 URL 或 storage_key
                storage_key = asset.file_url if asset.file_url else None
                if not storage_key:
                    continue

                local_file = temp_path / f"asset_{i}_{asset.id}.mp4"
                if _download_asset(storage_key, local_file):
                    downloaded_videos.append(str(local_file))
                    logger.info(f"Downloaded asset: {asset.name} -> {local_file}")
                else:
                    logger.warning(f"Failed to download asset: {asset.name}")

            return downloaded_videos
        finally:
            session.close()
    except Exception as e:
        logger.error(f"Error downloading library assets: {e}")
        return []


def _process_with_editing_mode(
    video_paths: list[str],
    audio_path: Optional[str],
    mode: str,
    output_path: Path,
) -> None:
    """根据剪辑模式处理视频"""
    from video_processing.editing_modes import (
        EditingMode,
        EditingModeConfig,
        EditingModeProcessor,
        PIPPosition,
    )

    config = EditingModeConfig(
        mode=EditingMode(mode),
        output_width=OUTPUT_WIDTH,
        output_height=OUTPUT_HEIGHT,
        output_fps=int(OUTPUT_FPS),
        pip_position=PIPPosition.TOP_RIGHT,
        pip_scale=0.25,
        transition_duration=0.5,
    )

    processor = EditingModeProcessor(config=config)
    processor.process(
        video_paths=video_paths,
        audio_path=audio_path,
        output_path=str(output_path),
    )


@celery_app.task(bind=True, name="worker.generate_video", max_retries=2)
def generate_video(self, task_id: str) -> dict:
    """
    生成视频任务

    Args:
        task_id: 任务 ID（从数据库加载完整任务信息）

    Returns:
        生成结果字典
    """
    from packages.domain import EditingMode, GeneratedVideo, GenerationTaskStatus
    from packages.adapters.sqlalchemy_impl.generation_task_repository import (
        SQLAlchemyGenerationTaskRepository,
    )
    from worker_app.db import SessionLocal

    # 从数据库加载任务信息
    session = SessionLocal()
    try:
        task_repo = SQLAlchemyGenerationTaskRepository(session)
        gen_task = task_repo.get(task_id)
        if gen_task is None:
            return {"status": "failed", "error": f"generation task {task_id} not found"}
        project_id = gen_task.project_id
        asset_library_id = gen_task.asset_library_id
        voice_library_id = gen_task.voice_library_id or ""
        mode = gen_task.strategy_id or "one_take"
    finally:
        session.close()

    try:
        editing_mode = EditingMode(mode)
    except ValueError:
        editing_mode = EditingMode.ONE_TAKE

    output_name = f"generated-{task_id}.mp4"
    storage_key = f"generated/projects/{project_id}/tasks/{task_id}/{output_name}"

    try:
        with tempfile.TemporaryDirectory(prefix="xiaoxia-generation-") as temp_dir:
            temp_path = Path(temp_dir)
            output_path = temp_path / output_name

            # 从素材库下载视频素材
            downloaded_videos = _download_library_assets(asset_library_id, temp_path)

            audio_path = None
            if voice_library_id:
                local_audio = temp_path / "voice.mp3"
                if _download_voice_asset(voice_library_id, local_audio):
                    audio_path = str(local_audio)

            if downloaded_videos:
                _process_with_editing_mode(
                    video_paths=downloaded_videos,
                    audio_path=audio_path,
                    mode=editing_mode.value,
                    output_path=output_path,
                )
            else:
                _create_fallback_clip(output_path, f"Generated Video {task_id[:8]}")

            file_size = output_path.stat().st_size
            duration = _probe_duration(output_path)

            return {
                "status": "completed",
                "task_id": task_id,
                "output_path": str(output_path),
                "file_size": file_size,
                "duration": duration,
                "width": OUTPUT_WIDTH,
                "height": OUTPUT_HEIGHT,
                "mode": editing_mode.value,
            }
    except Exception as error:
        logger.error(f"Video generation failed: {error}")
        return {
            "status": "failed",
            "task_id": task_id,
            "error": str(error),
        }