#
#            PySceneDetect: Python-Based Video Scene Detector
#   -------------------------------------------------------------------
#     [  Site:    https://scenedetect.com                           ]
#     [  Docs:    https://scenedetect.com/docs/                     ]
#     [  Github:  https://github.com/Breakthrough/PySceneDetect/    ]
#
# Copyright (C) 2025 Brandon Castellano <http://www.bcastell.com>.
# PySceneDetect is licensed under the BSD 3-Clause License; see the
# included LICENSE file, or visit one of the above pages for details.
#
"""Implements :func:`save_images` functionality."""

import logging
import math
import queue
import sys
import threading
import typing as ty
from pathlib import Path
from string import Template

import cv2
import numpy as np

from scenedetect.common import (
    FrameTimecode,
    Interpolation,
    SceneList,
    TimecodeLike,
)
from scenedetect.platform import StrPath, get_and_create_path, get_cv2_imwrite_params, tqdm
from scenedetect.video_stream import VideoStream

logger = logging.getLogger("pyscenedetect")


def _generate_timecode_list(
    scene_list: SceneList,
    num_images: int,
    frame_margin: TimecodeLike,
) -> list[list[FrameTimecode]]:
    """Generate per-scene image timecodes using PTS-accurate seconds-based timing.

    `frame_margin` accepts any :data:`TimecodeLike` value (e.g. ``int`` frames, ``float``
    seconds, or ``str`` such as ``"0.1s"``).
    """
    frame_rate = scene_list[0][0].frame_rate
    assert frame_rate is not None
    margin_secs = FrameTimecode(timecode=frame_margin, fps=frame_rate).seconds
    result = []
    for start, end in scene_list:
        duration_secs = (end - start).seconds
        if duration_secs <= 0:
            result.append([start] * num_images)
            continue
        segment_secs = duration_secs / num_images
        timecodes = []
        for j in range(num_images):
            seg_start = start.seconds + j * segment_secs
            seg_end = start.seconds + (j + 1) * segment_secs
            if num_images == 1:
                t = start.seconds + duration_secs / 2.0
            elif j == 0:
                t = min(seg_start + margin_secs, seg_end)
            elif j == num_images - 1:
                t = max(seg_end - margin_secs, seg_start)
            else:
                t = (seg_start + seg_end) / 2.0
            timecodes.append(FrameTimecode(t, fps=frame_rate))
        result.append(timecodes)
    return result


def _scale_image(
    image: np.ndarray,
    aspect_ratio: float | None,
    height: int | None,
    width: int | None,
    scale: float | None,
    interpolation: Interpolation,
) -> np.ndarray:
    # TODO: Combine this resize with the ones below.
    if aspect_ratio is not None:
        image = cv2.resize(
            image, (0, 0), fx=aspect_ratio, fy=1.0, interpolation=interpolation.value
        )
    image_height = image.shape[0]
    image_width = image.shape[1]

    # Figure out what kind of resizing needs to be done
    if height or width:
        if height and not width:
            factor = height / float(image_height)
            width = int(factor * image_width)
        elif width and not height:
            factor = width / float(image_width)
            height = int(factor * image_height)
        assert height is not None
        assert width is not None
        assert height > 0 and width > 0
        image = cv2.resize(image, (width, height), interpolation=interpolation.value)
    elif scale:
        image = cv2.resize(image, (0, 0), fx=scale, fy=scale, interpolation=interpolation.value)
    return image


class _ImageExtractor:
    def __init__(
        self,
        num_images: int = 3,
        frame_margin: TimecodeLike = 1,
        image_extension: str = "jpg",
        imwrite_param: list[int] | None = None,
        image_name_template: str = "$VIDEO_NAME-Scene-$SCENE_NUMBER-$IMAGE_NUMBER",
        scale: float | None = None,
        height: int | None = None,
        width: int | None = None,
        interpolation: Interpolation = Interpolation.CUBIC,
    ):
        """Multi-threaded implementation of save-images functionality. Uses background threads to
        handle image encoding and saving images to disk to improve parallelism.

        This object is thread-safe.

        Arguments:
            num_images: Number of images to generate for each scene.  Minimum is 1.
            frame_margin: Padding around the beginning and end of each scene used when
                selecting which frames to extract. Accepts an int (frames), float (seconds),
                or str (e.g. ``"0.1s"``, ``"00:00:00.100"``). Can be 0, but some video files
                may then fail to extract the very last frame.
            image_extension: Type of image to save (must be one of 'jpg', 'png', or 'webp').
            encoder_param: Quality/compression efficiency, based on type of image:
                'jpg' / 'webp':  Quality 0-100, higher is better quality.  100 is lossless for webp.
                'png': Compression from 1-9, where 9 achieves best filesize but is slower to encode.
            image_name_template: Template to use for output filanames. Can use template variables
                $VIDEO_NAME, $SCENE_NUMBER, $IMAGE_NUMBER, $TIMECODE, $FRAME_NUMBER, $TIMESTAMP_MS.
                *NOTE*: Should not include the image extension (set `image_extension` instead).
            scale: Optional factor by which to rescale saved images. A scaling factor of 1 would
                not result in rescaling. A value < 1 results in a smaller saved image, while a
                value > 1 results in an image larger than the original. This value is ignored if
                either the height or width values are specified.
            height: Optional value for the height of the saved images. Specifying both the height
                and width will resize images to an exact size, regardless of aspect ratio.
                Specifying only height will rescale the image to that number of pixels in height
                while preserving the aspect ratio.
            width: Optional value for the width of the saved images. Specifying both the width
                and height will resize images to an exact size, regardless of aspect ratio.
                Specifying only width will rescale the image to that number of pixels wide
                while preserving the aspect ratio.
            interpolation: Type of interpolation to use when resizing images.
        """
        self._num_images = num_images
        self._frame_margin = frame_margin
        self._image_extension = image_extension
        self._image_name_template = image_name_template
        self._scale = scale
        self._height = height
        self._width = width
        self._interpolation = interpolation
        self._imwrite_param: list[int] = imwrite_param if imwrite_param is not None else []

    def run(
        self,
        video: VideoStream,
        scene_list: SceneList,
        output_dir: StrPath | None = None,
        show_progress=False,
    ) -> dict[int, list[str]]:
        """Run image extraction on `video` using the current parameters. Thread-safe.

        Arguments:
            video: The video to process.
            scene_list: The scenes detected in the video.
            output_dir: Directory to write files to.
            show_progress: If `true` and tqdm is available, shows a progress bar.
        """
        # Setup flags and init progress bar if available.
        completed = True
        logger.info(
            f"Saving {self._num_images} images per scene [format={self._image_extension}]"
            f" {output_dir if output_dir else ''} "
        )
        progress_bar = None
        if show_progress:
            progress_bar = tqdm(
                total=len(scene_list) * self._num_images, unit="images", dynamic_ncols=True
            )

        timecode_list = self.generate_timecode_list(scene_list)
        image_filenames = {i: [] for i in range(len(timecode_list))}

        filename_template = Template(self._image_name_template)
        logger.debug("Writing images with template %s", filename_template.template)
        scene_num_format = "%0"
        scene_num_format += str(max(3, math.floor(math.log(len(scene_list), 10)) + 1)) + "d"
        image_num_format = "%0"
        image_num_format += str(math.floor(math.log(self._num_images, 10)) + 2) + "d"

        def format_filename(scene_number: int, image_number: int, image_timecode: FrameTimecode):
            return "{}.{}".format(
                filename_template.safe_substitute(
                    VIDEO_NAME=video.name,
                    SCENE_NUMBER=scene_num_format % (scene_number + 1),
                    IMAGE_NUMBER=image_num_format % (image_number + 1),
                    FRAME_NUMBER=image_timecode.frame_num,
                    TIMESTAMP_MS=int(image_timecode.seconds * 1000),
                    TIMECODE=image_timecode.get_timecode().replace(":", ";"),
                ),
                self._image_extension,
            )

        MAX_QUEUED_ENCODE_FRAMES = 4
        MAX_QUEUED_SAVE_IMAGES = 4
        encode_queue = queue.Queue(MAX_QUEUED_ENCODE_FRAMES)
        save_queue = queue.Queue(MAX_QUEUED_SAVE_IMAGES)
        error_queue = queue.Queue(2)  # Queue size must be the same as the # of worker threads!

        def check_error_queue():
            try:
                return error_queue.get(block=False)
            except queue.Empty:
                pass
            return None

        def launch_thread(callable, *args, **kwargs):
            def capture_errors(callable, *args, **kwargs):
                try:
                    return callable(*args, **kwargs)
                # Errors we capture in `error_queue` will be re-raised by this thread.
                except:  # noqa: E722
                    error_queue.put(sys.exc_info())
                return None

            thread = threading.Thread(
                target=capture_errors,
                args=(
                    callable,
                    *args,
                ),
                kwargs=kwargs,
                daemon=True,
            )
            thread.start()
            return thread

        def checked_put(work_queue: queue.Queue, item: ty.Any):
            error = None
            while True:
                try:
                    work_queue.put(item, timeout=0.1)
                    return
                except queue.Full:
                    error = check_error_queue()
                    if error is not None:
                        break
                    continue
            raise error[1].with_traceback(error[2])

        encode_thread = launch_thread(
            self.image_encode_thread,
            video,
            encode_queue,
            save_queue,
        )
        save_thread = launch_thread(self.image_save_thread, save_queue, progress_bar)

        for i, scene_timecodes in enumerate(timecode_list):
            for j, timecode in enumerate(scene_timecodes):
                video.seek(timecode)
                frame_im = video.read()
                if frame_im is not None and frame_im is not False:
                    file_path = format_filename(i, j, timecode)
                    image_filenames[i].append(file_path)
                    checked_put(
                        encode_queue, (frame_im, get_and_create_path(file_path, output_dir))
                    )
                else:
                    completed = False
                    break

        checked_put(encode_queue, (None, None))
        encode_thread.join()
        checked_put(save_queue, (None, None))
        save_thread.join()

        error = check_error_queue()
        if error is not None:
            raise error[1].with_traceback(error[2])

        if progress_bar is not None:
            progress_bar.close()
        if not completed:
            logger.error("Could not generate all output images.")

        return image_filenames

    def image_encode_thread(
        self,
        video: VideoStream,
        encode_queue: queue.Queue,
        save_queue: queue.Queue,
    ):
        aspect_ratio = video.aspect_ratio
        if abs(aspect_ratio - 1.0) < 0.01:
            aspect_ratio = None
        # TODO: Validate that encoder_param is within the proper range.
        # Should be between 0 and 100 (inclusive) for jpg/webp, and 1-9 for png.
        while True:
            frame_im, dest_path = encode_queue.get()
            if frame_im is None:
                return
            frame_im = self.resize_image(
                frame_im,
                aspect_ratio,
            )
            (is_ok, encoded) = cv2.imencode(
                f".{self._image_extension}", frame_im, self._imwrite_param
            )
            if not is_ok:
                continue
            save_queue.put((encoded, dest_path))

    def image_save_thread(self, save_queue: queue.Queue, progress_bar: tqdm):
        while True:
            encoded, dest_path = save_queue.get()
            if encoded is None:
                return
            if encoded is not False:
                encoded.tofile(Path(dest_path))
            if progress_bar is not None:
                progress_bar.update(1)

    def generate_timecode_list(self, scene_list: SceneList) -> list[list[FrameTimecode]]:
        """Generates a list of timecodes for each scene in `scene_list` based on the current config
        parameters.

        Uses PTS-accurate seconds-based timing so results are correct for both CFR and VFR video.
        """
        return _generate_timecode_list(scene_list, self._num_images, self._frame_margin)

    def resize_image(
        self,
        image: np.ndarray,
        aspect_ratio: float | None,
    ) -> np.ndarray:
        return _scale_image(
            image, aspect_ratio, self._height, self._width, self._scale, self._interpolation
        )


def save_images(
    scene_list: SceneList,
    video: VideoStream,
    num_images: int = 3,
    frame_margin: TimecodeLike = 1,
    image_extension: str = "jpg",
    encoder_param: int = 95,
    image_name_template: str = "$VIDEO_NAME-Scene-$SCENE_NUMBER-$IMAGE_NUMBER",
    output_dir: StrPath | None = None,
    show_progress: bool | None = False,
    scale: float | None = None,
    height: int | None = None,
    width: int | None = None,
    interpolation: Interpolation = Interpolation.CUBIC,
    threading: bool = True,
) -> dict[int, list[str]]:
    """Save a set number of images from each scene, given a list of scenes
    and the associated video/frame source.

    Arguments:
        scene_list: A list of scenes (pairs of FrameTimecode objects) returned
            from calling a SceneManager's detect_scenes() method.
        video: A VideoStream object corresponding to the scene list.
            Note that the video will be closed/re-opened and seeked through.
        num_images: Number of images to generate for each scene.  Minimum is 1.
        frame_margin: Padding around the beginning and end of each scene used when
            selecting which frames to extract. Accepts an int (frames), float (seconds),
            or str (e.g. ``"0.1s"``, ``"00:00:00.100"``). Can be 0, but some video files
            may then fail to extract the very last frame.
        image_extension: Type of image to save (must be one of 'jpg', 'png', or 'webp').
        encoder_param: Quality/compression efficiency, based on type of image:
            'jpg' / 'webp':  Quality 0-100, higher is better quality.  100 is lossless for webp.
            'png': Compression from 1-9, where 9 achieves best filesize but is slower to encode.
        image_name_template: Template to use for naming image files. Can use the template variables
            $VIDEO_NAME, $SCENE_NUMBER, $IMAGE_NUMBER, $TIMECODE, $FRAME_NUMBER, $TIMESTAMP_MS.
            Should not include an extension.
        output_dir: Directory to output the images into.  If not set, the output
            is created in the working directory.
        show_progress: If True, shows a progress bar if tqdm is installed.
        scale: Optional factor by which to rescale saved images. A scaling factor of 1 would
            not result in rescaling. A value < 1 results in a smaller saved image, while a
            value > 1 results in an image larger than the original. This value is ignored if
            either the height or width values are specified.
        height: Optional value for the height of the saved images. Specifying both the height
            and width will resize images to an exact size, regardless of aspect ratio.
            Specifying only height will rescale the image to that number of pixels in height
            while preserving the aspect ratio.
        width: Optional value for the width of the saved images. Specifying both the width
            and height will resize images to an exact size, regardless of aspect ratio.
            Specifying only width will rescale the image to that number of pixels wide
            while preserving the aspect ratio.
        interpolation: Type of interpolation to use when resizing images.
        threading: Offload image encoding and disk IO to background threads to improve performance.

    Returns:
        Dictionary of the format { scene_num : [image_paths] }, where scene_num is the
        number of the scene in scene_list (starting from 1), and image_paths is a list of
        the paths to the newly saved/created images.

    Raises:
        ValueError: Raised if any arguments are invalid or out of range (e.g.
        if num_images is negative).
    """

    if not scene_list:
        return {}
    if num_images <= 0:
        raise ValueError("num_images must be greater than 0")
    if isinstance(frame_margin, (int, float)) and frame_margin < 0:
        raise ValueError("frame_margin must be non-negative")

    # TODO: Validate that encoder_param is within the proper range.
    # Should be between 0 and 100 (inclusive) for jpg/webp, and 1-9 for png.
    imwrite_param = (
        [get_cv2_imwrite_params()[image_extension], encoder_param]
        if encoder_param is not None
        else []
    )
    video.reset()

    if threading:
        extractor = _ImageExtractor(
            num_images,
            frame_margin,
            image_extension,
            imwrite_param,
            image_name_template,
            scale,
            height,
            width,
            interpolation,
        )
        return extractor.run(video, scene_list, output_dir, bool(show_progress))

    # Setup flags and init progress bar if available.
    completed = True
    logger.info(
        f"Saving {num_images} images per scene [format={image_extension}]"
        f" {output_dir if output_dir else ''} "
    )
    progress_bar = None
    if show_progress:
        progress_bar = tqdm(total=len(scene_list) * num_images, unit="images", dynamic_ncols=True)

    filename_template = Template(image_name_template)

    scene_num_format = "%0"
    scene_num_format += str(max(3, math.floor(math.log(len(scene_list), 10)) + 1)) + "d"
    image_num_format = "%0"
    image_num_format += str(math.floor(math.log(num_images, 10)) + 2) + "d"

    timecode_list = _generate_timecode_list(scene_list, num_images, frame_margin)

    image_filenames = {i: [] for i in range(len(timecode_list))}
    aspect_ratio = video.aspect_ratio
    if abs(aspect_ratio - 1.0) < 0.01:
        aspect_ratio = None

    logger.debug("Writing images with template %s", filename_template.template)
    for i, scene_timecodes in enumerate(timecode_list):
        for j, image_timecode in enumerate(scene_timecodes):
            video.seek(image_timecode)
            frame_im = video.read()
            if isinstance(frame_im, np.ndarray):
                # TODO: Add extension to template.
                # TODO: Allow NUM to be a valid suffix in addition to NUMBER.
                file_path = "{}.{}".format(
                    filename_template.safe_substitute(
                        VIDEO_NAME=video.name,
                        SCENE_NUMBER=scene_num_format % (i + 1),
                        IMAGE_NUMBER=image_num_format % (j + 1),
                        FRAME_NUMBER=image_timecode.frame_num,
                        TIMESTAMP_MS=int(image_timecode.seconds * 1000),
                        TIMECODE=image_timecode.get_timecode().replace(":", ";"),
                    ),
                    image_extension,
                )
                image_filenames[i].append(file_path)
                # TODO: Combine this resize with the ones below.
                if aspect_ratio is not None:
                    frame_im = cv2.resize(
                        frame_im, (0, 0), fx=aspect_ratio, fy=1.0, interpolation=interpolation.value
                    )
                frame_height = frame_im.shape[0]
                frame_width = frame_im.shape[1]

                # Figure out what kind of resizing needs to be done
                if height or width:
                    if height and not width:
                        factor = height / float(frame_height)
                        width = int(factor * frame_width)
                    elif width and not height:
                        factor = width / float(frame_width)
                        height = int(factor * frame_height)
                    assert height is not None
                    assert width is not None
                    assert height > 0 and width > 0
                    frame_im = cv2.resize(
                        frame_im, (width, height), interpolation=interpolation.value
                    )
                elif scale:
                    frame_im = cv2.resize(
                        frame_im, (0, 0), fx=scale, fy=scale, interpolation=interpolation.value
                    )
                path = Path(get_and_create_path(file_path, output_dir))
                (is_ok, encoded) = cv2.imencode(f".{image_extension}", frame_im, imwrite_param)
                if is_ok:
                    encoded.tofile(path)
                else:
                    logger.error(f"Failed to encode image for {file_path}")
            #
            else:
                completed = False
                break
            if progress_bar is not None:
                progress_bar.update(1)

    if progress_bar is not None:
        progress_bar.close()

    if not completed:
        logger.error("Could not generate all output images.")

    return image_filenames
