# nodes.py
# 主要是各个节点的定义和业务逻辑
from datetime import datetime

import folder_paths

from .service import (
    extract_keyframes,
    split_and_save_video_segments
)


class ExtractKeyFrames:
    """
    用于从图像序列中提取关键帧的节点类。支持多种相似性计算方法（如 SSIM、MSE 和 pHash），
    并根据设定阈值判断是否为显著变化的关键帧，并进行合并与拆分控制以适应后续处理需求。
    """

    @classmethod
    def INPUT_TYPES(cls):
        """
        定义节点所需输入参数的信息。

        返回:
            dict: 包含 required 字段的字典，描述必填参数及属性：
                - image: 输入的图像序列，格式为 IMAGE 类型。
                - threshold: 浮点型，用于调节关键帧识别敏感度，默认 0.1。
                - merge_inter_frames: 整数型，定义合并相邻关键帧的最小间隔，默认 10。
                - max_frames: 整数型，单批次允许的最大帧数，默认 999999。
                - similarity_method: 枚举型，选择使用的图像相似度算法，可选 ["SSIM", "MSE", "pHash"]。
        """
        return {
            "required": {
                "image": ("IMAGE",),
                "threshold": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.00, "step": 0.01}),
                "merge_inter_frames": ("INT", {"default": 10, "min": 0, "max": 999, "step": 1}),
                "max_frames": ("INT", {"default": 999999, "min": 1, "max": 999999999, "step": 1}),
                "similarity_method": (["SSIM", "MSE", "pHash"], {"default": "SSIM"}),
            }
        }

    RETURN_TYPES = ("IMAGE", "STRING", "INT")
    RETURN_NAMES = ("KEYFRAMES", "indexes_string", "indexes_int")

    CATEGORY = "ExtractKeyFrames(关键帧处理)"
    FUNCTION = "execute"

    def execute(self, image, threshold, merge_inter_frames, max_frames, similarity_method):
        """
        执行关键帧提取流程。

        参数:
            image: 图像序列张量，维度为 [B, H, W, C]。
            threshold: 控制关键帧检测灵敏度的浮点值。
            merge_inter_frames: 最少合并间隔帧数。
            max_frames: 单个批次最大帧数限制。
            similarity_method: 选择使用的图像相似度算法。

        返回:
            tuple: 包含以下三个元素：
                - 提取的关键帧图像张量。
                - 关键帧索引组成的字符串。
                - 关键帧索引组成的整数列表。
        """
        # 调用service层的关键帧提取逻辑
        start=datetime.now()
        keyframe_indices = extract_keyframes(
            image, threshold, merge_inter_frames, max_frames, similarity_method
        )

        # 生成返回结果
        keyframe_images = image[keyframe_indices]
        indexes_string = ','.join(map(str, keyframe_indices))
        print(f"用时{datetime.now()-start}")
        return (keyframe_images, indexes_string, keyframe_indices)


class ExtractKeyFramesAndSplit:
    """
    根据关键帧提取结果，将原视频按关键帧位置切割成多个子视频片段。
    支持音频同步切割和多种视频格式输出。
    """

    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "image": ("IMAGE",),
                "threshold": ("FLOAT", {"default": 0.1, "min": -1.0, "max": 1.00, "step": 0.01}),
                "merge_inter_frames": ("INT", {"default": 10, "min": 0, "max": 999, "step": 1}),
                "max_frames": ("INT", {"default": 999999, "min": 1, "max": 999999999, "step": 1}),
                "similarity_method": (["SSIM", "MSE", "pHash"], {"default": "SSIM"}),
                "frame_rate": ("INT", {"default": 8, "min": 1, "max": 120, "step": 1}),
                "filename_prefix": ("STRING", {"default": "KeyFrameSegment"}),
                "format": (["h264-mp4", "h265-mp4", "vp9-webm"], {"default": "h264-mp4"}),
                "save_segments": ("BOOLEAN", {"default": True}),
                "video_quality": ("INT", {"default": 23, "min": 0, "max": 51, "step": 1}),
            },
            "optional": {
                "audio": ("AUDIO",),
            },
        }

    RETURN_TYPES = ("IMAGE", "STRING", "INT", "STRING")
    RETURN_NAMES = ("KEYFRAMES", "indexes_string", "indexes_int", "segment_files")
    OUTPUT_NODE = True
    FUNCTION = "execute"
    CATEGORY = "ExtractKeyFrames"

    def execute(self, image, threshold, merge_inter_frames, max_frames,
                similarity_method, frame_rate, filename_prefix, format,
                save_segments=True, video_quality=23, audio=None,):
        """
        执行关键帧提取和视频切割

        参数:
            image: 输入图像序列。
            threshold: 相似度检测阈值。
            merge_inter_frames: 合并相邻关键帧的最小间隔帧数。
            max_frames: 单个片段最大帧数。
            similarity_method: 相似度计算方法。
            frame_rate: 输出视频帧率。
            filename_prefix: 输出文件名前缀。
            format: 视频编码格式。
            save_segments: 是否保存分割后的视频片段。
            video_quality: 视频质量（CRF值）。
            audio: 可选音频数据。
            prompt: 元数据提示信息。
            extra_pnginfo: 额外PNG信息。

        返回:
            tuple: 包含关键帧图像、索引字符串、索引列表和视频片段路径列表。
        """

        # 1. 提取关键帧索引（复用service层逻辑）
        keyframe_indices = extract_keyframes(
            image, threshold, merge_inter_frames, max_frames, similarity_method
        )

        # 2. 获取关键帧图像
        keyframe_images = image[keyframe_indices]
        indexes_string = ','.join(map(str, keyframe_indices))

        # 3. 如果需要保存视频片段
        segment_files = []
        if save_segments:
            # 获取输出路径
            output_dir = folder_paths.get_output_directory()
            full_output_folder, filename, _, subfolder, _ = folder_paths.get_save_image_path(
                filename_prefix, output_dir
            )

            # 调用service层的视频切割逻辑（只传递必要参数）
            segment_files = split_and_save_video_segments(
                image=image,
                keyframe_indices=keyframe_indices,
                frame_rate=frame_rate,
                format=format,
                video_quality=video_quality,
                output_folder=full_output_folder,
                filename=filename,
                audio=audio
            )

        segment_files_str = '\n'.join(segment_files) if segment_files else ""

        print(f"\n✓ 完成！共提取 {len(keyframe_indices)} 个关键帧，生成 {len(segment_files)} 个视频片段")

        return (keyframe_images, indexes_string, keyframe_indices, segment_files_str)