xiaoxia-saas/apps/api/app/services/ai_avatar_render_service.py

"""AI数字人渲染合成 Service — #1798.

职责：
- 创建/查询/取消渲染任务
- 调用 Celery 异步任务执行渲染
- B-roll 合成 + 标题叠加 + 封面提取
- 用户隔离
"""

from __future__ import annotations

import logging
import os
import tempfile
import uuid
from datetime import datetime, timezone
from typing import Any, Optional

from sqlalchemy.orm import Session

from packages.adapters.sqlalchemy_impl.models import (
    AiAvatarRenderJob,
    LipsyncJobModel,
    ScriptModel,
)
from packages.domain.video_filter_builder import (
    build_cover_extract_command,
    build_title_drawtext_filter,
)
from packages.shared.storage import get_shared_storage_service

logger = logging.getLogger(__name__)


class AiAvatarRenderError(Exception):
    """渲染服务异常."""

    def __init__(self, message: str, code: str = "RenderError"):
        self.code = code
        super().__init__(message)


class AiAvatarRenderService:
    """AI数字人渲染合成 Service."""

    def __init__(self, db: Session):
        self.db = db

    # ── 创建任务 ──────────────────────────────────────────────────────────

    def create_render_job(
        self,
        *,
        user_id: str,
        lipsync_job_id: str,
        script_id: str,
        b_roll_segments: list[dict[str, Any]],
        title_config: dict[str, Any],
        cover_config: dict[str, Any],
        project_id: str = "",
    ) -> AiAvatarRenderJob:
        """创建渲染任务.

        Raises:
            AiAvatarRenderError: 校验失败
        """
        # 1. 验证对口型任务
        lipsync_job = (
            self.db.query(LipsyncJobModel)
            .filter(
                LipsyncJobModel.id == lipsync_job_id,
                LipsyncJobModel.user_id == user_id,
            )
            .first()
        )
        if lipsync_job is None:
            raise AiAvatarRenderError("对口型任务不存在", code="LipsyncJobNotFound")
        if lipsync_job.status != "completed":
            raise AiAvatarRenderError(
                f"对口型任务状态为 {lipsync_job.status}，仅 completed 状态可渲染",
                code="LipsyncJobNotCompleted",
            )
        if not lipsync_job.output_video_url:
            raise AiAvatarRenderError("对口型任务输出视频 URL 为空", code="LipsyncJobNoOutput")

        # 2. 验证文案归属
        script = (
            self.db.query(ScriptModel)
            .filter(
                ScriptModel.id == script_id,
                ScriptModel.user_id == user_id,
            )
            .first()
        )
        if script is None:
            raise AiAvatarRenderError("文案不存在或无权访问", code="ScriptNotFound")

        # 3. 创建渲染任务
        job_id = str(uuid.uuid4())
        job = AiAvatarRenderJob(
            id=job_id,
            user_id=user_id,
            project_id=project_id,
            lipsync_job_id=lipsync_job_id,
            script_id=script_id,
            b_roll_segments=[s if isinstance(s, dict) else s.model_dump() for s in b_roll_segments],
            title_config=title_config,
            cover_config=cover_config,
            status="pending",
        )
        self.db.add(job)
        self.db.flush()

        job.submitted_at = datetime.now(timezone.utc)
        self.db.commit()
        self.db.refresh(job)
        return job

    # ── 查询任务 ──────────────────────────────────────────────────────────

    def get_render_job(self, job_id: str, user_id: str) -> Optional[AiAvatarRenderJob]:
        """获取渲染任务详情（用户隔离）."""
        return (
            self.db.query(AiAvatarRenderJob)
            .filter(
                AiAvatarRenderJob.id == job_id,
                AiAvatarRenderJob.user_id == user_id,
            )
            .first()
        )

    def list_render_jobs(
        self,
        *,
        user_id: str,
        project_id: str = "",
        status: str = "",
        offset: int = 0,
        limit: int = 20,
    ) -> tuple[list[AiAvatarRenderJob], int]:
        """获取渲染任务列表（分页 + 用户隔离）."""
        query = self.db.query(AiAvatarRenderJob).filter(AiAvatarRenderJob.user_id == user_id)
        if project_id:
            query = query.filter(AiAvatarRenderJob.project_id == project_id)
        if status:
            query = query.filter(AiAvatarRenderJob.status == status)

        total = query.count()
        items = query.order_by(AiAvatarRenderJob.created_at.desc()).offset(offset).limit(limit).all()
        return items, total

    # ── 取消任务 ──────────────────────────────────────────────────────────

    def cancel_render_job(self, job_id: str, user_id: str) -> Optional[AiAvatarRenderJob]:
        """取消渲染任务（仅 pending 状态可取消）."""
        job = self.get_render_job(job_id, user_id)
        if job is None:
            return None
        if job.status in ("pending", "submitted"):
            job.status = "cancelled"
            job.updated_at = datetime.now(timezone.utc)
            self.db.commit()
            self.db.refresh(job)
        return job

    # ── 重试任务 ──────────────────────────────────────────────────────────

    def retry_render_job(self, job_id: str, user_id: str) -> Optional[AiAvatarRenderJob]:
        """重试失败的渲染任务."""
        job = self.get_render_job(job_id, user_id)
        if job is None:
            return None
        if job.status != "failed":
            return None
        job.status = "pending"
        job.progress = 0
        job.error_message = ""
        job.output_video_url = ""
        job.output_cover_url = ""
        job.output_duration = 0.0
        job.started_at = None
        job.completed_at = None
        job.updated_at = datetime.now(timezone.utc)
        self.db.commit()
        self.db.refresh(job)
        return job

    # ── 执行渲染（Celery 异步调用） ──────────────────────────────────────

    def execute_render(self, job_id: str) -> None:
        """执行渲染管线.

        由 Celery 异步任务调用，流程:
        1. 下载对口型输出视频 (20%)
        2. 构建 FFmpeg 滤镜链 (40%)
        3. 执行 FFmpeg 渲染 (80%)
        4. 提取封面 (90%)
        5. 上传到 OSS (95%)
        6. 更新任务状态 (100%)
        """
        job = self.db.query(AiAvatarRenderJob).filter(AiAvatarRenderJob.id == job_id).first()
        if job is None:
            logger.error("渲染任务不存在: %s", job_id)
            return

        if job.status == "cancelled":
            logger.info("渲染任务已取消: %s", job_id)
            return

        try:
            # 更新状态为 processing
            job.status = "processing"
            job.started_at = datetime.now(timezone.utc)
            job.progress = 5
            job.updated_at = datetime.now(timezone.utc)
            self.db.commit()

            # 获取对口型任务信息
            lipsync_job = self.db.query(LipsyncJobModel).filter(LipsyncJobModel.id == job.lipsync_job_id).first()
            if lipsync_job is None:
                raise AiAvatarRenderError("关联的对口型任务不存在", code="LipsyncJobNotFound")

            # 1. 下载对口型输出视频 (20%)
            input_video_path = self._download_video(lipsync_job.output_video_url)
            job.progress = 20
            self.db.commit()

            # 2. 构建 FFmpeg 滤镜链 (40%)
            from packages.domain.video_filter_builder import build_broll_overlay_filter

            filter_complex = build_broll_overlay_filter(
                b_roll_segments=job.b_roll_segments,
                video_duration=lipsync_job.output_duration,
            )

            # 标题叠加
            title_filter = build_title_drawtext_filter(job.title_config)
            if title_filter:
                if filter_complex:
                    filter_complex += f"[vout]{title_filter}[vout_titled];"
                else:
                    filter_complex = f"[0:v]{title_filter}[vout_titled];"

            # 清理末尾分号
            if filter_complex.endswith(";"):
                filter_complex = filter_complex[:-1]

            # 最终输出标签
            final_label = "vout_titled" if title_filter else ("vout" if filter_complex else None)

            job.progress = 40
            self.db.commit()

            # 3. 执行 FFmpeg 渲染 (80%)
            with tempfile.TemporaryDirectory() as tmpdir:
                output_video_path = os.path.join(tmpdir, "output.mp4")

                cmd = self._build_ffmpeg_command(
                    input_video=input_video_path,
                    b_roll_segments=job.b_roll_segments,
                    filter_complex=filter_complex,
                    final_label=final_label,
                    output_path=output_video_path,
                )

                exit_code = os.system(cmd)
                if exit_code != 0:
                    raise AiAvatarRenderError(f"FFmpeg 渲染失败，退出码: {exit_code}", code="FFmpegFailed")

                job.progress = 80
                self.db.commit()

                # 4. 提取封面 (90%)
                cover_path = ""
                if job.cover_config:
                    cover_path = os.path.join(tmpdir, "cover.jpg")
                    cover_cmd = build_cover_extract_command(job.cover_config, cover_path)
                    cover_cmd = cover_cmd.replace("INPUT_VIDEO", output_video_path)
                    cover_exit = os.system(cover_cmd)
                    if cover_exit != 0:
                        logger.warning("封面提取失败，跳过: %s", cover_cmd)
                        cover_path = ""

                job.progress = 90
                self.db.commit()

                # 5. 上传到 OSS (95%)
                output_video_url = self._upload_to_oss(output_video_path, f"ai-avatar/{job_id}/output.mp4")
                job.output_video_url = output_video_url

                if cover_path:
                    output_cover_url = self._upload_to_oss(cover_path, f"ai-avatar/{job_id}/cover.jpg")
                    job.output_cover_url = output_cover_url

                # 获取输出视频时长
                job.output_duration = lipsync_job.output_duration
                job.progress = 95
                self.db.commit()

            # 6. 完成
            job.status = "completed"
            job.progress = 100
            job.completed_at = datetime.now(timezone.utc)
            job.updated_at = datetime.now(timezone.utc)
            self.db.commit()
            logger.info("渲染任务完成: %s", job_id)

        except AiAvatarRenderError as exc:
            job.status = "failed"
            job.error_message = str(exc)
            job.updated_at = datetime.now(timezone.utc)
            self.db.commit()
            logger.error("渲染任务失败 [%s]: %s", job_id, exc)
        except Exception as exc:
            job.status = "failed"
            job.error_message = f"渲染异常: {str(exc)}"
            job.updated_at = datetime.now(timezone.utc)
            self.db.commit()
            logger.exception("渲染任务异常 [%s]", job_id)

    def _download_video(self, url: str) -> str:
        """下载视频到临时文件."""
        import httpx

        tmp = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
        try:
            with httpx.Client(timeout=120) as client:
                resp = client.get(url)
                resp.raise_for_status()
                tmp.write(resp.content)
            return tmp.name
        except Exception:
            if os.path.exists(tmp.name):
                os.unlink(tmp.name)
            raise

    def _build_ffmpeg_command(
        self,
        *,
        input_video: str,
        b_roll_segments: list[dict[str, Any]],
        filter_complex: str,
        final_label: Optional[str],
        output_path: str,
    ) -> str:
        """构建 FFmpeg 命令."""
        # 输入文件
        inputs = f"-i {input_video}"
        for seg in b_roll_segments:
            asset_url = seg.get("asset_url", "")
            if asset_url:
                inputs += f" -i {asset_url}"

        # 滤镜
        if filter_complex and final_label:
            filter_arg = f'-filter_complex "{filter_complex}" -map "[{final_label}]"'
        elif filter_complex:
            filter_arg = f'-filter_complex "{filter_complex}"'
        else:
            filter_arg = ""

        return f"ffmpeg {inputs} {filter_arg} -c:v libx264 -preset fast -crf 23 -y {output_path}"

    def _upload_to_oss(self, local_path: str, oss_key: str) -> str:
        """上传文件到 OSS，返回 URL.

        使用 SharedStorageService 统一存储服务。
        """
        storage = get_shared_storage_service()
        url = storage.upload_file_smart(local_path, oss_key)
        if url is None:
            raise AiAvatarRenderError(
                f"上传文件到 OSS 失败: {oss_key}",
                code="OSSUploadFailed",
            )
        logger.info("上传文件到 OSS 成功: %s -> %s", local_path, url)
        return url