import cython
from cython.cimports.av.codec.context import CodecContext, wrap_codec_context
from cython.cimports.av.container.streams import StreamContainer
from cython.cimports.av.dictionary import Dictionary
from cython.cimports.av.error import err_check
from cython.cimports.av.packet import Packet
from cython.cimports.av.stream import Stream, wrap_stream
from cython.cimports.av.utils import avdict_to_dict
from cython.cimports.libc.stdint import int64_t, uint8_t
from cython.cimports.libc.stdlib import free, malloc


@cython.cfunc
def close_input(self: InputContainer):
    self.streams = StreamContainer()
    with cython.nogil:
        if self._myflag & 2:
            # This causes `self.ptr` to be set to NULL.
            lib.avformat_close_input(cython.address(self.ptr))
            self._myflag &= ~2  # enum.input_was_opened = False


@cython.final
@cython.cclass
class InputContainer(Container):
    def __cinit__(self, *args, **kwargs):
        py_codec_context: CodecContext
        i: cython.uint
        stream: cython.pointer[lib.AVStream]
        codec: cython.pointer[cython.const[lib.AVCodec]]
        codec_context: cython.pointer[lib.AVCodecContext]

        # If we have either the global `options`, or a `stream_options`, prepare
        # a mashup of those options for each stream.
        c_options: cython.pointer[cython.pointer[lib.AVDictionary]] = cython.NULL
        base_dict: Dictionary
        stream_dict: Dictionary
        nb_streams_before: cython.uint = self.ptr.nb_streams
        if self.stream_options and nb_streams_before == 0:
            raise ValueError(
                "stream_options were provided, but this format does not expose "
                "its streams before avformat_find_stream_info (e.g. MPEG). "
                "Per-stream options cannot be applied."
            )
        # Only allocate c_options when streams are already known.
        if (self.options or self.stream_options) and nb_streams_before > 0:
            base_dict = Dictionary(self.options)
            c_options = cython.cast(
                cython.pointer[cython.pointer[lib.AVDictionary]],
                malloc(nb_streams_before * cython.sizeof(cython.p_void)),
            )
            for i in range(nb_streams_before):
                c_options[i] = cython.NULL
                if i < len(self.stream_options) and self.stream_options:
                    stream_dict = base_dict.copy()
                    stream_dict.update(self.stream_options[i])
                    lib.av_dict_copy(cython.address(c_options[i]), stream_dict.ptr, 0)
                else:
                    lib.av_dict_copy(cython.address(c_options[i]), base_dict.ptr, 0)

        self.set_timeout(self.open_timeout)
        self.start_timeout()
        with cython.nogil:
            ret = lib.avformat_find_stream_info(self.ptr, c_options)
        self.set_timeout(None)
        self.err_check(ret)

        if c_options:
            for i in range(nb_streams_before):
                lib.av_dict_free(cython.address(c_options[i]))
            free(c_options)

        at_least_one_accelerated_context = False

        self.streams = StreamContainer()
        for i in range(self.ptr.nb_streams):
            stream = self.ptr.streams[i]
            codec = lib.avcodec_find_decoder(stream.codecpar.codec_id)
            if codec:
                codec_context = lib.avcodec_alloc_context3(codec)
                err_check(
                    lib.avcodec_parameters_to_context(codec_context, stream.codecpar)
                )
                codec_context.pkt_timebase = stream.time_base
                py_codec_context = wrap_codec_context(
                    codec_context, codec, self.hwaccel
                )
                if py_codec_context.is_hwaccel:
                    at_least_one_accelerated_context = True
            else:
                # no decoder is available
                py_codec_context = None
            self.streams.add_stream(wrap_stream(self, stream, py_codec_context))

        if (
            self.hwaccel
            and not self.hwaccel.allow_software_fallback
            and not at_least_one_accelerated_context
        ):
            raise RuntimeError(
                "Hardware accelerated decode requested but no stream is compatible"
            )

        self.metadata = avdict_to_dict(
            self.ptr.metadata, self.metadata_encoding, self.metadata_errors
        )

    def __dealloc__(self):
        close_input(self)

    @property
    def start_time(self):
        self._assert_open()
        if self.ptr.start_time != lib.AV_NOPTS_VALUE:
            return self.ptr.start_time

    @property
    def duration(self):
        self._assert_open()
        if self.ptr.duration != lib.AV_NOPTS_VALUE:
            return self.ptr.duration

    @property
    def bit_rate(self):
        self._assert_open()
        return self.ptr.bit_rate

    @property
    def size(self):
        self._assert_open()
        return lib.avio_size(self.ptr.pb)

    def close(self):
        close_input(self)

    def demux(self, *args, **kwargs):
        """demux(streams=None, video=None, audio=None, subtitles=None, data=None)

        Yields a series of :class:`.Packet` from the given set of :class:`.Stream`::

            for packet in container.demux():
                # Do something with `packet`, often:
                for frame in packet.decode():
                    # Do something with `frame`.

        .. seealso:: :meth:`.StreamContainer.get` for the interpretation of
            the arguments.

        .. note:: The last packets are dummy packets that when decoded will flush the buffers.

        """
        self._assert_open()

        streams: list[Stream] = self.streams.get(*args, **kwargs)
        if self.ptr.nb_streams == 0:
            return
        include_stream: cython.pointer[uint8_t] = cython.cast(
            cython.pointer[uint8_t],
            malloc(self.ptr.nb_streams * cython.sizeof(uint8_t)),
        )
        if include_stream == cython.NULL:
            raise MemoryError()

        i: cython.uint
        packet: Packet
        read_packet: cython.pointer[lib.AVPacket]
        ret: cython.int

        self.set_timeout(self.read_timeout)
        try:
            for i in range(self.ptr.nb_streams):
                include_stream[i] = 0
            for stream in streams:
                i = stream.index
                if i >= self.ptr.nb_streams:
                    raise ValueError(f"stream index {i} out of range")
                include_stream[i] = 1

            # Pre-allocate a AVPacket that is reused as the read buffer.
            with cython.nogil:
                read_packet = lib.av_packet_alloc()
            if read_packet == cython.NULL:
                raise MemoryError("Could not allocate packet")

            while True:
                # Reset the read buffer
                with cython.nogil:
                    lib.av_packet_unref(read_packet)
                try:
                    self.start_timeout()
                    with cython.nogil:
                        ret = lib.av_read_frame(self.ptr, read_packet)
                    self.err_check(ret)
                except EOFError:
                    break

                if include_stream[read_packet.stream_index]:
                    # If AVFMTCTX_NOHEADER is set in ctx_flags, then new streams
                    # may also appear in av_read_frame().
                    # http://ffmpeg.org/doxygen/trunk/structAVFormatContext.html
                    # TODO: find better way to handle this
                    if read_packet.stream_index < len(self.streams):
                        # Move the encoded data out of the read buffer into a
                        # fresh Packet for the caller.
                        packet = Packet()
                        with cython.nogil:
                            lib.av_packet_move_ref(packet.ptr, read_packet)
                        packet._stream = self.streams[packet.ptr.stream_index]
                        # Keep track of this so that remuxing is easier.
                        packet.ptr.time_base = packet._stream.ptr.time_base
                        yield packet

            # Flush!
            for i in range(self.ptr.nb_streams):
                if include_stream[i]:
                    packet = Packet()
                    packet._stream = self.streams[i]
                    packet.ptr.time_base = packet._stream.ptr.time_base
                    yield packet

        finally:
            self.set_timeout(None)
            free(include_stream)
            if read_packet != cython.NULL:
                lib.av_packet_free(cython.address(read_packet))

    def decode(self, *args, **kwargs):
        """decode(streams=None, video=None, audio=None, subtitles=None, data=None)

        Yields a series of :class:`.Frame` from the given set of streams::

            for frame in container.decode():
                # Do something with `frame`.

        .. seealso:: :meth:`.StreamContainer.get` for the interpretation of
            the arguments.

        """
        self._assert_open()
        for packet in self.demux(*args, **kwargs):
            for frame in packet.decode():
                yield frame

    def seek(
        self,
        offset,
        *,
        backward: cython.bint = True,
        any_frame: cython.bint = False,
        stream: Stream | None = None,
        unsupported_frame_offset: cython.bint = False,
        unsupported_byte_offset: cython.bint = False,
    ):
        """seek(offset, *, backward=True, any_frame=False, stream=None)

        Seek to a (key)frame nearest to the given timestamp.

        :param int offset: Time to seek to, expressed in``stream.time_base`` if ``stream``
            is given, otherwise in :data:`av.time_base`.
        :param bool backward: If there is not a (key)frame at the given offset,
            look backwards for it.
        :param bool any_frame: Seek to any frame, not just a keyframe.
        :param Stream stream: The stream who's ``time_base`` the ``offset`` is in.

        :param bool unsupported_frame_offset: ``offset`` is a frame
            index instead of a time; not supported by any known format.
        :param bool unsupported_byte_offset: ``offset`` is a byte
            location in the file; not supported by any known format.

        After seeking, packets that you demux should correspond (roughly) to
        the position you requested.

        In most cases, the defaults of ``backwards = True`` and ``any_frame = False``
        are the best course of action, followed by you demuxing/decoding to
        the position that you want. This is because to properly decode video frames
        you need to start from the previous keyframe.

        .. seealso:: :ffmpeg:`avformat_seek_file` for discussion of the flags.

        """
        self._assert_open()

        if not isinstance(offset, int):
            raise TypeError("Container.seek only accepts integer offset.", type(offset))

        c_offset: int64_t = offset
        flags: cython.int = 0
        ret: cython.int

        if backward:
            flags |= lib.AVSEEK_FLAG_BACKWARD
        if any_frame:
            flags |= lib.AVSEEK_FLAG_ANY

        # If someone really wants (and to experiment), expose these.
        if unsupported_frame_offset:
            flags |= lib.AVSEEK_FLAG_FRAME
        if unsupported_byte_offset:
            flags |= lib.AVSEEK_FLAG_BYTE

        stream_index: cython.int = stream.index if stream else -1
        with cython.nogil:
            ret = lib.av_seek_frame(self.ptr, stream_index, c_offset, flags)
        err_check(ret)

        self.flush_buffers()

    @cython.cfunc
    def flush_buffers(self):
        self._assert_open()

        stream: Stream
        codec_context: CodecContext

        for stream in self.streams:
            codec_context = stream.codec_context
            if codec_context:
                codec_context.flush_buffers()