xiaoxia-saas/apps/worker/video_processing/dedup_helpers.py

"""查重辅助函数 — 渲染阶段指纹/查重预计算 + 兼容旧入库函数。

#2024: Worker 渲染+上传完成后**不直接创建 GeneratedVideo 成品记录**，改为：
1. ``compute_render_fingerprint_and_dedup``: 从本地视频计算指纹+查重（历史+批次），
   返回可序列化 dict（含 fingerprint_chunks），由 worker 写入
   ``GenerationTask.extra_meta["rendered_output"]``；
2. ``create_video_record_and_dedup``: 保留兼容——当传入 ``video_path`` 时会从本地视频
   计算指纹+查重并直接创建 GeneratedVideo 记录（供测试/旧路径使用）；
   当仅传 ``pre_dedup_result`` 时复用预计算结果，不再访问本地视频。

finalize 入口走 ``packages/application/generated_video_finalize.py`` 的
``finalize_generated_video``，不依赖本模块中数据库以外的 worker-only 逻辑。
"""

from __future__ import annotations

import logging
from uuid import uuid4

from sqlalchemy.orm import Session

logger = logging.getLogger(__name__)


def _safe_parse_fps(raw) -> float:
    if raw is None:
        return 25.0
    if isinstance(raw, (int, float)):
        return float(raw)
    s = str(raw).strip()
    if "/" in s:
        try:
            num, den = s.split("/", 1)
            return float(num) / float(den) if float(den) != 0 else 25.0
        except (ValueError, ZeroDivisionError):
            pass
    try:
        return float(s)
    except (ValueError, TypeError):
        return 25.0


def _compute_from_local(
    *,
    video_path: str,
    generation_task_id: str,
    project_id: str,
    user_id: str,
    batch_id: str,
    session: Session,
) -> dict:
    """从本地视频计算指纹+查重，返回可序列化结果 dict（不创建 DB 记录）。"""
    from video_processing.dedup import VideoDeduplicator
    from video_processing.ffmpeg_utils import probe_video_info

    result: dict = {
        "fingerprint_dict": None,
        "fingerprint_chunks": None,
        "duration": 0.0,
        "width": 1280,
        "height": 720,
        "fps": 25.0,
        "is_duplicate": False,
        "duplicate_of": None,
        "duplicate_rate": None,
        "match_count": None,
        "visual_similarity": None,
        "video_fingerprint_md5": "",
        "batch_similarity": None,
    }
    try:
        info = probe_video_info(video_path)
        result["duration"] = float(info.get("duration") or 0.0)
        result["width"] = int(info.get("width") or 1280)
        result["height"] = int(info.get("height") or 720)
        result["fps"] = _safe_parse_fps(info.get("fps"))
    except Exception as info_err:
        logger.warning("probe_video_info failed for task %s: %s", generation_task_id, info_err)

    try:
        deduplicator = VideoDeduplicator()
        fingerprint = deduplicator.compute_fingerprint(video_path)
        fp_dict = fingerprint.to_dict()
        result["fingerprint_dict"] = fp_dict
        result["video_fingerprint_md5"] = fingerprint.md5 or ""
        result["fingerprint_chunks"] = [
            {
                "start_time_ms": c.start_time_ms,
                "end_time_ms": c.end_time_ms,
                "phash_binary": c.phash_binary,
                "color_histogram": [float(v) for v in c.color_histogram],
                "frame_count": c.frame_count,
            }
            for c in fingerprint.chunks
        ]

        # 用 placeholder_id 占位（还没有真正的 video_id，不影响查重逻辑——
        # 因为查重排除的是 GeneratedVideo 表中的记录）
        placeholder_id = f"pre-{generation_task_id}"
        duration_sec = fingerprint.duration if fingerprint.duration else 0
        duplicate_result = deduplicator.check_duplicate(
            fingerprint,
            project_id,
            session,
            scope="user",
            user_id=user_id,
            duration_sec=duration_sec,
            exclude_video_id=placeholder_id,
        )
        batch_sim: float | None = None
        if not duplicate_result and batch_id:
            duplicate_result = deduplicator.check_batch_duplicate(fingerprint, batch_id, placeholder_id, session)
            if duplicate_result:
                batch_sim = float(duplicate_result.get("similarity", 0.0))
        result["batch_similarity"] = batch_sim
        if duplicate_result:
            result["is_duplicate"] = True
            result["duplicate_of"] = duplicate_result["duplicate_of"]
        else:
            result["is_duplicate"] = False

        try:
            rate_result = deduplicator.compute_duplicate_rate(
                fingerprint,
                project_id,
                placeholder_id,
                session,
                scope="user",
                user_id=user_id,
            )
            result["duplicate_rate"] = rate_result.get("duplicate_rate")
            result["match_count"] = rate_result.get("match_count")
            result["visual_similarity"] = rate_result.get("visual_similarity")
        except Exception as rate_err:
            logger.warning("compute_duplicate_rate failed for task %s: %s", generation_task_id, rate_err)
    except Exception as fp_err:
        logger.warning("Fingerprint compute failed for task %s: %s", generation_task_id, fp_err)

    return result


def compute_render_fingerprint_and_dedup(
    *,
    video_path: str,
    generation_task_id: str,
    project_id: str,
    user_id: str,
    batch_id: str,
    mode: str,
    session: Session,
) -> dict:
    """渲染+上传完成后的预计算：计算指纹+历史/批次查重，返回可序列化 dict。

    **不创建 GeneratedVideo 记录**。结果由调用方写入 extra_meta["rendered_output"]，
    finalize 时复用。mode 参数保留签名一致性（查重结果中不直接使用）。
    """
    _ = mode  # 保留在签名里便于调用方对齐；查重结果不含 mode
    return _compute_from_local(
        video_path=video_path,
        generation_task_id=generation_task_id,
        project_id=project_id,
        user_id=user_id,
        batch_id=batch_id,
        session=session,
    )


def create_video_record_and_dedup(
    *,
    generation_task_id: str,
    project_id: str,
    user_id: str = "",
    batch_id: str,
    file_url: str,
    file_size: int,
    duration: float | None = None,
    video_path: str | None,
    mode: str,
    session: Session,
    width: int = 1280,
    height: int = 720,
    fps: float = 25.0,
    name: str = "",
    thumbnail_url: str = "",
    pre_fingerprint_dict: dict | None = None,
    pre_fingerprint_chunks: list[dict] | None = None,
    pre_dedup_result: dict | None = None,
) -> dict:
    """创建 GeneratedVideo 记录 + 可选查重。

    两种用法：
    - 传入 ``video_path``（非 None）：从本地视频计算指纹+查重，直接创建记录（旧路径/测试）。
    - 仅传入 ``pre_*``：复用 worker 预计算结果，不访问本地视频（finalize 用）。

    Returns:
        {"video_id", "video_count", "is_duplicate", "batch_similarity", "duplicate_of"}
    """
    from packages.adapters.sqlalchemy_impl.generated_video_repository import (
        SQLAlchemyGeneratedVideoRepository,
    )
    from packages.adapters.sqlalchemy_impl.models import VideoFingerprintChunkModel
    from packages.domain.generated_video import GeneratedVideo

    try:
        video_id = uuid4().hex
        video_name = name.strip() if name else f"generated-{generation_task_id[:8]}.mp4"

        # 决定查重/元信息来源
        if video_path:
            pre = _compute_from_local(
                video_path=video_path,
                generation_task_id=generation_task_id,
                project_id=project_id,
                user_id=user_id,
                batch_id=batch_id,
                session=session,
            )
        else:
            pre = dict(pre_dedup_result or {})
            pre.setdefault("fingerprint_dict", pre_fingerprint_dict)
            pre.setdefault("fingerprint_chunks", pre_fingerprint_chunks)
            pre.setdefault("is_duplicate", False)
            pre.setdefault("duplicate_of", None)
            pre.setdefault("duplicate_rate", None)
            pre.setdefault("match_count", None)
            pre.setdefault("visual_similarity", None)
            pre.setdefault("batch_similarity", None)

        used_duration = float(duration if duration is not None else pre.get("duration", 0.0))
        used_width = int(pre.get("width", width) or width)
        used_height = int(pre.get("height", height) or height)
        used_fps = float(pre.get("fps", fps) or fps)

        generated_video = GeneratedVideo(
            id=video_id,
            project_id=project_id,
            user_id=user_id,
            generation_task_id=generation_task_id,
            name=video_name,
            file_url=file_url,
            file_size=file_size,
            duration=used_duration,
            width=used_width,
            height=used_height,
            fps=used_fps,
            status="completed",
            generation_params={"mode": mode},
            thumbnail_url=thumbnail_url or None,
            video_fingerprint=pre.get("fingerprint_dict"),
            is_duplicate=bool(pre.get("is_duplicate", False)),
            duplicate_of=pre.get("duplicate_of"),
            duplicate_rate=pre.get("duplicate_rate"),
            match_count=pre.get("match_count"),
            visual_similarity=pre.get("visual_similarity"),
        )

        # 写分片指纹表
        chunks = pre.get("fingerprint_chunks")
        if chunks:
            try:
                chunk_models = [
                    VideoFingerprintChunkModel(
                        id=uuid4().hex,
                        video_id=video_id,
                        project_id=project_id,
                        user_id=user_id,
                        start_time_ms=int(c.get("start_time_ms", 0)),
                        end_time_ms=int(c.get("end_time_ms", 0)),
                        phash_binary=str(c.get("phash_binary", "")),
                        color_histogram=[float(v) for v in (c.get("color_histogram") or [])],
                        frame_count=int(c.get("frame_count", 0)),
                    )
                    for c in chunks
                    if isinstance(c, dict)
                ]
                if chunk_models:
                    # 幂等：先清理旧分片
                    session.query(VideoFingerprintChunkModel).filter(
                        VideoFingerprintChunkModel.video_id == video_id
                    ).delete(synchronize_session=False)
                    session.bulk_save_objects(chunk_models)
            except Exception as chunk_err:
                logger.warning("Failed to save fingerprint chunks for %s: %s", video_id, chunk_err)

        repo = SQLAlchemyGeneratedVideoRepository(session)
        repo.create(generated_video)
        session.commit()
        return {
            "video_id": video_id,
            "video_count": 1,
            "is_duplicate": bool(pre.get("is_duplicate", False)),
            "batch_similarity": pre.get("batch_similarity"),
            "duplicate_of": pre.get("duplicate_of"),
        }
    except Exception as e:
        logger.error("Failed to create video record for task %s: %s", generation_task_id, e)
        session.rollback()
        return {
            "video_id": "",
            "video_count": 0,
            "is_duplicate": False,
            "batch_similarity": None,
            "duplicate_of": None,
        }