codeccontext.pyx

cimport libav as lib

from av.audio.format cimport AudioFormat, get_audio_format
from av.audio.layout cimport AudioLayout, get_audio_layout
from av.audio.frame cimport AudioFrame, alloc_audio_frame
from av.frame cimport Frame
from av.packet cimport Packet
from av.utils cimport err_check


cdef class AudioCodecContext(CodecContext):


    cdef _init(self, lib.AVCodecContext *ptr, lib.AVCodec *codec):
        CodecContext._init(self, ptr, codec)

        # Sometimes there isn't a layout set, but there are a number of
        # channels. Assume it is the default layout.
        # TODO: Put this behind `not bare_metal`.
        # TODO: Do this more efficiently.
        if self.ptr.channels and not self.ptr.channel_layout:
            self.ptr.channel_layout = get_audio_layout(self.ptr.channels, 0).layout

    cdef _encode(self, Frame input_frame):
        """Encodes a frame of audio, returns a packet if one is ready.
        The output packet does not necessarily contain data for the most recent frame, 
        as encoders can delay, split, and combine input frames internally as needed.
        If called with with no args it will flush out the encoder and return the buffered
        packets until there are none left, at which it will return None.
        """

        cdef bint is_flushing = input_frame is None
        cdef AudioFrame frame = input_frame

        # Resample. A None frame will flush the resampler, and then the fifo (if used).
        if not self.resampler:
            self.resampler = AudioResampler(
                self.format,
                self.layout,
                self.ptr.sample_rate
            )
        frame = self.resampler.resample(frame)

        cdef bint use_fifo = not (self.ptr.codec.capabilities & lib.CODEC_CAP_VARIABLE_FRAME_SIZE)
        if use_fifo:
            print 'USING FIFO; frame_size:', self.ptr.frame_size
            if not self.fifo:
                self.fifo = AudioFifo()
            if frame:
                self.fifo.write(frame)

            # Pull partial frames if we were requested to flush (via a None frame).
            frame = self.fifo.read(self.ptr.frame_size, partial=is_flushing)

        cdef Packet packet = Packet()
        cdef int got_packet = 0

        if frame is not None:

            # TODO: Centralize time handling.

            # If the frame has a valid pts, scale it to the codec's time_base.
            # Remember that the AudioFifo time_base is always 1/sample_rate!
            if frame.ptr.pts != lib.AV_NOPTS_VALUE:
                frame.ptr.pts = lib.av_rescale_q(
                    frame.ptr.pts, 
                    frame._time_base,
                    self.ptr.time_base
                )
            else:
                frame.ptr.pts = lib.av_rescale(
                    self.ptr.frame_number,
                    self.ptr.sample_rate,
                    self.ptr.frame_size,
                )
                
        # TODO codec-ctx: streams rebased pts/dts/duration from self.ptr.time_base to self._stream.time_base
        
        err_check(lib.avcodec_encode_audio2(
            self.ptr,
            &packet.struct,
            frame.ptr if frame is not None else NULL,
            &got_packet,
        ))

        if got_packet:
            return packet
        

    cdef _decode_one(self, lib.AVPacket *packet, int *data_consumed):

        if not self.next_frame:
            self.next_frame = alloc_audio_frame()

        cdef int completed_frame = 0
        data_consumed[0] = err_check(lib.avcodec_decode_audio4(self.ptr, self.next_frame.ptr, &completed_frame, packet))
        if not completed_frame:
            return
        
        cdef AudioFrame frame = self.next_frame
        self.next_frame = None
        
        frame._init_properties()

        return frame

    property frame_size:
        """Number of samples per channel in an audio frame."""
        def __get__(self): return self.ptr.frame_size
    
    property sample_rate:
        """Number samples of per second."""
        def __get__(self):
            return self.ptr.sample_rate
        def __set__(self, int value):
            self.ptr.sample_rate = value

    # TODO: Deprecate.
    property rate:
        """Number samples of per second."""
        def __get__(self):
            return self.ptr.sample_rate
        def __set__(self, int value):
            self.ptr.sample_rate = value

    # TODO: Integrate into AudioLayout.
    property channels:
        def __get__(self):
            return self.ptr.channels
        def __set__(self, value):
            self.ptr.channels = value
            self.ptr.channel_layout = lib.av_get_default_channel_layout(value)
    property channel_layout:
        def __get__(self):
            return self.ptr.channel_layout

    property layout:
        def __get__(self):
            return get_audio_layout(self.ptr.channels, self.ptr.channel_layout)
        def __set__(self, value):
            cdef AudioLayout layout = AudioLayout(value)
            print 'SETTING LAYOUT TO', layout
            self.ptr.channel_layout = layout.layout
            self.ptr.channels = layout.nb_channels

    property format:
        def __get__(self):
            return get_audio_format(self.ptr.sample_fmt)
        def __set__(self, value):
            cdef AudioFormat format = AudioFormat(value)
            self.ptr.sample_fmt = format.sample_fmt