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 typing import Optional
from urllib.parse import urlparse

import oss2
from worker_app.celery_app import celery_app
from worker_app.db import SessionLocal

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 _update_task_status(task_id: str, status_action: str, **kwargs) -> bool:
    """更新 GenerationTask 状态（独立 session，异常不向外抛出）。

    Args:
        task_id: 任务 ID
        status_action: 状态动作名，如 "mark_processing" / "mark_completed" / "mark_failed"
        **kwargs: 传递给对应方法的参数

    Returns:
        True 表示更新成功，False 表示更新失败
    """
    try:
        from packages.adapters.sqlalchemy_impl.generation_task_repository import (
            SQLAlchemyGenerationTaskRepository,
        )

        session = SessionLocal()
        try:
            repo = SQLAlchemyGenerationTaskRepository(session)
            task = repo.get(task_id)
            if task is None:
                logger.warning("更新任务状态失败：任务不存在 task_id=%s", task_id)
                return False

            action = getattr(task, status_action, None)
            if action is None:
                logger.warning("未知的状态动作: %s", status_action)
                return False

            action(**kwargs)
            repo.update(task)
            logger.info("GenerationTask 状态更新成功: task_id=%s action=%s", task_id, status_action)
            return True
        finally:
            session.close()
    except Exception as e:
        logger.error(
            "更新 GenerationTask 状态异常: task_id=%s action=%s error=%s",
            task_id,
            status_action,
            e,
            exc_info=True,
        )
        return False


# ── FFmpeg / OSS helpers ─────────────────────────────────────────────────────


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"),
    asset_ids: list[str] | None = None,
) -> list[str]:
    """
    从素材库下载视频素材

    Args:
        asset_library_id: 素材库 ID
        temp_path: 临时目录路径
        video_extensions: 支持的视频扩展名
        asset_ids: 指定素材 ID 列表，为空则下载全部 ready 视频素材

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

        session = SessionLocal()

        try:
            # 查询素材库中的视频素材
            query = session.query(AssetModel).filter(
                AssetModel.asset_library_id == asset_library_id,
                AssetModel.status == "ready",
                AssetModel.file_type.in_(["video", "video/mp4", "video/quicktime"]),
            )
            # 如果指定了 asset_ids，则只下载这些素材
            if asset_ids:
                query = query.filter(AssetModel.id.in_(asset_ids))
            assets = query.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 Task ──────────────────────────────────────────────────────────────


@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

    logger.info("开始生成视频任务: task_id=%s", task_id)

    # 从数据库加载任务信息
    session = SessionLocal()
    try:
        from packages.adapters.sqlalchemy_impl.generation_task_repository import (
            SQLAlchemyGenerationTaskRepository,
        )

        task_repo = SQLAlchemyGenerationTaskRepository(session)
        gen_task = task_repo.get(task_id)
        if gen_task is None:
            logger.error("生成任务不存在: task_id=%s", task_id)
            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"
        task_asset_ids = list(gen_task.asset_ids or [])
        batch_id = getattr(gen_task, "batch_id", "") or ""
    finally:
        session.close()

    # 标记任务为 running
    _update_task_status(task_id, "mark_processing")

    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

            # 从素材库下载视频素材（如果任务指定了 asset_ids 则只下载这些）
            downloaded_videos = _download_library_assets(asset_library_id, temp_path, asset_ids=task_asset_ids or None)

            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)

            # 上传到 OSS
            bucket = _oss_bucket()
            if bucket:
                try:
                    bucket.put_object_from_file(storage_key, str(output_path))
                except Exception as oss_err:
                    logger.warning(f"OSS upload failed: {oss_err}")

            # 构建视频 URL
            if bucket:
                file_url = f"{PUBLIC_API_BASE_URL}/{storage_key}"
            else:
                file_url = f"{GENERATED_FILES_URL_PREFIX}/{task_id}/{output_name}"

            # 创建 GeneratedVideo 记录 + 查重
            video_count = _create_video_record_and_dedup(
                task_id=task_id,
                project_id=project_id,
                batch_id=batch_id,
                file_url=file_url,
                file_size=file_size,
                duration=duration,
                video_path=str(output_path),
                mode=editing_mode.value,
            )

            # 标记任务为 completed
            _update_task_status(task_id, "mark_completed", result_count=video_count or 1)

            logger.info("视频生成完成: task_id=%s duration=%.2fs file_size=%d", task_id, duration, file_size)

            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}", exc_info=True)
        # 标记任务为 failed
        _update_task_status(task_id, "mark_failed", error_message=str(error))
        return {
            "status": "failed",
            "task_id": task_id,
            "error": str(error),
        }


def _create_video_record_and_dedup(
    *,
    task_id: str,
    project_id: str,
    batch_id: str,
    file_url: str,
    file_size: int,
    duration: float,
    video_path: str,
    mode: str,
) -> int:
    """创建 GeneratedVideo 记录，计算指纹并执行查重（历史 + 批次）。

    Returns:
        创建的视频记录数量（1 表示成功，0 表示失败）
    """
    from uuid import uuid4

    from video_processing.dedup import VideoDeduplicator

    from packages.adapters.sqlalchemy_impl.generated_video_repository import (
        SQLAlchemyGeneratedVideoRepository,
    )
    from packages.domain import GeneratedVideo

    session = SessionLocal()
    try:
        video_id = uuid4().hex
        generated_video = GeneratedVideo(
            id=video_id,
            project_id=project_id,
            generation_task_id=task_id,
            name=f"generated-{task_id[:8]}.mp4",
            file_url=file_url,
            file_size=file_size,
            duration=duration,
            width=OUTPUT_WIDTH,
            height=OUTPUT_HEIGHT,
            fps=OUTPUT_FPS,
            status="completed",
            generation_params={"mode": mode},
        )

        video_repo = SQLAlchemyGeneratedVideoRepository(session)
        video_repo.create(generated_video)

        # 计算视频指纹
        deduplicator = VideoDeduplicator()
        try:
            fingerprint = deduplicator.compute_fingerprint(video_path)
        except Exception as fp_err:
            logger.warning(f"Fingerprint computation failed for {video_id}: {fp_err}")
            session.commit()
            return 1

        generated_video.video_fingerprint = fingerprint.to_dict()

        # (a) 历史成片查重
        duplicate_result = deduplicator.check_duplicate(fingerprint, project_id, session)

        # (b) 批次内查重（仅当有 batch_id 时）
        if not duplicate_result and batch_id:
            duplicate_result = deduplicator.check_batch_duplicate(fingerprint, batch_id, video_id, session)

        if duplicate_result:
            generated_video.is_duplicate = True
            generated_video.duplicate_of = duplicate_result["duplicate_of"]
            logger.info(
                f"Duplicate detected: {video_id} -> {duplicate_result['duplicate_of']} "
                f"(reason={duplicate_result['reason']}, similarity={duplicate_result['similarity']:.3f})"
            )
        else:
            generated_video.is_duplicate = False
            generated_video.duplicate_of = None

        video_repo.update(generated_video)
        session.commit()
        logger.info(f"GeneratedVideo record created: {video_id} (task={task_id}, dup={generated_video.is_duplicate})")
        return 1
    except Exception as e:
        logger.error(f"Failed to create video record / dedup for task {task_id}: {e}")
        session.rollback()
        return 0
    finally:
        session.close()