from av.audio.format cimport get_audio_format
from av.descriptor cimport wrap_avclass
from av.enum cimport define_enum
from av.utils cimport avrational_to_fraction, flag_in_bitfield
from av.video.format cimport get_video_format


cdef object _cinit_sentinel = object()


cdef Codec wrap_codec(const lib.AVCodec *ptr):
    cdef Codec codec = Codec(_cinit_sentinel)
    codec.ptr = ptr
    codec.is_encoder = lib.av_codec_is_encoder(ptr)
    codec._init()
    return codec


Properties = define_enum('Properties', 'av.codec', (
    ('NONE', 0),
    ('INTRA_ONLY', lib.AV_CODEC_PROP_INTRA_ONLY,
        """Codec uses only intra compression.
        Video and audio codecs only."""),
    ('LOSSY', lib.AV_CODEC_PROP_LOSSY,
        """Codec supports lossy compression. Audio and video codecs only.

        Note: A codec may support both lossy and lossless
        compression modes."""),
    ('LOSSLESS', lib.AV_CODEC_PROP_LOSSLESS,
        """Codec supports lossless compression. Audio and video codecs only."""),
    ('REORDER', lib.AV_CODEC_PROP_REORDER,
        """Codec supports frame reordering. That is, the coded order (the order in which
        the encoded packets are output by the encoders / stored / input to the
        decoders) may be different from the presentation order of the corresponding
        frames.

        For codecs that do not have this property set, PTS and DTS should always be
        equal."""),
    ('BITMAP_SUB', lib.AV_CODEC_PROP_BITMAP_SUB,
        """Subtitle codec is bitmap based
        Decoded AVSubtitle data can be read from the AVSubtitleRect->pict field."""),
    ('TEXT_SUB', lib.AV_CODEC_PROP_TEXT_SUB,
        """Subtitle codec is text based.
        Decoded AVSubtitle data can be read from the AVSubtitleRect->ass field."""),
), is_flags=True)

Capabilities = define_enum('Capabilities', 'av.codec', (
    ('NONE', 0),
    ('DRAW_HORIZ_BAND', lib.AV_CODEC_CAP_DRAW_HORIZ_BAND,
        """Decoder can use draw_horiz_band callback."""),
    ('DR1', lib.AV_CODEC_CAP_DR1,
        """Codec uses get_buffer() for allocating buffers and supports custom allocators.
        If not set, it might not use get_buffer() at all or use operations that
        assume the buffer was allocated by avcodec_default_get_buffer."""),
    ('HWACCEL', 1 << 4),
    ('DELAY', lib.AV_CODEC_CAP_DELAY,
        """Encoder or decoder requires flushing with NULL input at the end in order to
        give the complete and correct output.

        NOTE: If this flag is not set, the codec is guaranteed to never be fed with
              with NULL data. The user can still send NULL data to the public encode
              or decode function, but libavcodec will not pass it along to the codec
              unless this flag is set.

        Decoders:
        The decoder has a non-zero delay and needs to be fed with avpkt->data=NULL,
        avpkt->size=0 at the end to get the delayed data until the decoder no longer
        returns frames.

        Encoders:
        The encoder needs to be fed with NULL data at the end of encoding until the
        encoder no longer returns data.

        NOTE: For encoders implementing the AVCodec.encode2() function, setting this
              flag also means that the encoder must set the pts and duration for
              each output packet. If this flag is not set, the pts and duration will
              be determined by libavcodec from the input frame."""),
    ('SMALL_LAST_FRAME', lib.AV_CODEC_CAP_SMALL_LAST_FRAME,
        """Codec can be fed a final frame with a smaller size.
        This can be used to prevent truncation of the last audio samples."""),
    ('HWACCEL_VDPAU', 1 << 7),
    ('SUBFRAMES', lib.AV_CODEC_CAP_SUBFRAMES,
        """Codec can output multiple frames per AVPacket
        Normally demuxers return one frame at a time, demuxers which do not do
        are connected to a parser to split what they return into proper frames.
        This flag is reserved to the very rare category of codecs which have a
        bitstream that cannot be split into frames without timeconsuming
        operations like full decoding. Demuxers carrying such bitstreams thus
        may return multiple frames in a packet. This has many disadvantages like
        prohibiting stream copy in many cases thus it should only be considered
        as a last resort."""),
    ('EXPERIMENTAL', lib.AV_CODEC_CAP_EXPERIMENTAL,
        """Codec is experimental and is thus avoided in favor of non experimental
        encoders"""),
    ('CHANNEL_CONF', lib.AV_CODEC_CAP_CHANNEL_CONF,
        """Codec should fill in channel configuration and samplerate instead of container"""),
    ('NEG_LINESIZES', 1 << 11),
    ('FRAME_THREADS', lib.AV_CODEC_CAP_FRAME_THREADS,
        """Codec supports frame-level multithreading""",),
    ('SLICE_THREADS', lib.AV_CODEC_CAP_SLICE_THREADS,
        """Codec supports slice-based (or partition-based) multithreading."""),
    ('PARAM_CHANGE', lib.AV_CODEC_CAP_PARAM_CHANGE,
        """Codec supports changed parameters at any point."""),
    ('OTHER_THREADS', lib.AV_CODEC_CAP_OTHER_THREADS,
        """Codec supports avctx->thread_count == 0 (auto)."""),
    ('VARIABLE_FRAME_SIZE', lib.AV_CODEC_CAP_VARIABLE_FRAME_SIZE,
        """Audio encoder supports receiving a different number of samples in each call."""),
    ('AVOID_PROBING', lib.AV_CODEC_CAP_AVOID_PROBING,
        """Decoder is not a preferred choice for probing.
        This indicates that the decoder is not a good choice for probing.
        It could for example be an expensive to spin up hardware decoder,
        or it could simply not provide a lot of useful information about
        the stream.
        A decoder marked with this flag should only be used as last resort
        choice for probing."""),
    ('HARDWARE', lib.AV_CODEC_CAP_HARDWARE,
        """Codec is backed by a hardware implementation. Typically used to
        identify a non-hwaccel hardware decoder. For information about hwaccels, use
        avcodec_get_hw_config() instead."""),
    ('HYBRID', lib.AV_CODEC_CAP_HYBRID,
        """Codec is potentially backed by a hardware implementation, but not
        necessarily. This is used instead of AV_CODEC_CAP_HARDWARE, if the
        implementation provides some sort of internal fallback."""),
    ('ENCODER_REORDERED_OPAQUE', 1 << 20,  # lib.AV_CODEC_CAP_ENCODER_REORDERED_OPAQUE,  # FFmpeg 4.2
        """This codec takes the reordered_opaque field from input AVFrames
        and returns it in the corresponding field in AVCodecContext after
        encoding."""),
    ('ENCODER_FLUSH', 1 << 21,  # lib.AV_CODEC_CAP_ENCODER_FLUSH  # FFmpeg 4.3
        """This encoder can be flushed using avcodec_flush_buffers(). If this
        flag is not set, the encoder must be closed and reopened to ensure that
        no frames remain pending."""),
), is_flags=True)


class UnknownCodecError(ValueError):
    pass


cdef class Codec(object):

    """Codec(name, mode='r')

    :param str name: The codec name.
    :param str mode: ``'r'`` for decoding or ``'w'`` for encoding.

    This object exposes information about an available codec, and an avenue to
    create a :class:`.CodecContext` to encode/decode directly.

    ::

        >>> codec = Codec('mpeg4', 'r')
        >>> codec.name
        'mpeg4'
        >>> codec.type
        'video'
        >>> codec.is_encoder
        False

    """

    def __cinit__(self, name, mode='r'):

        if name is _cinit_sentinel:
            return

        if mode == 'w':
            self.ptr = lib.avcodec_find_encoder_by_name(name)
            if not self.ptr:
                self.desc = lib.avcodec_descriptor_get_by_name(name)
                if self.desc:
                    self.ptr = lib.avcodec_find_encoder(self.desc.id)

        elif mode == 'r':
            self.ptr = lib.avcodec_find_decoder_by_name(name)
            if not self.ptr:
                self.desc = lib.avcodec_descriptor_get_by_name(name)
                if self.desc:
                    self.ptr = lib.avcodec_find_decoder(self.desc.id)

        else:
            raise ValueError('Invalid mode; must be "r" or "w".', mode)

        self._init(name)

        # Sanity check.
        if (mode == 'w') != self.is_encoder:
            raise RuntimeError("Found codec does not match mode.", name, mode)

    cdef _init(self, name=None):

        if not self.ptr:
            raise UnknownCodecError(name)

        if not self.desc:
            self.desc = lib.avcodec_descriptor_get(self.ptr.id)
            if not self.desc:
                raise RuntimeError('No codec descriptor for %r.' % name)

        self.is_encoder = lib.av_codec_is_encoder(self.ptr)

        # Sanity check.
        if self.is_encoder and lib.av_codec_is_decoder(self.ptr):
            raise RuntimeError('%s is both encoder and decoder.')

    def create(self):
        """Create a :class:`.CodecContext` for this codec."""
        from .context import CodecContext
        return CodecContext.create(self)

    property is_decoder:
        def __get__(self):
            return not self.is_encoder

    property descriptor:
        def __get__(self): return wrap_avclass(self.ptr.priv_class)

    property name:
        def __get__(self): return self.ptr.name or ''
    property long_name:
        def __get__(self): return self.ptr.long_name or ''

    @property
    def type(self):
        """
        The media type of this codec.

        E.g: ``'audio'``, ``'video'``, ``'subtitle'``.

        """
        return lib.av_get_media_type_string(self.ptr.type)

    property id:
        def __get__(self): return self.ptr.id

    @property
    def frame_rates(self):
        """A list of supported frame rates (:class:`fractions.Fraction`), or ``None``."""
        if not self.ptr.supported_framerates:
            return

        ret = []
        cdef int i = 0
        while self.ptr.supported_framerates[i].denum:
            ret.append(avrational_to_fraction(&self.ptr.supported_framerates[i]))
            i += 1
        return ret

    @property
    def audio_rates(self):
        """A list of supported audio sample rates (``int``), or ``None``."""
        if not self.ptr.supported_samplerates:
            return

        ret = []
        cdef int i = 0
        while self.ptr.supported_samplerates[i]:
            ret.append(self.ptr.supported_samplerates[i])
            i += 1
        return ret

    @property
    def video_formats(self):
        """A list of supported :class:`.VideoFormat`, or ``None``."""
        if not self.ptr.pix_fmts:
            return

        ret = []
        cdef int i = 0
        while self.ptr.pix_fmts[i] != -1:
            ret.append(get_video_format(self.ptr.pix_fmts[i], 0, 0))
            i += 1
        return ret

    @property
    def audio_formats(self):
        """A list of supported :class:`.AudioFormat`, or ``None``."""
        if not self.ptr.sample_fmts:
            return

        ret = []
        cdef int i = 0
        while self.ptr.sample_fmts[i] != -1:
            ret.append(get_audio_format(self.ptr.sample_fmts[i]))
            i += 1
        return ret

    # NOTE: there are some overlaps, which we defer to how `ffmpeg -codecs`
    # handles them (by prefering the capablity to the property).
    # Also, LOSSLESS and LOSSY don't have to agree.

    @Properties.property
    def properties(self):
        """Flag property of :class:`.Properties`"""
        return self.desc.props

    intra_only = properties.flag_property('INTRA_ONLY')
    lossy = properties.flag_property('LOSSY')  # Defer to capability.
    lossless = properties.flag_property('LOSSLESS')  # Defer to capability.
    reorder = properties.flag_property('REORDER')
    bitmap_sub = properties.flag_property('BITMAP_SUB')
    text_sub = properties.flag_property('TEXT_SUB')

    @Capabilities.property
    def capabilities(self):
        """Flag property of :class:`.Capabilities`"""
        return self.ptr.capabilities

    draw_horiz_band = capabilities.flag_property('DRAW_HORIZ_BAND')
    dr1 = capabilities.flag_property('DR1')
    hwaccel = capabilities.flag_property('HWACCEL')
    delay = capabilities.flag_property('DELAY')
    small_last_frame = capabilities.flag_property('SMALL_LAST_FRAME')
    hwaccel_vdpau = capabilities.flag_property('HWACCEL_VDPAU')
    subframes = capabilities.flag_property('SUBFRAMES')
    experimental = capabilities.flag_property('EXPERIMENTAL')
    channel_conf = capabilities.flag_property('CHANNEL_CONF')
    neg_linesizes = capabilities.flag_property('NEG_LINESIZES')
    frame_threads = capabilities.flag_property('FRAME_THREADS')
    slice_threads = capabilities.flag_property('SLICE_THREADS')
    param_change = capabilities.flag_property('PARAM_CHANGE')
    other_threads = capabilities.flag_property('OTHER_THREADS')
    variable_frame_size = capabilities.flag_property('VARIABLE_FRAME_SIZE')
    avoid_probing = capabilities.flag_property('AVOID_PROBING')
    hardware = capabilities.flag_property('HARDWARE')
    hybrid = capabilities.flag_property('HYBRID')
    encoder_reordered_opaque = capabilities.flag_property('ENCODER_REORDERED_OPAQUE')
    encoder_flush = capabilities.flag_property('ENCODER_FLUSH')


cdef get_codec_names():
    names = set()
    cdef const lib.AVCodec *ptr
    cdef void *opaque = NULL
    while True:
        ptr = lib.av_codec_iterate(&opaque)
        if ptr:
            names.add(ptr.name)
        else:
            break
    return names

codecs_available = get_codec_names()


codec_descriptor = wrap_avclass(lib.avcodec_get_class())


def dump_codecs():
    """Print information about available codecs."""

    print '''Codecs:
 D..... = Decoding supported
 .E.... = Encoding supported
 ..V... = Video codec
 ..A... = Audio codec
 ..S... = Subtitle codec
 ...I.. = Intra frame-only codec
 ....L. = Lossy compression
 .....S = Lossless compression
 ------'''

    for name in sorted(codecs_available):

        try:
            e_codec = Codec(name, 'w')
        except ValueError:
            e_codec = None

        try:
            d_codec = Codec(name, 'r')
        except ValueError:
            d_codec = None

        # TODO: Assert these always have the same properties.
        codec = e_codec or d_codec

        try:
            print ' %s%s%s%s%s%s %-18s %s' % (
                '.D'[bool(d_codec)],
                '.E'[bool(e_codec)],
                codec.type[0].upper(),
                '.I'[codec.intra_only],
                '.L'[codec.lossy],
                '.S'[codec.lossless],
                codec.name,
                codec.long_name
            )
        except Exception as e:
            print '...... %-18s ERROR: %s' % (codec.name, e)
