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

"""
视频素材分析器 - 基于 FFmpeg + NumPy 的轻量级智能分析

提供：
1. 素材分类 - 基于视频特征的多维度分析
2. 质量评分 - 基于分辨率、帧率、码率、清晰度、稳定性的综合评分
"""

from __future__ import annotations

import json
import logging
import os
import tempfile
from dataclasses import dataclass, field

import numpy as np
from PIL import Image

from packages.domain.classification import AssetClassification

logger = logging.getLogger(__name__)


@dataclass
class VideoInfo:
    """视频基本信息"""

    width: int = 0
    height: int = 0
    fps: float = 0.0
    duration: float = 0.0
    bitrate: int = 0
    codec: str = ""
    has_audio: bool = False
    file_size: int = 0


@dataclass
class ColorAnalysis:
    """色彩分析结果"""

    dominant_hue: float = 0.0  # 主色调 (0-360)
    green_ratio: float = 0.0  # 绿色占比
    warm_ratio: float = 0.0  # 暖色调占比
    cool_ratio: float = 0.0  # 冷色调占比
    avg_saturation: float = 0.0
    avg_brightness: float = 0.0


@dataclass
class MotionAnalysis:
    """运动分析结果"""

    motion_score: float = 0.0  # 运动幅度 (0-1)
    scene_changes: int = 0  # 场景切换次数


@dataclass
class AudioAnalysis:
    """音频分析结果"""

    has_audio: bool = False
    speech_ratio: float = 0.0  # 人声比例
    music_ratio: float = 0.0  # 音乐比例
    ambient_ratio: float = 0.0  # 环境音比例


@dataclass
class ClassificationResult:
    """分类结果"""

    category: AssetClassification
    confidence: float
    scores: dict[str, float] = field(default_factory=dict)


@dataclass
class QualityScore:
    """质量评分结果"""

    total: float
    resolution_score: float = 0.0
    fps_score: float = 0.0
    bitrate_score: float = 0.0
    clarity_score: float = 0.0
    stability_score: float = 0.0


class AssetAnalyzer:
    """
    轻量级视频素材分析器

    使用 FFmpeg + NumPy 进行视频特征分析，不依赖外部 AI API。
    """

    def __init__(self, video_path: str, temp_dir: str | None = None):
        """
        初始化分析器

        Args:
            video_path: 视频文件路径
            temp_dir: 临时目录，用于存储提取的帧
        """
        self.video_path = video_path
        self._video_info: VideoInfo | None = None
        self._frames: list[np.ndarray] | None = None
        self._temp_dir = temp_dir or tempfile.mkdtemp(prefix="asset_analyzer_")

    def __del__(self):
        """清理临时文件"""
        self._cleanup_temp_dir()

    def _cleanup_temp_dir(self):
        """清理临时目录"""
        try:
            import shutil

            if os.path.exists(self._temp_dir):
                shutil.rmtree(self._temp_dir)
        except Exception as e:
            logger.warning(f"Operation failed in apps/worker/worker_app/tasks/asset_analyzer.py: {e}", exc_info=True)

    def get_video_info(self) -> VideoInfo:
        """获取视频基本信息"""
        if self._video_info is not None:
            return self._video_info

        info = VideoInfo()

        try:
            from video_processing.ffmpeg_utils import run_ffprobe

            cmd = [
                "ffprobe",
                "-v",
                "quiet",
                "-print_format",
                "json",
                "-show_format",
                "-show_streams",
                self.video_path,
            ]
            stdout, _ = run_ffprobe(cmd, timeout=30)
            data = json.loads(stdout)
            streams = data.get("streams", [])
            format_info = data.get("format", {})

            for stream in streams:
                if stream.get("codec_type") == "video":
                    info.width = int(stream.get("width", 0))
                    info.height = int(stream.get("height", 0))
                    info.codec = stream.get("codec_name", "")

                    # 解析帧率
                    fps_str = stream.get("r_frame_rate", "0/1")
                    if "/" in fps_str:
                        num, denom = fps_str.split("/")
                        info.fps = float(num) / float(denom) if float(denom) != 0 else 0.0
                    else:
                        info.fps = float(fps_str)

                elif stream.get("codec_type") == "audio":
                    info.has_audio = True

            info.duration = float(format_info.get("duration", 0))
            info.bitrate = int(format_info.get("bit_rate", 0))
            info.file_size = int(format_info.get("size", 0))

        except Exception as e:
            logger.warning(f"Failed to get video info: {e}")

        self._video_info = info
        return info

    def extract_frames(self, count: int = 10) -> list[np.ndarray]:
        """
        从视频中均匀抽取帧

        Args:
            count: 抽取的帧数

        Returns:
            帧数据列表 (RGB 格式)
        """
        if self._frames is not None:
            return self._frames

        frames: list[np.ndarray] = []
        info = self.get_video_info()

        if info.duration <= 0:
            logger.warning("Video duration is 0, cannot extract frames")
            return frames

        try:
            # 计算采样间隔
            interval = max(1.0, info.duration / count)

            for i in range(count):
                timestamp = i * interval

                # 提取单帧为 PNG
                output_path = os.path.join(self._temp_dir, f"frame_{i:03d}.png")
                cmd = [
                    "ffmpeg",
                    "-y",  # 覆盖输出文件
                    "-ss",
                    str(timestamp),
                    "-i",
                    self.video_path,
                    "-vframes",
                    "1",
                    "-q:v",
                    "2",  # 高质量
                    "-f",
                    "image2",
                    output_path,
                ]

                from video_processing.ffmpeg_utils import run_ffmpeg

                try:
                    run_ffmpeg(cmd, timeout=10)
                except Exception:
                    continue

                if os.path.exists(output_path):
                    # 读取帧并转换为 numpy 数组
                    img = self._load_image_as_array(output_path)
                    if img is not None:
                        frames.append(img)

        except Exception as e:
            logger.warning(f"Failed to extract frames: {e}")

        self._frames = frames
        return frames

    def _load_image_as_array(self, path: str) -> np.ndarray | None:
        """加载图片为 numpy 数组 (RGB 格式)"""
        try:
            from PIL import Image

            with Image.open(path) as img:
                if img.mode != "RGB":
                    img = img.convert("RGB")
                return np.array(img)
        except Exception as e:
            logger.warning(f"Failed to load image {path}: {e}")
            return None

    def analyze_color_distribution(self, frames: list[np.ndarray] | None = None) -> ColorAnalysis:
        """
        分析色彩分布 (HSV 空间)

        Returns:
            ColorAnalysis 对象
        """
        if frames is None:
            frames = self.extract_frames()

        if not frames:
            return ColorAnalysis()

        result = ColorAnalysis()
        all_hsv = []

        try:
            for frame in frames:
                # RGB 转 HSV
                rgb = frame.astype(float) / 255.0
                r, g, b = rgb[:, :, 0], rgb[:, :, 1], rgb[:, :, 2]

                maxc = np.maximum(np.maximum(r, g), b)
                minc = np.minimum(np.minimum(r, g), b)
                v = maxc
                s = np.where(maxc > 0, (maxc - minc) / maxc, 0)

                # 计算色相
                rc = np.where(maxc == r, (maxc - g - (maxc - b)) / (maxc - minc + 1e-10), 0)
                gc = np.where(maxc == g, 2.0 + (maxc - b - (maxc - r)) / (maxc - minc + 1e-10), 0)
                bc = np.where(maxc == b, 4.0 + (maxc - r - (maxc - g)) / (maxc - minc + 1e-10), 0)
                h = (rc + gc + bc) * 60
                h = np.where(h < 0, h + 360, h)

                frame_hsv = np.stack([h.flatten(), s.flatten(), v.flatten()], axis=1)
                all_hsv.append(frame_hsv)

            if all_hsv:
                all_hsv = np.vstack(all_hsv)

                # 主色调
                result.dominant_hue = float(np.median(all_hsv[:, 0]))

                # 饱和度
                result.avg_saturation = float(np.mean(all_hsv[:, 1]))

                # 亮度
                result.avg_brightness = float(np.mean(all_hsv[:, 2]))

                # 计算颜色比例
                # 绿色: 60-180 度
                green_mask = (all_hsv[:, 0] >= 60) & (all_hsv[:, 0] <= 180)
                result.green_ratio = float(np.mean(green_mask))

                # 暖色调 (红/黄/橙): 0-60, 300-360 度
                warm_mask = (all_hsv[:, 0] <= 60) | (all_hsv[:, 0] >= 300)
                result.warm_ratio = float(np.mean(warm_mask))

                # 冷色调 (蓝/青): 180-300 度
                cool_mask = (all_hsv[:, 0] >= 180) & (all_hsv[:, 0] <= 300)
                result.cool_ratio = float(np.mean(cool_mask))

        except Exception as e:
            logger.warning(f"Failed to analyze color distribution: {e}")

        return result

    def analyze_motion(self, frames: list[np.ndarray] | None = None) -> MotionAnalysis:
        """
        分析画面运动幅度

        Returns:
            MotionAnalysis 对象
        """
        if frames is None:
            frames = self.extract_frames()

        if len(frames) < 2:
            return MotionAnalysis()

        result = MotionAnalysis()
        motion_scores = []
        scene_changes = 0

        try:
            for i in range(len(frames) - 1):
                # 计算相邻帧差异
                diff = np.abs(frames[i + 1].astype(float) - frames[i].astype(float))
                mean_diff = np.mean(diff) / 255.0

                motion_scores.append(mean_diff)

                # 检测场景切换 (帧差异 > 30%)
                if mean_diff > 0.3:
                    scene_changes += 1

            if motion_scores:
                # 使用中位数避免异常值影响
                result.motion_score = float(np.median(motion_scores))
                # 归一化到 0-1
                result.motion_score = min(1.0, result.motion_score * 5)
                result.scene_changes = scene_changes

        except Exception as e:
            logger.warning(f"Failed to analyze motion: {e}")

        return result

    def analyze_audio(self) -> AudioAnalysis:
        """
        分析音频特征

        Returns:
            AudioAnalysis 对象
        """
        result = AudioAnalysis()
        info = self.get_video_info()

        if not info.has_audio:
            return result

        result.has_audio = True

        try:
            # 提取音频并分析频率特征
            audio_path = os.path.join(self._temp_dir, "audio.wav")
            cmd = [
                "ffmpeg",
                "-y",
                "-i",
                self.video_path,
                "-vn",  # 不要视频
                "-ac",
                "1",  # 单声道
                "-ar",
                "8000",  # 降低采样率
                "-f",
                "wav",
                audio_path,
            ]

            from video_processing.ffmpeg_utils import run_ffmpeg

            try:
                run_ffmpeg(cmd, timeout=30)
            except Exception:
                # 音频提取失败，返回默认分析结果
                return AudioAnalysis(  # type: ignore[call-arg]
                    has_speech=False,
                    speech_ratio=0.0,
                    avg_volume=0.0,
                )

            if os.path.exists(audio_path):
                # 读取音频数据
                import struct

                with open(audio_path, "rb") as f:
                    # 跳过 WAV 头
                    f.read(44)
                    audio_data = f.read()

                if len(audio_data) >= 2:
                    # 转换为 numpy 数组
                    audio_samples = np.array(struct.unpack(f"<{len(audio_data)//2}h", audio_data), dtype=float)
                    audio_samples = audio_samples / 32768.0

                    if len(audio_samples) > 0:
                        # 简单频谱分析
                        fft = np.abs(np.fft.rfft(audio_samples[: min(len(audio_samples), 8000)]))
                        freqs = np.fft.rfftfreq(min(len(audio_samples), 8000), 1 / 8000)

                        # 人声频率: 300-3400 Hz
                        speech_mask = (freqs >= 300) & (freqs <= 3400)
                        speech_energy = np.mean(fft[speech_mask]) if speech_mask.any() else 0

                        # 音乐低频: 60-250 Hz
                        bass_mask = (freqs >= 60) & (freqs <= 250)
                        bass_energy = np.mean(fft[bass_mask]) if bass_mask.any() else 0

                        # 环境音 (高频): > 4000 Hz
                        high_mask = freqs > 4000
                        high_energy = np.mean(fft[high_mask]) if high_mask.any() else 0

                        total_energy = speech_energy + bass_energy + high_energy + 1e-10

                        result.speech_ratio = float(speech_energy / total_energy)
                        result.music_ratio = float(bass_energy / total_energy)
                        result.ambient_ratio = float(high_energy / total_energy)

        except Exception as e:
            logger.warning(f"Failed to analyze audio: {e}")

        return result

    def classify(self) -> ClassificationResult:
        """
        综合分析得出分类结果

        Returns:
            ClassificationResult 对象
        """
        # 提取分析数据
        frames = self.extract_frames()
        color = self.analyze_color_distribution(frames)
        motion = self.analyze_motion(frames)
        audio = self.analyze_audio()

        # 计算各类别得分
        scores = self._calculate_category_scores(color, motion, audio)

        # 找最高分
        if not scores:
            return ClassificationResult(
                category=AssetClassification.OTHER,
                confidence=0.3,
                scores={},
            )

        best_category = max(scores.items(), key=lambda x: x[1])
        category = AssetClassification(best_category[0])
        confidence = min(0.95, max(0.3, best_category[1]))

        return ClassificationResult(
            category=category,
            confidence=confidence,
            scores=scores,
        )

    def _calculate_category_scores(
        self,
        color: ColorAnalysis,
        motion: MotionAnalysis,
        audio: AudioAnalysis,
    ) -> dict[str, float]:
        """
        计算各类别的置信度得分

        Args:
            color: 色彩分析结果
            motion: 运动分析结果
            audio: 音频分析结果

        Returns:
            各类别得分字典
        """
        scores = {}

        # 1. 风景 (scenic) - 绿色、户外、自然
        scenic_score = 0.0
        if color.green_ratio > 0.3:
            scenic_score += 0.4 * color.green_ratio
        if color.avg_saturation > 0.3:
            scenic_score += 0.2 * color.avg_saturation
        if color.avg_brightness > 0.4:
            scenic_score += 0.2
        if motion.motion_score > 0.1 and motion.motion_score < 0.5:
            scenic_score += 0.2  # 适度运动（如云朵、树叶）
        if not audio.has_audio or audio.ambient_ratio > 0.5:
            scenic_score += 0.2  # 自然环境音
        scores[AssetClassification.SCENIC.value] = min(1.0, scenic_score)

        # 2. 产品 (product) - 中等亮度、均匀色彩、低运动
        product_score = 0.0
        if 0.3 < color.avg_brightness < 0.7:
            product_score += 0.3
        if color.avg_saturation < 0.5:
            product_score += 0.2
        if motion.motion_score < 0.15:
            product_score += 0.4  # 低运动 = 产品展示
        if color.cool_ratio > 0.3:
            product_score += 0.2  # 冷色调 = 科技感
        scores[AssetClassification.PRODUCT.value] = min(1.0, product_score)

        # 3. 人物 (person) - 中等运动、有时有人声
        person_score = 0.0
        if 0.1 < motion.motion_score < 0.4:
            person_score += 0.3  # 适度运动
        if audio.has_audio and audio.speech_ratio > 0.3:
            person_score += 0.5  # 有人声
        if color.avg_brightness > 0.3:
            person_score += 0.2
        scores[AssetClassification.PERSON.value] = min(1.0, person_score)

        # 4. 动物 (animal) - 高运动、有时自然音
        animal_score = 0.0
        if motion.motion_score > 0.3:
            animal_score += 0.4  # 高运动
        if motion.scene_changes > 2:
            animal_score += 0.2
        if audio.has_audio and (audio.ambient_ratio > 0.3 or audio.speech_ratio > 0.2):
            animal_score += 0.3
        scores[AssetClassification.ANIMAL.value] = min(1.0, animal_score)

        # 5. 美食 (food) - 暖色调、高饱和度
        food_score = 0.0
        if color.warm_ratio > 0.4:
            food_score += 0.5
        if color.avg_saturation > 0.5:
            food_score += 0.3
        if 0.4 < color.avg_brightness < 0.8:
            food_score += 0.2
        scores[AssetClassification.FOOD.value] = min(1.0, food_score)

        # 6. 科技 (tech) - 冷色调、低饱和度、低运动
        tech_score = 0.0
        if color.cool_ratio > 0.4:
            tech_score += 0.4
        if color.avg_saturation < 0.4:
            tech_score += 0.3
        if motion.motion_score < 0.2:
            tech_score += 0.3
        scores[AssetClassification.TECH.value] = min(1.0, tech_score)

        # 7. 运动 (sport) - 高运动
        sport_score = 0.0
        if motion.motion_score > 0.4:
            sport_score += 0.6
        if motion.scene_changes > 3:
            sport_score += 0.2
        if color.avg_brightness > 0.4:
            sport_score += 0.2
        scores[AssetClassification.SPORT.value] = min(1.0, sport_score)

        # 8. 音乐 (music) - 有节奏性音乐
        music_score = 0.0
        if audio.has_audio and audio.music_ratio > 0.4:
            music_score += 0.6
        # 纯视觉判断：色彩丰富但非自然
        if color.avg_saturation > 0.5 and color.green_ratio < 0.2:
            music_score += 0.3
        scores[AssetClassification.MUSIC.value] = min(1.0, music_score)

        # 9. 其他 (other) - 默认最低分
        scores[AssetClassification.OTHER.value] = 0.1

        return scores

    def calculate_quality_score(self) -> QualityScore:
        """
        计算视频质量综合评分 (0-100)

        评分维度:
        1. 分辨率得分 (25分)
        2. 帧率得分 (20分)
        3. 码率得分 (20分)
        4. 清晰度得分 (20分) - Laplacian 方差
        5. 稳定性得分 (15分) - 帧间位移方差
        """
        info = self.get_video_info()
        frames = self.extract_frames()

        # 1. 分辨率得分
        resolution_score = self._score_resolution(info.width, info.height)

        # 2. 帧率得分
        fps_score = self._score_framerate(info.fps)

        # 3. 码率得分
        bitrate_score = self._score_bitrate(info.bitrate)

        # 4. 清晰度得分
        clarity_score = self._score_clarity(frames)

        # 5. 稳定性得分
        stability_score = self._score_stability(frames)

        total = resolution_score + fps_score + bitrate_score + clarity_score + stability_score

        return QualityScore(
            total=round(min(100, max(0, total)), 1),
            resolution_score=resolution_score,
            fps_score=fps_score,
            bitrate_score=bitrate_score,
            clarity_score=clarity_score,
            stability_score=stability_score,
        )

    def _score_resolution(self, width: int, height: int) -> float:
        """分辨率评分 (满分 25)"""
        pixels = width * height

        if pixels >= 3840 * 2160:  # 4K
            return 25.0
        elif pixels >= 2560 * 1440:  # 2K
            return 22.0
        elif pixels >= 1920 * 1080:  # 1080p
            return 20.0
        elif pixels >= 1280 * 720:  # 720p
            return 15.0
        elif pixels >= 854 * 480:  # 480p
            return 8.0
        else:
            return 3.0

    def _score_framerate(self, fps: float) -> float:
        """帧率评分 (满分 20)"""
        if fps >= 60:
            return 20.0
        elif fps >= 30:
            return 15.0
        elif fps >= 24:
            return 10.0
        elif fps >= 15:
            return 7.0
        else:
            return 5.0

    def _score_bitrate(self, bitrate: int) -> float:
        """码率评分 (满分 20)"""
        bitrate_mbps = bitrate / 1_000_000

        if bitrate_mbps > 10:
            return 20.0
        elif bitrate_mbps >= 5:
            return 15.0
        elif bitrate_mbps >= 2:
            return 10.0
        elif bitrate_mbps >= 0.5:
            return 5.0
        else:
            return 3.0

    def _score_clarity(self, frames: list[np.ndarray]) -> float:
        """
        清晰度评分 (满分 20)

        使用 Laplacian 方差评估画面清晰度
        高方差 = 细节丰富 = 高分
        """
        if not frames:
            return 10.0  # 默认中等分

        try:
            variances = []

            for frame in frames[:5]:  # 只分析前 5 帧
                if len(frame.shape) == 3:
                    # 转灰度
                    gray = np.dot(frame[..., :3], [0.299, 0.587, 0.114]).astype(np.uint8)
                else:
                    gray = frame

                # Laplacian 算子
                laplacian = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]], dtype=np.float32)

                # 手动计算卷积
                from scipy import signal

                laplacian_img = signal.convolve2d(gray.astype(float), laplacian, mode="same")
                variance = np.var(laplacian_img)
                variances.append(variance)

            # 归一化方差到 0-20 分
            avg_variance = np.mean(variances)
            # 根据经验值调整
            score = min(20.0, avg_variance / 100)
            return float(score)

        except ImportError:
            # 如果没有 scipy，使用简化方法
            return 10.0
        except Exception:
            return 10.0

    def _score_stability(self, frames: list[np.ndarray]) -> float:
        """
        稳定性评分 (满分 15)

        分析帧间位移方差
        画面稳定 = 高分
        剧烈抖动 = 低分
        """
        if len(frames) < 2:
            return 10.0  # 默认中等分

        try:
            displacements = []

            for i in range(len(frames) - 1):
                # 缩小帧以加速处理
                scale = 0.25
                new_h = int(frames[i].shape[0] * scale)
                new_w = int(frames[i].shape[1] * scale)
                frame1_small = np.array(Image.fromarray(frames[i]).resize((new_w, new_h)))
                new_h2 = int(frames[i + 1].shape[0] * scale)
                new_w2 = int(frames[i + 1].shape[1] * scale)
                frame2_small = np.array(Image.fromarray(frames[i + 1]).resize((new_w2, new_h2)))

                # 简单位移检测：灰度差
                gray1 = np.mean(frame1_small, axis=2) if len(frame1_small.shape) == 3 else frame1_small
                gray2 = np.mean(frame2_small, axis=2) if len(frame2_small.shape) == 3 else frame2_small

                diff = np.abs(gray2.astype(float) - gray1.astype(float))
                displacement = np.mean(diff) / 255.0
                displacements.append(displacement)

            # 高位移方差 = 不稳定
            if displacements:
                displacement_variance = np.var(displacements)
                # 归一化
                instability = min(1.0, displacement_variance * 10)
                score = 15.0 * (1.0 - instability)
                return float(max(0.0, score))

            return 10.0

        except Exception:
            return 10.0


def classify_asset_real(video_path: str) -> tuple[str, float]:
    """
    真实分类入口函数

    Args:
        video_path: 视频文件路径

    Returns:
        (分类类别, 置信度)
    """
    try:
        analyzer = AssetAnalyzer(video_path)
        result = analyzer.classify()
        return result.category.value, result.confidence
    except Exception as e:
        logger.warning(f"Classification failed, using fallback: {e}")
        return AssetClassification.OTHER.value, 0.3


def calculate_quality_score_real(video_path: str) -> float:
    """
    质量评分入口函数

    Args:
        video_path: 视频文件路径

    Returns:
        质量评分 (0-100)
    """
    try:
        analyzer = AssetAnalyzer(video_path)
        result = analyzer.calculate_quality_score()
        return result.total
    except Exception as e:
        logger.warning(f"Quality scoring failed, using fallback: {e}")
        return 50.0