"""
视频帧提取工具

这个模块提供了简单高效的视频帧提取功能。主要特点：
1. 使用ffmpeg进行视频处理，支持硬件加速
2. 按指定时间间隔提取视频关键帧
3. 支持多种视频格式
4. 支持高清视频帧输出
5. 直接从原视频提取高质量关键帧

不依赖OpenCV和sklearn等库，只使用ffmpeg作为外部依赖，降低了安装和使用的复杂度。
"""
import json
import os
import subprocess
from typing import List, Dict

from app.utils import ffmpeg_utils
from loguru import logger
from tqdm import tqdm


class VideoProcessor:
    def __init__(self, video_path: str):
        """
        初始化视频处理器

        Args:
            video_path: 视频文件路径
        """
        if not os.path.exists(video_path):
            raise FileNotFoundError(f"视频文件不存在: {video_path}")

        self.video_path = video_path
        self.video_info = self._get_video_info()
        print(f'video_info----{self.video_info}')
        self.fps = float(self.video_info.get('fps', 25))
        self.duration = float(self.video_info.get('duration', 0))
        self.width = int(self.video_info.get('width', 0))
        self.height = int(self.video_info.get('height', 0))
        self.total_frames = int(self.fps * self.duration)

    def _get_video_info(self) -> Dict[str, str]:
        cmd = [
            "ffprobe",
            "-v", "error",
            "-select_streams", "v:0",
            "-show_entries", "stream=width,height,r_frame_rate,duration:format=duration",
            "-of", "json",
            self.video_path
        ]

        try:
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                check=True
            )

            data = json.loads(result.stdout)

            stream = data.get("streams", [{}])[0]
            fmt = data.get("format", {})

            info = {}

            info["width"] = str(stream.get("width") or 1280)
            info["height"] = str(stream.get("height") or 720)

            # MKV 的 stream.duration 可能是 N/A，所以用 format.duration 兜底
            duration = stream.get("duration")
            if not duration or str(duration).upper() == "N/A":
                duration = fmt.get("duration")

            if not duration or str(duration).upper() == "N/A":
                duration = 0

            info["duration"] = str(duration)

            r_frame_rate = stream.get("r_frame_rate", "25/1")
            try:
                num, den = map(int, r_frame_rate.split("/"))
                info["fps"] = str(num / den) if den != 0 else "25"
            except Exception:
                info["fps"] = "25"

            return info

        except subprocess.CalledProcessError as e:
            logger.error(f"获取视频信息失败: {e.stderr}")
            return {
                "width": "1280",
                "height": "720",
                "fps": "25",
                "duration": "0"
            }
        except Exception as e:
            logger.error(f"获取视频信息异常: {e}")
            return {
                "width": "1280",
                "height": "720",
                "fps": "25",
                "duration": "0"
            }

    def extract_frames_by_interval(self, output_dir: str, interval_seconds: float = 5.0,
                                   use_hw_accel: bool = True, signals=None, one_image=False,
                                   scale_width=1280,
                                   scale_height=720) -> List[int]:
        """
        按指定时间间隔提取视频帧

        优化了 Windows 系统兼容性，特别是 N 卡硬件加速的滤镜链问题

        Args:
            output_dir: 输出目录
            interval_seconds: 帧提取间隔（秒）
            use_hw_accel: 是否使用硬件加速

        Returns:
            List[int]: 提取的帧号列表
        """
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # 计算起始时间和帧提取点
        start_time = 0
        end_time = self.duration
        extraction_times = []

        current_time = start_time
        while current_time < end_time:
            extraction_times.append(current_time)
            current_time += interval_seconds

        if not extraction_times:
            logger.warning("未找到需要提取的帧")
            return []

        # 获取硬件加速信息
        hwaccel_info = ffmpeg_utils.get_ffmpeg_hwaccel_info()
        hwaccel_type = hwaccel_info.get("type", "software")

        # 提取帧 - 使用优化的进度条
        frame_numbers = []
        successful_extractions = 0
        failed_extractions = 0

        logger.info(f"开始提取 {len(extraction_times)} 个关键帧，使用 {hwaccel_type} 加速")

        with tqdm(total=len(extraction_times), desc="🎬 提取视频帧", unit="帧",
                  bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]") as pbar:
            for i, timestamp in enumerate(extraction_times):
                frame_number = int(timestamp * self.fps)
                frame_numbers.append(frame_number)

                # 格式化时间戳字符串 (HHMMSSmmm)
                hours = int(timestamp // 3600)
                minutes = int((timestamp % 3600) // 60)
                seconds = int(timestamp % 60)
                milliseconds = int((timestamp % 1) * 1000)
                time_str = f"{hours:02d}{minutes:02d}{seconds:02d}{milliseconds:03d}"

                output_path = os.path.join(output_dir, f"keyframe_{frame_number:06d}_{time_str}.webp")

                # 构建 FFmpeg 命令 - 针对 Windows N 卡优化
                success = self._extract_single_frame_optimized(
                    timestamp, output_path, use_hw_accel, hwaccel_type,
                    scale_width, scale_height
                )

                if success:
                    successful_extractions += 1
                    pbar.set_postfix({
                        "✅": successful_extractions,
                        "❌": failed_extractions,
                        "时间": f"{timestamp:.1f}s"
                    })
                else:
                    failed_extractions += 1
                    pbar.set_postfix({
                        "✅": successful_extractions,
                        "❌": failed_extractions,
                        "时间": f"{timestamp:.1f}s"
                    })
                if signals:
                    signals.update_label.emit(
                        f"10% - 正在提取关键帧，{failed_extractions + successful_extractions}/{len(extraction_times)}")
                pbar.update(1)
                if one_image:
                    break

        # 统计结果
        total_attempts = len(extraction_times)
        success_rate = (successful_extractions / total_attempts) * 100 if total_attempts > 0 else 0

        logger.info(f"关键帧提取完成: 成功 {successful_extractions}/{total_attempts} 帧 ({success_rate:.1f}%)")

        if failed_extractions > 0:
            logger.warning(f"有 {failed_extractions} 帧提取失败，可能是硬件加速兼容性问题")

        # 验证实际生成的文件
        actual_files = [f for f in os.listdir(output_dir) if f.endswith('.webp')]
        logger.info(f"实际生成文件数量: {len(actual_files)} 个")

        if len(actual_files) == 0:
            logger.error("未生成任何关键帧文件，可能需要禁用硬件加速")
            raise Exception("关键帧提取完全失败，请检查视频文件和 FFmpeg 配置")

        return frame_numbers

    def _extract_single_frame_optimized(self, timestamp: float, output_path: str,
                                        use_hw_accel: bool, hwaccel_type: str,
                                        scale_width=1280,
                                        scale_height=720) -> bool:
        """
        优化的单帧提取方法，解决 Windows N 卡硬件加速兼容性问题

        Args:
            timestamp: 时间戳（秒）
            output_path: 输出文件路径
            use_hw_accel: 是否使用硬件加速
            hwaccel_type: 硬件加速类型

        Returns:
            bool: 是否成功提取
        """
        # 策略1: 优先尝试纯编码器方案（避免硬件解码滤镜链问题）
        if use_hw_accel and hwaccel_type in ["nvenc", "cuda"]:
            # 对于 NVIDIA 显卡，优先使用纯软件解码 + NVENC 编码
            if self._try_extract_with_software_decode(timestamp, output_path, scale_width, scale_height):
                return True

        # 策略2: 尝试标准硬件加速
        if use_hw_accel and ffmpeg_utils.is_ffmpeg_hwaccel_available():
            hw_accel = ffmpeg_utils.get_ffmpeg_hwaccel_args()
            if self._try_extract_with_hwaccel(timestamp, output_path, hw_accel):
                return True

        # 策略3: 软件方案
        if self._try_extract_with_software(timestamp, output_path):
            return True

        # 策略4: 超级兼容性方案（Windows 特殊处理）
        return self._try_extract_with_ultra_compatibility(timestamp, output_path)

    def _try_extract_with_software_decode(self, timestamp: float, output_path: str,
                               scale_width=1280,
                               scale_height=720) -> bool:
        """
        使用纯软件解码提取帧（推荐用于 Windows N 卡）
        参考 clip_video.py 中的成功实现

        Args:
            timestamp: 时间戳
            output_path: 输出路径

        Returns:
            bool: 是否成功
        """
        # 参考 clip_video.py 中的兼容性方案，专门针对图片输出优化
        cmd = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel", "error",
            "-ss", str(timestamp),  # 先定位时间戳
            "-i", self.video_path,
            "-vframes", "1",  # 只提取一帧
            "-vf", f"scale='min({scale_width},iw)':-2",
            "-compression_level", "6",
            "-q:v", "75",  # 高质量
            "-pix_fmt", "yuv420p",  # 明确指定像素格式
            "-y",
            output_path
        ]
        print(f'cmd: {cmd}')
        return self._execute_ffmpeg_command(cmd, f"软件解码提取帧 {timestamp:.1f}s")

    def _try_extract_with_hwaccel(self, timestamp: float, output_path: str, hw_accel: List[str]) -> bool:
        """
        使用硬件加速提取帧

        Args:
            timestamp: 时间戳
            output_path: 输出路径
            hw_accel: 硬件加速参数

        Returns:
            bool: 是否成功
        """
        cmd = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel", "error",
        ]

        # 添加硬件加速参数
        cmd.extend(hw_accel)

        cmd.extend([
            "-ss", str(timestamp),
            "-i", self.video_path,
            "-vframes", "1",
            "-q:v", "2",
            "-pix_fmt", "yuv420p",
            "-y",
            output_path
        ])

        return self._execute_ffmpeg_command(cmd, f"硬件加速提取帧 {timestamp:.1f}s")

    def _try_extract_with_software(self, timestamp: float, output_path: str) -> bool:
        """
        使用纯软件方案提取帧（最后的备用方案）
        参考 clip_video.py 中的基本编码方案

        Args:
            timestamp: 时间戳
            output_path: 输出路径

        Returns:
            bool: 是否成功
        """
        # 最基本的兼容性方案，参考 clip_video.py 的 try_basic_fallback
        cmd = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel", "warning",  # 更详细的日志用于调试
            "-ss", str(timestamp),
            "-i", self.video_path,
            "-vframes", "1",
            "-q:v", "3",  # 稍微降低质量以提高兼容性
            "-pix_fmt", "yuv420p",
            "-avoid_negative_ts", "make_zero",  # 避免时间戳问题
            "-y",
            output_path
        ]

        return self._execute_ffmpeg_command(cmd, f"软件方案提取帧 {timestamp:.1f}s")

    def _try_extract_with_ultra_compatibility(self, timestamp: float, output_path: str) -> bool:
        """
        超级兼容性方案，专门解决 Windows 系统的 MJPEG 编码问题

        Args:
            timestamp: 时间戳
            output_path: 输出路径

        Returns:
            bool: 是否成功
        """
        # 方案1: 使用 PNG 格式避免 MJPEG 问题
        png_output = output_path.replace('.jpg', '.png')
        cmd1 = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel", "error",
            "-ss", str(timestamp),
            "-i", self.video_path,
            "-vframes", "1",
            "-f", "image2",  # 明确指定图片格式
            "-y",
            png_output
        ]

        if self._execute_ffmpeg_command(cmd1, f"PNG格式提取帧 {timestamp:.1f}s"):
            # 如果 PNG 成功，转换为 JPG
            try:
                from PIL import Image
                with Image.open(png_output) as img:
                    # 转换为 RGB 模式（去除 alpha 通道）
                    if img.mode in ('RGBA', 'LA'):
                        background = Image.new('RGB', img.size, (255, 255, 255))
                        background.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None)
                        img = background
                    img.save(output_path, 'JPEG', quality=90)

                # 删除临时 PNG 文件
                os.remove(png_output)
                return True
            except Exception as e:
                logger.debug(f"PNG 转 JPG 失败: {e}")
                # 如果转换失败，直接重命名 PNG 为 JPG
                try:
                    os.rename(png_output, output_path)
                    return True
                except Exception:
                    pass

        # 方案2: 使用最简单的参数
        cmd2 = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel", "error",
            "-i", self.video_path,
            "-ss", str(timestamp),  # 把 -ss 放在 -i 后面
            "-vframes", "1",
            "-f", "mjpeg",  # 明确指定 MJPEG 格式
            "-q:v", "5",  # 降低质量要求
            "-y",
            output_path
        ]

        if self._execute_ffmpeg_command(cmd2, f"MJPEG格式提取帧 {timestamp:.1f}s"):
            return True

        # 方案3: 最后的尝试 - 使用 BMP 格式
        bmp_output = output_path.replace('.jpg', '.bmp')
        cmd3 = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel", "error",
            "-i", self.video_path,
            "-ss", str(timestamp),
            "-vframes", "1",
            "-f", "bmp",
            "-y",
            bmp_output
        ]

        if self._execute_ffmpeg_command(cmd3, f"BMP格式提取帧 {timestamp:.1f}s"):
            # 尝试转换 BMP 为 JPG
            try:
                from PIL import Image
                with Image.open(bmp_output) as img:
                    img.save(output_path, 'JPEG', quality=90)
                os.remove(bmp_output)
                return True
            except Exception:
                # 如果转换失败，直接重命名
                try:
                    os.rename(bmp_output, output_path)
                    return True
                except Exception:
                    pass

        return False

    def _execute_ffmpeg_command(self, cmd: List[str], description: str) -> bool:
        """
        执行 FFmpeg 命令并处理结果
        参考 clip_video.py 中的错误处理机制

        Args:
            cmd: FFmpeg 命令列表
            description: 操作描述

        Returns:
            bool: 是否成功
        """
        try:
            # 参考 clip_video.py 中的 Windows 处理方式
            is_windows = os.name == 'nt'
            process_kwargs = {
                "stdout": subprocess.PIPE,
                "stderr": subprocess.PIPE,
                "text": True,
                "check": True,
                "timeout": 30  # 30秒超时
            }

            if is_windows:
                process_kwargs["encoding"] = 'utf-8'
            print(f'cmd---{" ".join(cmd)}')
            result = subprocess.run(cmd, **process_kwargs)

            # 验证输出文件
            output_path = cmd[-1]
            if os.path.exists(output_path) and os.path.getsize(output_path) > 0:
                return True
            else:
                return False

        except subprocess.CalledProcessError as e:
            print(f'_execute_ffmpeg_command---{e}')
            # 简化错误日志，仅记录关键信息
            return False
        except subprocess.TimeoutExpired:
            return False
        except Exception as e:
            print(f'_execute_ffmpeg_command---{e}')
            return False

    def _detect_hw_accelerator(self) -> List[str]:
        """
        检测系统可用的硬件加速器

        Returns:
            List[str]: 硬件加速器ffmpeg命令参数
        """
        # 使用集中式硬件加速检测
        if ffmpeg_utils.is_ffmpeg_hwaccel_available():
            return ffmpeg_utils.get_ffmpeg_hwaccel_args()
        return []

    def process_video_pipeline(self,
                               output_dir: str,
                               interval_seconds: float = 5.0,  # 帧提取间隔（秒）
                               use_hw_accel: bool = True,
                               signals=None,
                               one_image=False,
                               scale_width=1280,
                               scale_height=720) -> None:
        """
        执行简化的视频处理流程，直接从原视频按固定时间间隔提取帧

        Args:
            output_dir: 输出目录
            interval_seconds: 帧提取间隔（秒）
            use_hw_accel: 是否使用硬件加速
        """
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)

        try:
            # 直接从原视频提取关键帧
            print(f"从视频间隔 {interval_seconds} 秒提取关键帧...")
            self.extract_frames_by_interval(
                output_dir,
                interval_seconds=interval_seconds,
                use_hw_accel=use_hw_accel,
                signals=signals,
                one_image=one_image,
                scale_width=scale_width,
                scale_height=scale_height
            )

            logger.info(f"处理完成！视频帧已保存在: {output_dir}")

        except Exception as e:
            import traceback
            logger.error(f"视频处理失败: \n{traceback.format_exc()}")
            raise

    def extract_frames_by_interval_ultra_compatible(self, output_dir: str, interval_seconds: float = 5.0) -> List[int]:
        """
        使用超级兼容性方案按指定时间间隔提取视频帧
        
        直接使用PNG格式提取，避免MJPEG编码问题，确保最高兼容性
        
        Args:
            output_dir: 输出目录
            interval_seconds: 帧提取间隔（秒）
            
        Returns:
            List[int]: 提取的帧号列表
        """
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        # 计算起始时间和帧提取点
        start_time = 0
        end_time = self.duration
        extraction_times = []

        current_time = start_time
        while current_time < end_time:
            extraction_times.append(current_time)
            current_time += interval_seconds

        if not extraction_times:
            logger.warning("未找到需要提取的帧")
            return []

        # 提取帧 - 使用美化的进度条
        frame_numbers = []
        successful_extractions = 0
        failed_extractions = 0

        logger.info(f"开始提取 {len(extraction_times)} 个关键帧，使用超级兼容性方案")

        with tqdm(total=len(extraction_times), desc="🎬 提取关键帧", unit="帧",
                  bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]") as pbar:
            for i, timestamp in enumerate(extraction_times):
                frame_number = int(timestamp * self.fps)
                frame_numbers.append(frame_number)

                # 格式化时间戳字符串 (HHMMSSmmm)
                hours = int(timestamp // 3600)
                minutes = int((timestamp % 3600) // 60)
                seconds = int(timestamp % 60)
                milliseconds = int((timestamp % 1) * 1000)
                time_str = f"{hours:02d}{minutes:02d}{seconds:02d}{milliseconds:03d}"

                output_path = os.path.join(output_dir, f"keyframe_{frame_number:06d}_{time_str}.jpg")

                # 直接使用超级兼容性方案
                success = self._extract_frame_ultra_compatible(timestamp, output_path)

                if success:
                    successful_extractions += 1
                    pbar.set_postfix({
                        "✅": successful_extractions,
                        "❌": failed_extractions,
                        "时间": f"{timestamp:.1f}s"
                    })
                else:
                    failed_extractions += 1
                    pbar.set_postfix({
                        "✅": successful_extractions,
                        "❌": failed_extractions,
                        "时间": f"{timestamp:.1f}s"
                    })

                pbar.update(1)

        # 统计结果
        total_attempts = len(extraction_times)
        success_rate = (successful_extractions / total_attempts) * 100 if total_attempts > 0 else 0

        logger.info(f"关键帧提取完成: 成功 {successful_extractions}/{total_attempts} 帧 ({success_rate:.1f}%)")

        if failed_extractions > 0:
            logger.warning(f"有 {failed_extractions} 帧提取失败")

        # 验证实际生成的文件
        actual_files = [f for f in os.listdir(output_dir) if f.endswith('.jpg')]
        logger.info(f"实际生成文件数量: {len(actual_files)} 个")

        if len(actual_files) == 0:
            logger.error("未生成任何关键帧文件")
            raise Exception("关键帧提取完全失败，请检查视频文件")

        return frame_numbers

    def _extract_frame_ultra_compatible(self, timestamp: float, output_path: str) -> bool:
        """
        超级兼容性方案提取单帧
        
        Args:
            timestamp: 时间戳（秒）
            output_path: 输出文件路径
            
        Returns:
            bool: 是否成功提取
        """
        # 使用 PNG 格式避免 MJPEG 问题
        png_output = output_path.replace('.jpg', '.png')
        cmd = [
            "ffmpeg",
            "-hide_banner",
            "-loglevel", "error",
            "-ss", str(timestamp),
            "-i", self.video_path,
            "-vframes", "1",
            "-f", "image2",  # 明确指定图片格式
            "-y",
            png_output
        ]
        print(cmd)

        try:
            # 执行FFmpeg命令
            result = subprocess.run(cmd, capture_output=True, text=True, check=True, timeout=30)

            # 验证PNG文件是否成功生成
            if os.path.exists(png_output) and os.path.getsize(png_output) > 0:
                # 转换PNG为JPG
                try:
                    from PIL import Image
                    with Image.open(png_output) as img:
                        # 转换为 RGB 模式（去除 alpha 通道）
                        if img.mode in ('RGBA', 'LA'):
                            background = Image.new('RGB', img.size, (255, 255, 255))
                            background.paste(img, mask=img.split()[-1] if img.mode == 'RGBA' else None)
                            img = background
                        img.save(output_path, 'JPEG', quality=90)

                    # 删除临时 PNG 文件
                    os.remove(png_output)
                    return True
                except Exception as e:
                    logger.warning(f"PNG 转 JPG 失败: {e}")
                    # 如果转换失败，直接重命名 PNG 为 JPG
                    try:
                        os.rename(png_output, output_path)
                        return True
                    except Exception:
                        return False
            else:
                return False

        except subprocess.CalledProcessError as e:
            logger.warning(f"超级兼容性方案提取帧 {timestamp:.1f}s 失败: {e}")
            return False
        except subprocess.TimeoutExpired:
            logger.warning(f"超级兼容性方案提取帧 {timestamp:.1f}s 超时")
            return False
        except Exception as e:
            logger.warning(f"超级兼容性方案提取帧 {timestamp:.1f}s 异常: {e}")
            return False


if __name__ == "__main__":
    import time

    start_time = time.time()

    # 使用示例
    processor = VideoProcessor("./resource/videos/test.mp4")

    # 设置间隔为3秒提取帧
    processor.process_video_pipeline(
        output_dir="output",
        interval_seconds=3.0,
        use_hw_accel=True
    )

    end_time = time.time()
    print(f"处理完成！总耗时: {end_time - start_time:.2f} 秒")
