#
#            PySceneDetect: Python-Based Video Scene Detector
#   -------------------------------------------------------------------
#     [  Site:    https://scenedetect.com                           ]
#     [  Docs:    https://scenedetect.com/docs/                     ]
#     [  Github:  https://github.com/Breakthrough/PySceneDetect/    ]
#
# Copyright (C) 2014-2024 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.
#
"""PySceneDetect scenedetect.stats_manager Tests

This file includes unit tests for the scenedetect.stats_manager module (specifically,
the StatsManager object, used to coordinate caching of frame metrics to/from a CSV
file to speed up subsequent calls to detect_scenes on a SceneManager.

These tests rely on the SceneManager, VifdeoStreamCv2, and ContentDetector classes.

These tests also require the testvideo.mp4 (see test_scene_manager.py for download
instructions), however any other valid video file can be used as well by modifying
the fixture test_video_file in conftest.py.

Additionally, these tests will create, write to, and read from files which use names
TEST_STATS_FILE_XXXXXXXXXXXX.csv, where the X's will be replaced with random digits.
These files will be deleted, if possible, after the tests are completed running.
"""

import csv
import os
import random
from pathlib import Path

import pytest

from scenedetect.backends.opencv import VideoStreamCv2
from scenedetect.detectors import ContentDetector
from scenedetect.frame_timecode import FrameTimecode
from scenedetect.scene_manager import SceneManager
from scenedetect.stats_manager import (
    COLUMN_NAME_FRAME_NUMBER,
    COLUMN_NAME_TIMECODE,
    StatsFileCorrupt,
    StatsManager,
)


def test_metrics():
    """Test StatsManager metric registration/setting/getting with a set of pre-defined
    key-value pairs (metric_dict).
    """
    metric_dict = {"some_metric": 1.2345, "another_metric": 6.7890}
    metric_keys = list(metric_dict.keys())

    stats = StatsManager()
    frame_key = 100
    assert not stats.is_save_required()

    stats.register_metrics(metric_keys)

    assert not stats.is_save_required()

    assert not stats.metrics_exist(frame_key, metric_keys)
    assert stats.get_metrics(frame_key, metric_keys) == [None] * len(metric_keys)

    stats.set_metrics(frame_key, metric_dict)

    assert stats.is_save_required()

    assert stats.metrics_exist(frame_key, metric_keys)
    assert stats.metrics_exist(frame_key, metric_keys[1:])

    assert stats.get_metrics(frame_key, metric_keys) == [
        metric_dict[metric_key] for metric_key in metric_keys
    ]


def test_detector_metrics(test_video_file):
    """Test passing StatsManager to a SceneManager and using it for storing the frame metrics
    from a ContentDetector.
    """
    video = VideoStreamCv2(test_video_file)
    stats_manager = StatsManager()
    scene_manager = SceneManager(stats_manager)
    scene_manager.add_detector(ContentDetector())
    video_fps = video.frame_rate
    duration = FrameTimecode("00:00:05", video_fps)
    scene_manager.auto_downscale = True
    scene_manager.detect_scenes(video=video, duration=duration)
    # Check that metrics were written to the StatsManager.
    assert stats_manager.get_metrics(0, ContentDetector.METRIC_KEYS)


def test_load_empty_stats(tmp_path: Path):
    """Test loading an empty stats file, ensuring it results in no errors."""
    path = tmp_path.joinpath("stats.csv")
    with open(path, "w"):
        pass
    stats_manager = StatsManager()
    stats_manager.load_from_csv(path)


def test_save_no_detect_scenes(tmp_path: Path):
    """Test saving without calling detect_scenes."""
    path = tmp_path.joinpath("stats.csv")
    stats_manager = StatsManager()
    stats_manager.save_to_csv(path)


def test_load_hardcoded_file(tmp_path: Path):
    """Test loading a stats file with some hard-coded data generated by this test case."""

    path = tmp_path.joinpath("stats.csv")
    stats_manager = StatsManager()
    with open(path, "w") as stats_file:
        stats_writer = csv.writer(stats_file, lineterminator="\n")

        some_metric_key = "some_metric"
        some_metric_value = 1.2
        some_frame_key = 100
        base_timecode = FrameTimecode(0, 29.97)
        some_frame_timecode = base_timecode + some_frame_key

        # Write out a valid file.
        stats_writer.writerow([COLUMN_NAME_FRAME_NUMBER, COLUMN_NAME_TIMECODE, some_metric_key])
        stats_writer.writerow(
            [some_frame_key + 1, some_frame_timecode.get_timecode(), str(some_metric_value)]
        )

    stats_manager.load_from_csv(path)

    # Check that we decoded the correct values.
    assert stats_manager.metrics_exist(some_frame_key, [some_metric_key])
    assert stats_manager.get_metrics(some_frame_key, [some_metric_key])[0] == pytest.approx(
        some_metric_value
    )


def test_save_load_from_video(test_video_file, tmp_path: Path):
    """Test generating and saving some frame metrics from TEST_VIDEO_FILE to a file on disk, and
    loading the file back to ensure the loaded frame metrics agree with those that were saved.
    """
    video = VideoStreamCv2(test_video_file)
    stats_manager = StatsManager()
    scene_manager = SceneManager(stats_manager)

    scene_manager.add_detector(ContentDetector())

    video_fps = video.frame_rate
    duration = FrameTimecode("00:00:05", video_fps)

    scene_manager.auto_downscale = True
    scene_manager.detect_scenes(video, duration=duration)

    path = tmp_path.joinpath("stats.csv")
    stats_manager.save_to_csv(csv_file=path)

    metrics = stats_manager.metric_keys

    stats_manager_new = StatsManager()

    stats_manager_new.load_from_csv(path)

    # Compare the first 5 frames. Frame 0 won't have any metrics for this detector.
    for frame in range(1, 5 + 1):
        assert stats_manager.metrics_exist(frame, metrics)
        orig_metrics = stats_manager.get_metrics(frame, metrics)
        new_metrics = stats_manager_new.get_metrics(frame, metrics)
        for i, metric_val in enumerate(orig_metrics):
            assert metric_val == pytest.approx(new_metrics[i])


def test_load_corrupt_stats(tmp_path: Path):
    """Test loading a corrupted stats file created by outputting data in the wrong format."""

    stats_manager = StatsManager()

    path = tmp_path.joinpath("stats.csv")
    with open(path, "w") as stats_file:
        stats_writer = csv.writer(stats_file, lineterminator="\n")

        some_metric_key = "some_metric"
        some_metric_value = str(1.2)
        some_frame_key = 100
        base_timecode = FrameTimecode(0, 29.97)
        some_frame_timecode = base_timecode + some_frame_key

        # Write out some invalid files.

        # File #0: Wrong Header Names [StatsFileCorrupt]
        # Swapped timecode & frame number.
        stats_writer.writerow([COLUMN_NAME_TIMECODE, COLUMN_NAME_FRAME_NUMBER, some_metric_key])
        stats_writer.writerow(
            [some_frame_key, some_frame_timecode.get_timecode(), some_metric_value]
        )

        stats_file.close()

        with pytest.raises(StatsFileCorrupt):
            stats_manager.load_from_csv(path)
