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

"""Video deduplication module - compute fingerprints and detect duplicates."""

import hashlib
import logging
import os
import tempfile
from dataclasses import dataclass, field
from typing import Optional
from uuid import uuid4

import cv2
import numpy as np
from celery import Task
from sqlalchemy.orm import Session
from worker_app.celery_app import celery_app
from worker_app.db import SessionLocal

from packages.adapters.sqlalchemy_impl.generated_video_repository import SQLAlchemyGeneratedVideoRepository
from packages.adapters.sqlalchemy_impl.models import VideoFingerprintChunkModel
from packages.shared.storage import get_storage_service

logger = logging.getLogger(__name__)

# 分片策略常量
SHORT_VIDEO_CHUNK_SEC = 2  # ≤60秒视频，每 2 秒一个分片
LONG_VIDEO_CHUNK_SEC = 5  # >60秒视频，每 5 秒一个分片
SHORT_VIDEO_THRESHOLD_SEC = 60


def compute_phash(image: np.ndarray, hash_size: int = 8) -> str:
    """计算图像的感知哈希（pHash），基于 DCT（离散余弦变换）。

    算法步骤：
    1. 将图像缩放到 hash_size*4 × hash_size*4（默认 32×32）
    2. 转为灰度图，应用 2D DCT 提取频率分量
    3. 取左上角 hash_size×hash_size 的低频分量（默认 8×8 = 64 bit）
    4. 排除 DC 分量（[0,0] 位置），计算中位数
    5. 每个分量与中位数比较，生成二值 hash

    Args:
        image: BGR 格式的 numpy 图像数组
        hash_size: 哈希边长，默认 8（生成 64-bit hash）

    Returns:
        十六进制字符串表示的感知哈希
    """
    # Resize to 32x32 for DCT
    resized = cv2.resize(image, (hash_size * 4, hash_size * 4))
    gray = cv2.cvtColor(resized, cv2.COLOR_BGR2GRAY).astype(np.float32)
    # Apply 2D DCT
    dct = cv2.dct(gray)
    # Take top-left 8x8 low-frequency components
    dct_low = dct[:hash_size, :hash_size]
    # Compute median (excluding DC component at [0,0])
    dct_low[0, 0] = 0
    median = np.median(dct_low)
    # Generate hash based on comparison with median
    diff = (dct_low > median).astype(int)
    hash_str = "".join(str(b) for row in diff for b in row)
    return hex(int(hash_str, 2))[2:]


def hamming_distance(hash1: str, hash2: str) -> int:
    """计算两个十六进制哈希之间的汉明距离（不同 bit 位数）。

    使用 XOR 异或 + bit 计数：bin(h1 ^ h2).count("1")。
    例如：hamming_distance("00", "ff") = 8（8 个 bit 全不同）。

    Args:
        hash1: 十六进制字符串
        hash2: 十六进制字符串

    Returns:
        不同 bit 的数量
    """
    h1, h2 = int(hash1, 16), int(hash2, 16)
    return bin(h1 ^ h2).count("1")


def compute_color_histogram(image: np.ndarray, bins: int = 32) -> list[float]:
    """Compute color histogram for an image."""
    hist = []
    for i in range(3):
        h = cv2.calcHist([image], [i], None, [bins], [0, 256])
        h = cv2.normalize(h, h).flatten()
        hist.extend(h)
    return hist


def compute_chunk_interval(duration: float) -> float:
    """根据视频时长返回分片间隔（秒）。

    短视频（≤60秒）：每 2 秒一个分片
    长视频（>60秒）：每 5 秒一个分片
    """
    if duration <= SHORT_VIDEO_THRESHOLD_SEC:
        return SHORT_VIDEO_CHUNK_SEC
    return LONG_VIDEO_CHUNK_SEC


@dataclass
class FingerprintChunk:
    """单个分片指纹数据。"""

    start_time_ms: int
    end_time_ms: int
    phash_binary: str
    color_histogram: list[float]
    frame_count: int = 1


@dataclass
class VideoFingerprint:
    """Video fingerprint containing multiple similarity metrics."""

    md5: str
    keyframe_phashes: list[str]
    color_histograms: list[list[float]]
    duration: float
    resolution: tuple[int, int]
    chunks: list[FingerprintChunk] = field(default_factory=list)

    def to_dict(self) -> dict:
        # 注意：color_histograms 里的值可能是 np.float32（来自 cv2.normalize），
        # 直接存进 dict 后 SQLAlchemy JSON 序列化会报 "float32 is not JSON serializable"。
        # 这里统一转成 Python 原生 float。
        native_histograms = [[float(v) for v in hist] for hist in self.color_histograms]
        return {
            "md5": self.md5,
            "keyframe_phashes": self.keyframe_phashes,
            "color_histograms": native_histograms,
            "duration": float(self.duration),
            "resolution": [int(self.resolution[0]), int(self.resolution[1])],
            "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 self.chunks
            ],
        }

    def to_chunk_models(self, video_id: str, project_id: str, user_id: str = "") -> list[VideoFingerprintChunkModel]:
        """将分片数据转为 SQLAlchemy Model 列表，用于批量写入 video_fingerprint_chunks 表。"""
        models = []
        for chunk in self.chunks:
            models.append(
                VideoFingerprintChunkModel(
                    id=uuid4().hex,
                    video_id=video_id,
                    project_id=project_id,
                    user_id=user_id,
                    start_time_ms=chunk.start_time_ms,
                    end_time_ms=chunk.end_time_ms,
                    phash_binary=chunk.phash_binary,
                    color_histogram=[float(v) for v in chunk.color_histogram],
                    frame_count=chunk.frame_count,
                )
            )
        return models


class VideoDeduplicator:
    """Video deduplication using multiple fingerprint methods."""

    PHASH_THRESHOLD = 10
    HISTOGRAM_THRESHOLD = 0.85

    def compute_fingerprint(self, video_path: str) -> VideoFingerprint:
        """Compute video fingerprint using MD5, pHash, and color histogram.

        按时间分片抽帧：短视频（≤60s）每 2s 一片，长视频每 5s 一片。
        每片取 1 帧计算 pHash + color_histogram。
        同时保留 keyframe_phashes/color_histograms 聚合字段（向后兼容）。
        """
        cap = cv2.VideoCapture(video_path)
        if not cap.isOpened():
            raise RuntimeError(f"Cannot open video: {video_path}")

        fps = cap.get(cv2.CAP_PROP_FPS)
        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        duration = frame_count / fps if fps > 0 else 0
        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

        md5_hash = hashlib.md5(usedforsecurity=False)
        chunks: list[FingerprintChunk] = []

        # 分片间隔（秒）
        chunk_interval_sec = compute_chunk_interval(duration)
        chunk_interval_ms = int(chunk_interval_sec * 1000)
        duration_ms = int(duration * 1000)

        # 遍历每个分片时间窗口，取 1 帧
        start_ms = 0
        while start_ms < duration_ms:
            end_ms = min(start_ms + chunk_interval_ms, duration_ms)
            # 定位到分片中点
            seek_ms = (start_ms + end_ms) / 2
            cap.set(cv2.CAP_PROP_POS_MSEC, seek_ms)
            ret, frame = cap.read()
            if ret:
                # MD5 计算
                _, buffer = cv2.imencode(".jpg", frame)
                md5_hash.update(buffer)

                phash = compute_phash(frame)
                hist = compute_color_histogram(frame)

                chunks.append(
                    FingerprintChunk(
                        start_time_ms=start_ms,
                        end_time_ms=end_ms,
                        phash_binary=phash,
                        color_histogram=hist,
                        frame_count=1,
                    )
                )

            start_ms = end_ms

        cap.release()

        # 向后兼容：聚合 keyframe_phashes / color_histograms
        keyframe_phashes = [c.phash_binary for c in chunks]
        color_histograms = [c.color_histogram for c in chunks]

        return VideoFingerprint(
            md5=md5_hash.hexdigest(),
            keyframe_phashes=keyframe_phashes,
            color_histograms=color_histograms,
            duration=duration,
            resolution=(width, height),
            chunks=chunks,
        )

    def _get_existing_chunks(self, video_id: str, session: Session) -> list[dict]:
        """从 video_fingerprint_chunks 表读取分片数据。返回空列表表示无分片数据。"""
        rows = (
            session.query(VideoFingerprintChunkModel)
            .filter(VideoFingerprintChunkModel.video_id == video_id)
            .order_by(VideoFingerprintChunkModel.start_time_ms)
            .all()
        )
        return [
            {
                "phash_binary": r.phash_binary,
                "color_histogram": r.color_histogram,
                "start_time_ms": r.start_time_ms,
                "end_time_ms": r.end_time_ms,
            }
            for r in rows
        ]

    def check_duplicate(self, fingerprint: VideoFingerprint, project_id: str, session: Session) -> Optional[dict]:
        """检查视频是否与项目中已有视频重复。

        查重逻辑：
        1. MD5 精确匹配 → similarity=1.0
        2. pHash 相似度（优先从分片表读取，回退到 JSON 字段）

        判定阈值：avg_distance < PHASH_THRESHOLD(10)

        Args:
            fingerprint: 待检测视频的指纹
            project_id: 项目 ID，仅在同一项目内搜索
            session: 数据库会话

        Returns:
            重复信息字典（含 duplicate, duplicate_of, reason, similarity），
            或 None 表示未找到重复。
        """
        video_repo = SQLAlchemyGeneratedVideoRepository(session)
        existing_videos = video_repo.list_by_project(project_id)

        for existing in existing_videos:
            if not existing.video_fingerprint:
                continue

            ef = existing.video_fingerprint

            # 精确匹配：MD5 完全一致
            if fingerprint.md5 == ef.get("md5"):
                return {"duplicate": True, "duplicate_of": existing.id, "reason": "exact_md5_match", "similarity": 1.0}

            # 优先从分片表读取已有视频的分片 phash
            existing_phashes = []
            chunk_data = self._get_existing_chunks(existing.id, session)
            if chunk_data:
                existing_phashes = [c["phash_binary"] for c in chunk_data]
            else:
                # 回退：从 JSON 字段读取（存量旧视频）
                existing_phashes = ef.get("keyframe_phashes", [])

            if not existing_phashes:
                continue

            # 计算每个新关键帧到已有关键帧的最小汉明距离，取平均
            min_distances = []
            for phash in fingerprint.keyframe_phashes:
                distances = [hamming_distance(phash, ep) for ep in existing_phashes]
                min_distances.append(min(distances))
            avg_distance = sum(min_distances) / len(min_distances) if min_distances else 100

            if avg_distance >= self.PHASH_THRESHOLD:
                continue

            phash_similarity = 1.0 - (avg_distance / 64)

            return {
                "duplicate": True,
                "duplicate_of": existing.id,
                "reason": "phash_similar",
                "similarity": phash_similarity,
            }

        return None

    def check_batch_duplicate(
        self,
        fingerprint: VideoFingerprint,
        batch_id: str,
        current_video_id: str,
        session: Session,
    ) -> Optional[dict]:
        """检查视频是否与同批次内其他视频重复。

        逻辑与 check_duplicate 一致（MD5 + pHash），但搜索范围限定为同 batch_id 的视频。

        Args:
            fingerprint: 待检测视频的指纹
            batch_id: 批次 ID
            current_video_id: 当前视频 ID（排除自身）
            session: 数据库会话

        Returns:
            重复信息字典，或 None 表示未找到重复
        """
        video_repo = SQLAlchemyGeneratedVideoRepository(session)
        batch_videos = video_repo.list_by_batch(batch_id)

        for existing in batch_videos:
            if existing.id == current_video_id:
                continue
            if not existing.video_fingerprint:
                continue

            ef = existing.video_fingerprint

            if fingerprint.md5 == ef.get("md5"):
                return {
                    "duplicate": True,
                    "duplicate_of": existing.id,
                    "reason": "batch_exact_md5_match",
                    "similarity": 1.0,
                }

            # 优先从分片表读取
            existing_phashes = []
            chunk_data = self._get_existing_chunks(existing.id, session)
            if chunk_data:
                existing_phashes = [c["phash_binary"] for c in chunk_data]
            else:
                existing_phashes = ef.get("keyframe_phashes", [])

            if not existing_phashes:
                continue

            min_distances = []
            for phash in fingerprint.keyframe_phashes:
                distances = [hamming_distance(phash, ep) for ep in existing_phashes]
                min_distances.append(min(distances))
            avg_distance = sum(min_distances) / len(min_distances) if min_distances else 100

            if avg_distance >= self.PHASH_THRESHOLD:
                continue

            phash_similarity = 1.0 - (avg_distance / 64)
            return {
                "duplicate": True,
                "duplicate_of": existing.id,
                "reason": "batch_phash_similar",
                "similarity": phash_similarity,
            }

        return None

    @staticmethod
    def _average_histogram_similarity(histograms_a: list[list[float]], histograms_b: list[list[float]]) -> float:
        """
        计算两组颜色直方图之间的平均余弦相似度。

        对每组直方图对取最小长度对齐，计算余弦相似度后取平均。

        Args:
            histograms_a: 第一组直方图（每帧一个 list）
            histograms_b: 第二组直方图

        Returns:
            平均余弦相似度，范围 [0, 1]
        """
        if not histograms_a or not histograms_b:
            return 0.0

        similarities = []
        for ha in histograms_a:
            best = 0.0
            vec_a = np.array(ha, dtype=np.float64)
            norm_a = np.linalg.norm(vec_a)
            if norm_a == 0:
                continue
            for hb in histograms_b:
                vec_b = np.array(hb, dtype=np.float64)
                # 对齐长度
                min_len = min(len(vec_a), len(vec_b))
                va, vb = vec_a[:min_len], vec_b[:min_len]
                norm_b = np.linalg.norm(vb)
                if norm_b == 0:
                    continue
                sim = float(np.dot(va, vb) / (norm_a * norm_b))
                best = max(best, sim)
            similarities.append(best)

        return sum(similarities) / len(similarities) if similarities else 0.0

    def compute_duplicate_rate(
        self,
        fingerprint: VideoFingerprint,
        project_id: str,
        current_video_id: str | None,
        session: Session,
        *,
        user_id: str = "",
    ) -> float:
        """计算当前视频与用户库内已有视频的最高相似度百分比。

        优先按 user_id 全局比较（跨项目），user_id 为空时回退到项目级比较。
        遍历最近 200 个其他有指纹的视频，对每个计算相似度：
        - MD5 精确匹配 → 100%
        - pHash 相似度 → (1.0 - avg_distance / 64) * 100
        取最高值作为 duplicate_rate（0~100）。
        如果没有其他视频可比较，返回 0.0。

        Args:
            fingerprint: 当前视频的指纹
            project_id: 项目 ID（user_id 为空时的回退范围）
            current_video_id: 当前视频 ID（排除自身，可为 None）
            session: 数据库会话
            user_id: 用户 ID（优先按用户全局比较）

        Returns:
            duplicate_rate: 0~100 的浮点数
        """
        # 限制查询最近 200 个视频，避免大库内存溢出
        from packages.adapters.sqlalchemy_impl.models import GeneratedVideoModel

        # 优先按 user_id 全局比较（跨项目），否则回退到项目级
        if user_id:
            query = session.query(GeneratedVideoModel).filter(
                GeneratedVideoModel.user_id == user_id,
            )
            logger.debug("compute_duplicate_rate: user-level scope user_id=%s", user_id)
        else:
            query = session.query(GeneratedVideoModel).filter(
                GeneratedVideoModel.project_id == project_id,
            )
            logger.debug("compute_duplicate_rate: project-level fallback project_id=%s", project_id)

        # 排除当前视频自身（记录可能已写入 DB，必须在查询层排除）
        if current_video_id:
            query = query.filter(GeneratedVideoModel.id != current_video_id)

        recent_models = query.order_by(GeneratedVideoModel.generated_at.desc()).limit(200).all()
        video_repo = SQLAlchemyGeneratedVideoRepository(session)
        existing_videos = [video_repo._to_domain(m) for m in recent_models]

        max_similarity = 0.0
        for existing in existing_videos:
            if current_video_id and existing.id == current_video_id:
                continue
            if not existing.video_fingerprint:
                continue

            ef = existing.video_fingerprint

            # MD5 精确匹配 → 100%
            if fingerprint.md5 == ef.get("md5"):
                return 100.0

            # 优先从分片表读取
            existing_phashes = []
            chunk_data = self._get_existing_chunks(existing.id, session)
            if chunk_data:
                existing_phashes = [c["phash_binary"] for c in chunk_data]
            else:
                existing_phashes = ef.get("keyframe_phashes", [])

            if not existing_phashes or not fingerprint.keyframe_phashes:
                continue

            min_distances = []
            for phash in fingerprint.keyframe_phashes:
                distances = [hamming_distance(phash, ep) for ep in existing_phashes]
                min_distances.append(min(distances))
            avg_distance = sum(min_distances) / len(min_distances) if min_distances else 64
            similarity = (1.0 - avg_distance / 64) * 100
            max_similarity = max(max_similarity, similarity)

        return round(max(max_similarity, 0.0), 2)


def _save_fingerprint_chunks(
    fingerprint: VideoFingerprint,
    video_id: str,
    project_id: str,
    user_id: str,
    session: Session,
) -> None:
    """将指纹分片数据批量写入 video_fingerprint_chunks 表。幂等：已有数据时跳过。"""
    # 幂等检查：已有分片数据则跳过
    existing_count = (
        session.query(VideoFingerprintChunkModel).filter(VideoFingerprintChunkModel.video_id == video_id).count()
    )
    if existing_count > 0:
        logger.debug("Fingerprint chunks already exist for video %s (%d chunks), skipping", video_id, existing_count)
        return

    if not fingerprint.chunks:
        logger.warning("No chunks in fingerprint for video %s, skipping chunk save", video_id)
        return

    chunk_models = fingerprint.to_chunk_models(video_id, project_id, user_id)
    session.bulk_save_objects(chunk_models)
    logger.info("Saved %d fingerprint chunks for video %s", len(chunk_models), video_id)


@celery_app.task(bind=True, max_retries=3, name="worker.check_duplicate")
def check_duplicate_task(self: Task, generated_video_id: str) -> dict:
    """Celery task to check if generated video is a duplicate."""
    session = SessionLocal()
    temp_dir = tempfile.mkdtemp()

    try:
        video_repo = SQLAlchemyGeneratedVideoRepository(session)
        storage_service = get_storage_service()
        deduplicator = VideoDeduplicator()

        video = video_repo.get(generated_video_id)
        if video is None:
            raise ValueError(f"Generated video {generated_video_id} not found")

        local_path = os.path.join(temp_dir, f"{generated_video_id}.mp4")
        storage_service.download_file(
            f"projects/{video.project_id}/generated/{generated_video_id}/{generated_video_id}.mp4", local_path
        )

        fingerprint = deduplicator.compute_fingerprint(local_path)

        duplicate_result = deduplicator.check_duplicate(fingerprint, video.project_id, session)

        video.video_fingerprint = fingerprint.to_dict()
        if duplicate_result:
            video.is_duplicate = True
            video.duplicate_of = duplicate_result["duplicate_of"]
        else:
            video.is_duplicate = False
            video.duplicate_of = None

        video_repo.update(video)

        # 写入分片表
        _save_fingerprint_chunks(fingerprint, generated_video_id, video.project_id, video.user_id, session)

        session.commit()

        logger.info(f"Duplicate check completed for video {generated_video_id}: is_duplicate={video.is_duplicate}")

        return {
            "ok": True,
            "video_id": generated_video_id,
            "is_duplicate": video.is_duplicate,
            "duplicate_of": video.duplicate_of,
            "fingerprint": fingerprint.to_dict(),
        }
    except Exception as e:
        logger.error(f"Duplicate check failed for {generated_video_id}: {str(e)}")
        session.rollback()
        raise self.retry(exc=e, countdown=60) from e
    finally:
        session.close()
        import shutil

        shutil.rmtree(temp_dir, ignore_errors=True)