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

import subprocess
import tempfile
from datetime import datetime, timezone
from pathlib import Path

from celery.utils.log import get_task_logger
from video_processing.oss_helpers import download_asset
from worker_app.celery_app import celery_app
from worker_app.core.asset_types import infer_mime_type_from_storage_key
from worker_app.db import SessionLocal

from packages.adapters.sqlalchemy_impl import (
    SQLAlchemyAssetRepository,
    SQLAlchemyIngestJobRepository,
)
from packages.domain import Asset, AssetStatus, IngestJobStatus

logger = get_task_logger(__name__)


# 最小有效文件大小（字节）：小于此值的直接判为无效，避免文本/空文件伪装成媒体
MIN_VIDEO_FILE_SIZE = 1024  # 1KB
MIN_AUDIO_FILE_SIZE = 100  # 100B
MIN_IMAGE_FILE_SIZE = 100  # 100B

# 支持的视频编码格式（白名单，尽可能放宽）
# 渲染引擎会在 concat 前统一转码为 h264，因此只要 ffprobe 能识别的视频编码都允许 ingested
SUPPORTED_VIDEO_CODECS = {
    "h264",
    "avc1",
    "avc",  # H.264 / AVC
    "hevc",
    "h265",
    "hev1",
    "hvc1",  # H.265 / HEVC
    "vp9",
    "vp09",  # VP9
    "av1",
    "av01",  # AV1
    "vp8",
    "vp08",  # VP8
    "mpeg4",
    "mp4v",  # MPEG-4
    "mpeg2video",
    "mpg2",  # MPEG-2
    "wmv2",
    "wmv1",
    "vc1",  # WMV / VC-1
    "flv1",
    "flv",
    "vp6f",  # Flash / FLV
    "theora",
    "ogg",  # Theora
    "prores",
    "prores_ks",
    "apcn",
    "apch",
    "apco",
    "apcs",
    "ap4h",
    "ap4x",  # Apple ProRes
    "dnxhd",
    "dnxhr",  # DNxHD / DNxHR
}


def _safe_parse_fps(fps_str: str) -> float:
    """Safely parse fps from a fraction string like \"30/1\" or \"30000/1001\"."""
    try:
        if "/" in fps_str:
            num, den = fps_str.split("/", 1)
            den_val = float(den)
            if den_val == 0:
                return 0.0
            return float(num) / den_val
        return float(fps_str)
    except (ValueError, ZeroDivisionError):
        return 0.0


def extract_media_metadata(file_url: str, media_type: str) -> tuple[dict, bool]:
    """
    提取媒体文件的元数据。

    Args:
        file_url: 媒体文件 URL 或本地路径
        media_type: 媒体类型 (video, audio, image)

    Returns:
        (metadata_dict, success)
        - metadata_dict: 提取的元数据字典，失败时返回空字典
        - success: 是否成功提取到有效元数据
    """
    metadata = {}
    success = False

    try:
        if media_type == "video":
            # 使用 ffprobe 提取视频元数据
            from video_processing.ffmpeg_utils import run_ffprobe

            cmd = [
                "ffprobe",
                "-v",
                "quiet",
                "-print_format",
                "json",
                "-show_format",
                "-show_streams",
                file_url,
            ]
            try:
                stdout, _ = run_ffprobe(cmd, timeout=30)
                import json as json_lib

                probe_data = json_lib.loads(stdout)

                # 提取视频流信息
                for stream in probe_data.get("streams", []):
                    if stream.get("codec_type") == "video":
                        metadata["width"] = int(stream.get("width", 0))
                        metadata["height"] = int(stream.get("height", 0))
                        metadata["codec"] = stream.get("codec_name", "")
                        metadata["fps"] = (
                            _safe_parse_fps(stream.get("r_frame_rate", "0/1")) if stream.get("r_frame_rate") else 0  # type: ignore[assignment]
                        )
                        break

                # 提取格式信息
                format_info = probe_data.get("format", {})
                metadata["duration"] = float(format_info.get("duration", 0))  # type: ignore[assignment]
                metadata["size_bytes"] = int(format_info.get("size", 0))
                metadata["bitrate"] = int(format_info.get("bit_rate", 0))

                # 有效性判断：有视频流 + duration > 0 + size > 0
                has_video_stream = any(s.get("codec_type") == "video" for s in probe_data.get("streams", []))
                has_audio_stream = any(s.get("codec_type") == "audio" for s in probe_data.get("streams", []))
                if (has_video_stream or has_audio_stream) and metadata.get("duration", 0) > 0:
                    success = True

            except Exception as e:
                logger.warning("视频元数据提取失败: %s", e)

        elif media_type == "audio":
            from video_processing.ffmpeg_utils import run_ffprobe

            cmd = [
                "ffprobe",
                "-v",
                "quiet",
                "-print_format",
                "json",
                "-show_format",
                "-show_streams",
                file_url,
            ]
            try:
                stdout, _ = run_ffprobe(cmd, timeout=30)
                import json as json_lib

                probe_data = json_lib.loads(stdout)

                has_audio_stream = any(s.get("codec_type") == "audio" for s in probe_data.get("streams", []))
                format_info = probe_data.get("format", {})
                metadata["duration"] = float(format_info.get("duration", 0))  # type: ignore[assignment]
                metadata["size_bytes"] = int(format_info.get("size", 0))
                metadata["bitrate"] = int(format_info.get("bit_rate", 0))
                metadata["codec"] = next(
                    (s.get("codec_name", "") for s in probe_data.get("streams", []) if s.get("codec_type") == "audio"),
                    "",
                )

                if has_audio_stream and metadata.get("duration", 0) > 0:
                    success = True

            except Exception as e:
                logger.warning("音频元数据提取失败: %s", e)

        elif media_type == "image":
            # 使用 Pillow 提取图片元数据
            try:
                from PIL import Image

                with Image.open(file_url) as img:
                    img.verify()  # 验证文件完整性
                    # verify后需要重新打开才能读尺寸
                    with Image.open(file_url) as img2:
                        metadata["width"] = img2.width
                        metadata["height"] = img2.height
                        metadata["format"] = img2.format
                        metadata["mode"] = img2.mode
                        if img2.width > 0 and img2.height > 0:
                            success = True

            except ImportError:
                logger.warning("Pillow not available for image metadata extraction")
            except Exception as e:
                logger.warning(f"Failed to extract image metadata: {e}")

    except subprocess.TimeoutExpired:
        logger.warning(f"Timeout extracting metadata from {file_url}")
    except FileNotFoundError:
        logger.warning("ffprobe not found, cannot extract video metadata")
    except Exception as e:
        logger.warning(f"Failed to extract metadata: {e}")

    return metadata, success


def _is_valid_media(metadata: dict, media_type: str) -> bool:
    """根据元数据判断文件是否为有效媒体文件。

    Args:
        metadata: extract_media_metadata 返回的元数据
        media_type: 媒体类型

    Returns:
        True 表示文件有效
    """
    size = int(metadata.get("size_bytes", 0))

    if media_type == "video":
        duration = float(metadata.get("duration", 0))
        if size < MIN_VIDEO_FILE_SIZE or duration <= 0:
            return False
        # 编码格式校验：只排除明确非视频的编码格式，只要 ffprobe 能识别的视频编码都允许
        # 渲染引擎会在 concat 前统一转码为 h264 yuv420p，ingest 层不再做严格的编码拦截
        codec = str(metadata.get("codec", "")).lower()
        if codec and codec not in SUPPORTED_VIDEO_CODECS:
            logger.info("检测到非白名单视频编码 %s，仍允许 ingested，渲染层会统一转码", codec)
        return True
    if media_type == "audio":
        duration = float(metadata.get("duration", 0))
        return size >= MIN_AUDIO_FILE_SIZE and duration > 0
    if media_type == "image":
        width = int(metadata.get("width", 0))
        height = int(metadata.get("height", 0))
        return size >= MIN_IMAGE_FILE_SIZE and width > 0 and height > 0
    return False


@celery_app.task(name="worker.ingest_asset")
def ingest_asset(job_id: str) -> dict:
    """
    Ingest asset task.

    Steps:
    1. Fetch IngestJob from repository
    2. Extract metadata from storage_key
    3. Create Asset entity
    4. Update IngestJob status to COMPLETED
    5. Return result
    """
    db = SessionLocal()
    try:
        job_repo = SQLAlchemyIngestJobRepository(db)
        asset_repo = SQLAlchemyAssetRepository(db)

        job = job_repo.get(job_id)
        if job is None:
            return {"status": "failed", "error": "job not found"}

        # Update job status to PROCESSING
        job.status = IngestJobStatus.PROCESSING
        job.updated_at = datetime.now(timezone.utc)
        job_repo.update(job)
        db.commit()

        # Extract real metadata from media file
        filename = job.storage_key.split("/")[-1]
        mime_type = infer_mime_type_from_storage_key(job.storage_key)

        # Determine media type from mime_type
        media_type = "video"
        if mime_type.startswith("image/"):
            media_type = "image"
        elif mime_type.startswith("audio/"):
            media_type = "audio"

        # 先从 OSS 下载文件到本地临时目录，再提取元数据
        # （storage_key 是 OSS 内部路径，不能直接传给 ffprobe/Pillow）
        local_file = None
        thumbnail_url = None
        try:
            suffix = Path(job.storage_key).suffix or ".bin"
            with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
                local_file = Path(tmp.name)

            download_ok = download_asset(job.storage_key, local_file)
            if not download_ok:
                logger.warning("素材下载失败，无法提取元数据: job_id=%s storage_key=%s", job_id, job.storage_key)
                metadata, extract_success = {}, False
            else:
                metadata, extract_success = extract_media_metadata(str(local_file), media_type)

            # 视频类型：生成缩略图（文件还在的时候生成）
            thumbnail_url = None
            if media_type == "video" and extract_success:
                try:
                    from video_processing.thumbnail_generator import generate_and_upload_thumbnail

                    thumb_storage_key = f"assets/{job.project_id}/thumbnails/{job_id}.jpg"
                    thumbnail_url = generate_and_upload_thumbnail(str(local_file), thumb_storage_key)
                    if thumbnail_url:
                        logger.info(
                            "素材缩略图生成成功: job_id=%s url=%s",
                            job_id,
                            thumbnail_url[:80],
                        )
                except Exception as thumb_err:
                    logger.warning(
                        "素材缩略图生成失败（不影响主流程）: job_id=%s error=%s",
                        job_id,
                        thumb_err,
                    )
        finally:
            if local_file and local_file.exists():
                try:
                    local_file.unlink()
                except OSError:
                    pass

        # 有效性校验：ffprobe/Pillow 必须成功，且文件大小/时长/尺寸满足最小要求
        is_valid = extract_success and _is_valid_media(metadata, media_type)

        if not is_valid:
            # 文件无效，创建 ERROR 状态的 asset 并标记 job 失败
            error_reason = "metadata extraction failed" if not extract_success else "media validation failed"
            logger.warning(
                "素材有效性校验失败，标记为ERROR: job_id=%s storage_key=%s media_type=%s reason=%s",
                job_id,
                job.storage_key,
                media_type,
                error_reason,
            )

            asset = Asset.create(
                project_id=job.project_id,
                library_id=job.library_id,
                name=filename,
                storage_key=job.storage_key,
                mime_type=mime_type,
                metadata={"ingest_error": error_reason},
                file_size=int(metadata.get("size_bytes", 0)),
                duration=float(metadata.get("duration", 0)),
                width=int(metadata.get("width", 0)),
                height=int(metadata.get("height", 0)),
                codec=metadata.get("codec") or None,
                status=AssetStatus.ERROR,
                file_hash=job.file_hash,
            )
            asset_repo.create(asset)

            # Update job status to FAILED
            job.status = IngestJobStatus.FAILED
            job.error_message = f"Invalid media file: {error_reason}"
            job.result_asset_id = asset.id
            job.updated_at = datetime.now(timezone.utc)
            job_repo.update(job)
            db.commit()

            return {
                "status": "failed",
                "job_id": job.id,
                "asset_id": asset.id,
                "error": error_reason,
            }

        # Create Asset
        asset = Asset.create(
            project_id=job.project_id,
            library_id=job.library_id,
            name=filename,
            storage_key=job.storage_key,
            mime_type=mime_type,
            metadata=metadata,
            file_size=int(metadata.get("size_bytes", 0)),
            duration=float(metadata.get("duration", 0)),
            width=int(metadata.get("width", 0)),
            height=int(metadata.get("height", 0)),
            codec=metadata.get("codec") or None,
            status=AssetStatus.READY,
            file_hash=job.file_hash,
            thumbnail_url=thumbnail_url,
        )
        asset_repo.create(asset)

        # Update job status to COMPLETED
        job.status = IngestJobStatus.COMPLETED
        job.result_asset_id = asset.id
        job.updated_at = datetime.now(timezone.utc)
        job_repo.update(job)

        db.commit()

        return {
            "status": "completed",
            "job_id": job.id,
            "asset_id": asset.id,
        }
    except Exception as e:
        db.rollback()
        logger.error(f"Failed to ingest asset {job_id}: {e}")

        # Update job status to FAILED
        try:
            job_repo = SQLAlchemyIngestJobRepository(db)
            job = job_repo.get(job_id)
            if job:
                job.status = IngestJobStatus.FAILED
                job.error_message = str(e)
                job.updated_at = datetime.now(timezone.utc)
                job_repo.update(job)
                db.commit()
        except Exception:
            db.rollback()

        return {
            "status": "failed",
            "job_id": job_id,
            "error": str(e),
        }
    finally:
        db.close()