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

from .asset_quality_scoring import (
    AudioAnalysis,
    ClassificationResult,
    ColorAnalysis,
    MotionAnalysis,
    QualityScore,
    VideoInfo,
    calculate_category_scores,
    calculate_quality_score,
    classify_from_analysis,
)

logger = logging.getLogger(__name__)


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:
        """
        综合分析得出分类结果

        评分逻辑在 asset_quality_scoring.calculate_category_scores / classify_from_analysis
        纯函数中，此处只负责采集分析数据后委托计算。

        Returns:
            ClassificationResult 对象
        """
        frames = self.extract_frames()
        color = self.analyze_color_distribution(frames)
        motion = self.analyze_motion(frames)
        audio = self.analyze_audio()

        return classify_from_analysis(color, motion, audio)

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

        评分逻辑在 asset_quality_scoring.calculate_quality_score 纯函数中，
        此处只负责采集数据后委托计算。

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

        return calculate_quality_score(info, frames)


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