import cython
from cython.cimports.av.audio.frame import alloc_audio_frame
from cython.cimports.av.error import err_check


@cython.cclass
class AudioFifo:
    """A simple audio sample FIFO (First In First Out) buffer."""

    def __repr__(self):
        try:
            result = (
                f"<av.{self.__class__.__name__} {self.samples} samples of "
                f"{self.sample_rate}hz {self.layout} {self.format} at 0x{id(self):x}>"
            )
        except AttributeError:
            result = (
                f"<av.{self.__class__.__name__} uninitialized, use fifo.write(frame),"
                f" at 0x{id(self):x}>"
            )
        return result

    def __dealloc__(self):
        if self.ptr:
            lib.av_audio_fifo_free(self.ptr)

    @cython.ccall
    def write(self, frame: AudioFrame | None):
        """write(frame)

        Push a frame of samples into the queue.

        :param AudioFrame frame: The frame of samples to push.

        The FIFO will remember the attributes from the first frame, and use those
        to populate all output frames.

        If there is a :attr:`~.Frame.pts` and :attr:`~.Frame.time_base` and
        :attr:`~.AudioFrame.sample_rate`, then the FIFO will assert that the incoming
        timestamps are continuous.

        """

        if frame is None:
            raise TypeError("AudioFifo must be given an AudioFrame.")

        if not frame.ptr.nb_samples:
            return

        if not self.ptr:
            # Hold onto a copy of the attributes of the first frame to populate
            # output frames with.
            self.template = alloc_audio_frame()
            self.template._copy_internal_attributes(frame)
            self.template._init_user_attributes()

            # Figure out our "time_base".
            if frame._time_base.num and frame.ptr.sample_rate:
                self.pts_per_sample = frame._time_base.den / float(frame._time_base.num)
                self.pts_per_sample /= frame.ptr.sample_rate
            else:
                self.pts_per_sample = 0

            self.ptr = lib.av_audio_fifo_alloc(
                cython.cast(lib.AVSampleFormat, frame.ptr.format),
                frame.layout.nb_channels,
                frame.ptr.nb_samples
                * 2,  # Just a default number of samples; it will adjust.
            )

            if not self.ptr:
                raise RuntimeError("Could not allocate AVAudioFifo.")

        # Make sure nothing changed.
        elif (
            frame.ptr.format != self.template.ptr.format
            or frame.ptr.sample_rate != self.template.ptr.sample_rate
            or (
                frame._time_base.num
                and self.template._time_base.num
                and (
                    frame._time_base.num != self.template._time_base.num
                    or frame._time_base.den != self.template._time_base.den
                )
            )
        ):
            raise ValueError("Frame does not match AudioFifo parameters.")

        # Assert that the PTS are what we expect.
        expected_pts = cython.declare(int64_t)
        if self.pts_per_sample and frame.ptr.pts != lib.AV_NOPTS_VALUE:
            expected_pts = cython.cast(
                int64_t, self.pts_per_sample * self.samples_written
            )
            if frame.ptr.pts != expected_pts:
                raise ValueError(
                    "Frame.pts (%d) != expected (%d); fix or set to None."
                    % (frame.ptr.pts, expected_pts)
                )

        err_check(
            lib.av_audio_fifo_write(
                self.ptr,
                cython.cast(cython.pointer[cython.p_void], frame.ptr.extended_data),
                frame.ptr.nb_samples,
            )
        )

        self.samples_written += frame.ptr.nb_samples

    @cython.ccall
    def read(self, samples: cython.int = 0, partial: cython.bint = False):
        """read(samples=0, partial=False)

        Read samples from the queue.

        :param int samples: The number of samples to pull; 0 gets all.
        :param bool partial: Allow returning less than requested.
        :returns: New :class:`AudioFrame` or ``None`` (if empty).

        If the incoming frames had valid a :attr:`~.Frame.time_base`,
        :attr:`~.AudioFrame.sample_rate` and :attr:`~.Frame.pts`, the returned frames
        will have accurate timing.

        """

        if not self.ptr:
            return

        buffered_samples: cython.int = lib.av_audio_fifo_size(self.ptr)
        if buffered_samples < 1:
            return

        samples = samples or buffered_samples

        if buffered_samples < samples:
            if partial:
                samples = buffered_samples
            else:
                return

        frame: AudioFrame = alloc_audio_frame()
        frame._copy_internal_attributes(self.template)
        frame._init(
            cython.cast(lib.AVSampleFormat, self.template.ptr.format),
            cython.cast(lib.AVChannelLayout, self.template.ptr.ch_layout),
            samples,
            1,  # Align?
        )

        err_check(
            lib.av_audio_fifo_read(
                self.ptr,
                cython.cast(cython.pointer[cython.p_void], frame.ptr.extended_data),
                samples,
            )
        )

        if self.pts_per_sample:
            frame.ptr.pts = cython.cast(
                uint64_t, self.pts_per_sample * self.samples_read
            )
        else:
            frame.ptr.pts = lib.AV_NOPTS_VALUE

        self.samples_read += samples
        return frame

    @cython.ccall
    def read_many(self, samples: cython.int, partial: cython.bint = False):
        """read_many(samples, partial=False)

        Read as many frames as we can.

        :param int samples: How large for the frames to be.
        :param bool partial: If we should return a partial frame.
        :returns: A ``list`` of :class:`AudioFrame`.

        """

        frame: AudioFrame
        frames: list = []
        while True:
            frame = self.read(samples, partial=partial)
            if frame is not None:
                frames.append(frame)
            else:
                break

        return frames

    @property
    def format(self):
        """The :class:`.AudioFormat` of this FIFO."""
        if not self.ptr:
            raise AttributeError(
                f"'{__name__}.AudioFifo' object has no attribute 'format'"
            )
        return self.template.format

    @property
    def layout(self):
        """The :class:`.AudioLayout` of this FIFO."""
        if not self.ptr:
            raise AttributeError(
                f"'{__name__}.AudioFifo' object has no attribute 'layout'"
            )
        return self.template.layout

    @property
    def sample_rate(self):
        if not self.ptr:
            raise AttributeError(
                f"'{__name__}.AudioFifo' object has no attribute 'sample_rate'"
            )
        return self.template.sample_rate

    @property
    def samples(self):
        """Number of audio samples (per channel) in the buffer."""
        return lib.av_audio_fifo_size(self.ptr) if self.ptr else 0