import functools
import os
import pathlib
from fractions import Fraction
from typing import cast

import numpy as np
import pytest

import av

from .common import TestCase, fate_suite


@functools.cache
def make_h264_test_video(path: str) -> None:
    """Generates a black H264 test video with two streams for testing hardware decoding."""

    # We generate a file here that's designed to be as compatible as possible with hardware
    # encoders. Hardware encoders are sometimes very picky and the errors we get are often
    # opaque, so there is nothing much we (PyAV) can do. The user needs to figure that out
    # if they want to use hwaccel. We only want to test the PyAV plumbing here.
    # Our video is H264, 1280x720p (note that some decoders have a minimum resolution limit), 24fps,
    # 8-bit yuv420p.
    pathlib.Path(path).parent.mkdir(parents=True, exist_ok=True)
    output_container = av.open(path, "w")

    streams = []
    for _ in range(2):
        stream = output_container.add_stream("libx264", rate=24)
        assert isinstance(stream, av.VideoStream)
        stream.width = 1280
        stream.height = 720
        stream.pix_fmt = "yuv420p"
        streams.append(stream)

    for _ in range(24):
        frame = av.VideoFrame.from_ndarray(
            np.zeros((720, 1280, 3), dtype=np.uint8), format="rgb24"
        )
        for stream in streams:
            for packet in stream.encode(frame):
                output_container.mux(packet)

    for stream in streams:
        for packet in stream.encode():
            output_container.mux(packet)

    output_container.close()


class TestDecode(TestCase):
    def test_decoded_video_frame_count(self) -> None:
        container = av.open(fate_suite("h264/interlaced_crop.mp4"))
        video_stream = next(s for s in container.streams if s.type == "video")

        assert video_stream is container.streams.video[0]

        frame_count = 0
        for frame in container.decode(video_stream):
            frame_count += 1

        assert frame_count == video_stream.frames

    def test_decode_audio_corrupt(self) -> None:
        # write an empty file
        path = self.sandboxed("empty.flac")
        with open(path, "wb"):
            pass

        packet_count = 0
        frame_count = 0

        with av.open(path) as container:
            for packet in container.demux(audio=0):
                for frame in packet.decode():
                    frame_count += 1
                packet_count += 1

        assert packet_count == 1
        assert frame_count == 0

    def test_decode_audio_sample_count(self) -> None:
        container = av.open(fate_suite("audio-reference/chorusnoise_2ch_44kHz_s16.wav"))
        audio_stream = next(s for s in container.streams if s.type == "audio")

        assert audio_stream is container.streams.audio[0]
        assert isinstance(audio_stream, av.AudioStream)

        sample_count = 0

        for frame in container.decode(audio_stream):
            sample_count += frame.samples

        assert audio_stream.duration is not None
        assert audio_stream.time_base is not None
        total_samples = (
            audio_stream.duration
            * audio_stream.sample_rate.numerator
            / audio_stream.time_base.denominator
        )
        assert sample_count == total_samples

    def test_decoded_time_base(self) -> None:
        container = av.open(fate_suite("h264/interlaced_crop.mp4"))
        stream = container.streams.video[0]

        assert stream.time_base == Fraction(1, 25)

        for packet in container.demux(stream):
            for frame in packet.decode():
                assert not isinstance(frame, av.subtitles.subtitle.SubtitleSet)
                assert packet.time_base == frame.time_base
                assert stream.time_base == frame.time_base
                return

    def test_decoded_motion_vectors(self) -> None:
        container = av.open(fate_suite("h264/interlaced_crop.mp4"))
        stream = container.streams.video[0]
        codec_context = stream.codec_context
        codec_context.options = {"flags2": "+export_mvs"}

        for frame in container.decode(stream):
            vectors = frame.side_data.get("MOTION_VECTORS")
            if frame.key_frame:
                # Key frame don't have motion vectors
                assert vectors is None
            else:
                assert vectors is not None and len(vectors) > 0
                return

    def test_decoded_motion_vectors_no_flag(self) -> None:
        container = av.open(fate_suite("h264/interlaced_crop.mp4"))
        stream = container.streams.video[0]

        for frame in container.decode(stream):
            vectors = frame.side_data.get("MOTION_VECTORS")
            if not frame.key_frame:
                assert vectors is None
                return

    def test_decoded_video_enc_params(self) -> None:
        container = av.open(fate_suite("h264/interlaced_crop.mp4"))
        stream = container.streams.video[0]
        stream.codec_context.options = {"export_side_data": "venc_params"}

        for frame in container.decode(stream):
            video_enc_params = cast(
                av.sidedata.encparams.VideoEncParams,
                frame.side_data.get("VIDEO_ENC_PARAMS"),
            )
            assert video_enc_params is not None
            assert video_enc_params.nb_blocks == 40 * 24

            first_block = video_enc_params.block_params(0)
            assert video_enc_params.qp + first_block.delta_qp == 29
            return

    def test_decoded_video_enc_params_no_flag(self) -> None:
        container = av.open(fate_suite("h264/interlaced_crop.mp4"))
        stream = container.streams.video[0]
        # When no additional flag is given, there should be no side data with the video encoding params

        for frame in container.decode(stream):
            video_enc_params = frame.side_data.get("VIDEO_ENC_PARAMS")
            assert video_enc_params is None

    def test_decode_video_corrupt(self) -> None:
        # write an empty file
        path = self.sandboxed("empty.h264")
        with open(path, "wb"):
            pass

        packet_count = 0
        frame_count = 0

        with av.open(path) as container:
            for packet in container.demux(video=0):
                for frame in packet.decode():
                    frame_count += 1
                packet_count += 1

        assert packet_count == 1
        assert frame_count == 0

    def test_decode_close_then_use(self) -> None:
        container = av.open(fate_suite("h264/interlaced_crop.mp4"))
        container.close()

        # Check accessing every attribute either works or raises
        # an `AssertionError`.
        for attr in dir(container):
            with self.subTest(attr=attr):
                try:
                    getattr(container, attr)
                except AssertionError:
                    pass

    def test_flush_decoded_video_frame_count(self) -> None:
        container = av.open(fate_suite("h264/interlaced_crop.mp4"))
        video_stream = container.streams.video[0]

        # Decode the first GOP, which requires a flush to get all frames
        have_keyframe = False
        input_count = 0
        output_count = 0

        for packet in container.demux(video_stream):
            if packet.is_keyframe:
                if have_keyframe:
                    break
                have_keyframe = True

            input_count += 1

            for frame in video_stream.decode(packet):
                output_count += 1

        # Check the test works as expected and requires a flush
        assert output_count < input_count

        for frame in video_stream.decode(None):
            # The Frame._time_base is not set by PyAV
            assert frame.time_base is None
            output_count += 1

        assert output_count == input_count

    def test_no_side_data(self) -> None:
        container = av.open(fate_suite("h264/interlaced_crop.mp4"))
        frame = next(container.decode(video=0))
        assert frame.rotation == 0

    def test_side_data(self) -> None:
        container = av.open(fate_suite("mov/displaymatrix.mov"))
        frame = next(container.decode(video=0))
        assert frame.rotation == -90

    def test_hardware_decode(self) -> None:
        hwdevices_available = av.codec.hwaccel.hwdevices_available()
        if "HWACCEL_DEVICE_TYPE" not in os.environ:
            pytest.skip(
                "Set the HWACCEL_DEVICE_TYPE to run this test. "
                f"Options are {' '.join(hwdevices_available)}"
            )

        HWACCEL_DEVICE_TYPE = os.environ["HWACCEL_DEVICE_TYPE"]
        assert HWACCEL_DEVICE_TYPE in hwdevices_available, (
            f"{HWACCEL_DEVICE_TYPE} not available"
        )

        test_video_path = "tests/assets/black.mp4"
        make_h264_test_video(test_video_path)

        hwaccel = av.codec.hwaccel.HWAccel(
            device_type=HWACCEL_DEVICE_TYPE, allow_software_fallback=False
        )

        container = av.open(test_video_path, hwaccel=hwaccel)
        video_stream = container.streams.video[0]
        assert video_stream.codec_context.is_hwaccel

        frame_count = 0
        for frame in container.decode(video_stream):
            frame_count += 1

        assert frame_count == video_stream.frames