"""
ComfyUI video trimming node.

Usage example inside ComfyUI:
    Load Video -> Video Trimmer -> Save Video

The node accepts a native ``VIDEO`` object, trims it by time or frame range,
and returns another native ``VIDEO`` object that remains compatible with
downstream video nodes such as Save Video.
"""

from __future__ import annotations

from fractions import Fraction
from typing import Any, Dict, Optional, Tuple


try:
    from comfy_api.latest import InputImpl, Types
except ImportError as exc:  # pragma: no cover - import error depends on host ComfyUI version
    raise ImportError(
        "ComfyUI's VIDEO API is unavailable. This node requires a ComfyUI build "
        "that exposes `comfy_api.latest` and native VIDEO components."
    ) from exc


def _safe_float(value: Any, default: float = 0.0) -> float:
    """Convert an arbitrary numeric value to float with a safe fallback."""
    try:
        return float(value)
    except (TypeError, ValueError):
        return default


def _copy_audio_segment(
    audio: Optional[Dict[str, Any]],
    start_frame: int,
    end_frame: int,
    fps: float,
) -> Optional[Dict[str, Any]]:
    """
    Trim audio so it stays synchronized with the selected video frames.

    Args:
        audio: ComfyUI audio payload or ``None``.
        start_frame: First included video frame index.
        end_frame: Exclusive ending video frame index.
        fps: Video frame rate.

    Returns:
        A trimmed audio payload compatible with ComfyUI, or ``None`` if there is
        no usable audio track.
    """
    if not audio or fps <= 0:
        return None

    waveform = audio.get("waveform")
    sample_rate = int(audio.get("sample_rate", 44100))

    if waveform is None or sample_rate <= 0:
        return None

    start_sample = int((start_frame / fps) * sample_rate)
    end_sample = int((end_frame / fps) * sample_rate)

    total_samples = waveform.shape[-1]
    start_sample = max(0, min(start_sample, total_samples))
    end_sample = max(start_sample, min(end_sample, total_samples))

    return {
        "waveform": waveform[..., start_sample:end_sample],
        "sample_rate": sample_rate,
    }


class VideoTrimmer:
    """
    Trim a native ComfyUI VIDEO input by timestamp or by frame range.

    The implementation works directly on ComfyUI's in-memory video components,
    which keeps the node easy to integrate with existing Load Video / Save Video
    workflows and avoids unnecessary file-system round-trips.
    """

    @classmethod
    def INPUT_TYPES(cls) -> Dict[str, Dict[str, Tuple[str, Dict[str, Any]]]]:
        """
        Describe the node inputs shown in the ComfyUI interface.

        ``trim_mode`` lets users choose an intuitive way to define the cut:
        either seconds or frame indices. Numeric controls are left available in
        the UI so users can quickly switch modes without reconnecting anything.
        """
        return {
            "required": {
                "video": ("VIDEO",),
                "trim_mode": (
                    ["time", "frames"],
                    {
                        "default": "time",
                        "tooltip": "Choose whether the trim range is defined in seconds or frame numbers.",
                    },
                ),
                "start_time": (
                    "FLOAT",
                    {
                        "default": 0.0,
                        "min": 0.0,
                        "max": 86400.0,
                        "step": 0.01,
                        "display": "number",
                        "tooltip": "Trim start in seconds when trim_mode is set to time.",
                    },
                ),
                "end_time": (
                    "FLOAT",
                    {
                        "default": 0.0,
                        "min": 0.0,
                        "max": 86400.0,
                        "step": 0.01,
                        "display": "number",
                        "tooltip": "Trim end in seconds when trim_mode is set to time. Use 0 to keep everything to the end.",
                    },
                ),
                "start_frame": (
                    "INT",
                    {
                        "default": 0,
                        "min": 0,
                        "max": 100000000,
                        "step": 1,
                        "display": "number",
                        "tooltip": "Trim start frame when trim_mode is set to frames.",
                    },
                ),
                "end_frame": (
                    "INT",
                    {
                        "default": 0,
                        "min": 0,
                        "max": 100000000,
                        "step": 1,
                        "display": "number",
                        "tooltip": "Trim end frame when trim_mode is set to frames. Use 0 to keep everything to the end.",
                    },
                ),
            }
        }

    RETURN_TYPES = ("VIDEO", "FLOAT", "INT", "INT", "INT")
    RETURN_NAMES = (
        "trimmed_video",
        "duration_seconds",
        "frame_count",
        "resolved_start_frame",
        "resolved_end_frame",
    )
    FUNCTION = "trim_video"
    CATEGORY = "Video Editing"

    def trim_video(
        self,
        video: Any,
        trim_mode: str,
        start_time: float,
        end_time: float,
        start_frame: int,
        end_frame: int,
    ) -> Tuple[Any, float, int, int, int]:
        """
        Trim an incoming VIDEO object and return a new VIDEO object.

        Args:
            video: Native ComfyUI VIDEO input from a node such as Load Video.
            trim_mode: Either ``"time"`` or ``"frames"``.
            start_time: Start timestamp in seconds.
            end_time: End timestamp in seconds, or 0 to use the video end.
            start_frame: Start frame index.
            end_frame: Exclusive end frame index, or 0 to use the video end.

        Returns:
            A tuple containing the trimmed video, resulting duration, resulting
            frame count, and the resolved start/end frame indices.

        Raises:
            ValueError: If the incoming video object is invalid or the range is invalid.
        """
        components = video.get_components()
        frames = getattr(components, "images", None)
        audio = getattr(components, "audio", None)
        metadata = getattr(components, "metadata", None)
        fps = _safe_float(getattr(components, "frame_rate", 0.0))

        if frames is None:
            raise ValueError("The input VIDEO object does not contain image frames.")

        if getattr(frames, "ndim", 0) < 4:
            raise ValueError(
                "The input VIDEO frames are malformed. Expected a tensor shaped like [frames, height, width, channels]."
            )

        total_frames = int(frames.shape[0])
        if total_frames <= 0:
            raise ValueError("The input VIDEO object is empty and cannot be trimmed.")

        if fps <= 0:
            raise ValueError("The input VIDEO object does not expose a valid frame rate.")

        resolved_start, resolved_end = self._resolve_frame_range(
            trim_mode=trim_mode,
            start_time=start_time,
            end_time=end_time,
            start_frame=start_frame,
            end_frame=end_frame,
            fps=fps,
            total_frames=total_frames,
        )

        trimmed_frames = frames[resolved_start:resolved_end]
        trimmed_audio = _copy_audio_segment(audio, resolved_start, resolved_end, fps)

        trimmed_video = InputImpl.VideoFromComponents(
            Types.VideoComponents(
                images=trimmed_frames,
                audio=trimmed_audio,
                frame_rate=Fraction(str(fps)),
                metadata=metadata,
            )
        )

        output_frame_count = int(trimmed_frames.shape[0])
        output_duration = output_frame_count / fps

        return (
            trimmed_video,
            output_duration,
            output_frame_count,
            resolved_start,
            resolved_end,
        )

    def _resolve_frame_range(
        self,
        trim_mode: str,
        start_time: float,
        end_time: float,
        start_frame: int,
        end_frame: int,
        fps: float,
        total_frames: int,
    ) -> Tuple[int, int]:
        """
        Convert UI inputs into a validated frame slice.

        ``end_*`` values are treated as exclusive bounds, which lines up with
        Python slicing and makes downstream calculations straightforward.
        """
        if trim_mode == "frames":
            resolved_start = int(start_frame)
            resolved_end = int(end_frame) if end_frame > 0 else total_frames
        elif trim_mode == "time":
            resolved_start = int(_safe_float(start_time) * fps)
            resolved_end = int(_safe_float(end_time) * fps) if end_time > 0 else total_frames
        else:
            raise ValueError(f"Unsupported trim_mode '{trim_mode}'. Expected 'time' or 'frames'.")

        if resolved_start < 0:
            raise ValueError("Start position must be greater than or equal to 0.")

        if resolved_end < 0:
            raise ValueError("End position must be greater than or equal to 0.")

        resolved_start = min(resolved_start, total_frames - 1)
        resolved_end = min(resolved_end, total_frames)
        resolved_end = max(resolved_end, resolved_start + 1)

        if resolved_start >= total_frames:
            raise ValueError(
                f"Start position {resolved_start} is outside the video. Total frames: {total_frames}."
            )

        if resolved_start >= resolved_end:
            raise ValueError(
                f"Invalid trim range: start frame {resolved_start} must be smaller than end frame {resolved_end}."
            )

        return resolved_start, resolved_end


NODE_CLASS_MAPPINGS = {
    "VideoTrimmer": VideoTrimmer,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "VideoTrimmer": "Video Trimmer",
}