#include "packet_ext.h"

void nn_packet_write_av_codec_parameters(struct nn_packet *packet, AVCodecParameters *codecpar)
{
    NNWT_PACKET_WRITE_TYPE(packet, enum AVMediaType, codecpar->codec_type);
    NNWT_PACKET_WRITE_TYPE(packet, enum AVCodecID, codecpar->codec_id);
    NNWT_PACKET_WRITE_TYPE(packet, u32, codecpar->codec_tag);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->extradata_size);
    NNWT_PACKET_WRITE_DATA(packet, codecpar->extradata, codecpar->extradata_size);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->format);
    NNWT_PACKET_WRITE_TYPE(packet, s64, codecpar->bit_rate);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->bits_per_coded_sample);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->bits_per_raw_sample);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->profile);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->level);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->width);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->height);
    NNWT_PACKET_WRITE_TYPE(packet, AVRational, codecpar->sample_aspect_ratio);
#if !defined CAMU_OLD_FFMPEG && LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(60, 11, 100)
    NNWT_PACKET_WRITE_TYPE(packet, AVRational, codecpar->framerate);
#endif
    NNWT_PACKET_WRITE_TYPE(packet, enum AVFieldOrder, codecpar->field_order);
    NNWT_PACKET_WRITE_TYPE(packet, enum AVColorRange, codecpar->color_range);
    NNWT_PACKET_WRITE_TYPE(packet, enum AVColorPrimaries, codecpar->color_primaries);
    NNWT_PACKET_WRITE_TYPE(packet, enum AVColorTransferCharacteristic, codecpar->color_trc);
    NNWT_PACKET_WRITE_TYPE(packet, enum AVColorSpace, codecpar->color_space);
    NNWT_PACKET_WRITE_TYPE(packet, enum AVChromaLocation, codecpar->chroma_location);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->video_delay);
#ifndef CAMU_OLD_FFMPEG
    // Incomplete.
    NNWT_PACKET_WRITE_TYPE(packet, enum AVChannelOrder, codecpar->ch_layout.order);
    if (codecpar->ch_layout.order == AV_CHANNEL_ORDER_NATIVE) {
        NNWT_PACKET_WRITE_TYPE(packet, u64, codecpar->ch_layout.u.mask);
    } else {
        NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->ch_layout.nb_channels);
    }
#else
    NNWT_PACKET_WRITE_TYPE(packet, u64, codecpar->channel_layout);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->channels);
#endif
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->sample_rate);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->block_align);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->frame_size);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->initial_padding);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->trailing_padding);
    NNWT_PACKET_WRITE_TYPE(packet, s32, codecpar->seek_preroll);
}

void nn_packet_write_av_codec_id(struct nn_packet *packet, enum AVCodecID codec_id)
{
    NNWT_PACKET_WRITE_TYPE(packet, enum AVCodecID, codec_id);
}

void nn_packet_write_av_dictionary(struct nn_packet *packet, AVDictionary *dict)
{
    s32 count = av_dict_count(dict);
    NNWT_PACKET_WRITE_TYPE(packet, s32, count);
    const AVDictionaryEntry *entry = NULL;
    while ((entry = av_dict_get(dict, "", entry, AV_DICT_IGNORE_SUFFIX))) {
        size_t len = al_strlen(entry->key);
        al_assert(len <= UINT32_MAX);
        NNWT_PACKET_WRITE_TYPE(packet, u32, len);
        NNWT_PACKET_WRITE_DATA(packet, entry->key, len);
        len = al_strlen(entry->value);
        al_assert(len <= UINT32_MAX);
        NNWT_PACKET_WRITE_TYPE(packet, u32, len);
        NNWT_PACKET_WRITE_DATA(packet, entry->value, len);
    }
}

void nn_packet_write_av_stream(struct nn_packet *packet, AVStream *stream)
{
    NNWT_PACKET_WRITE_TYPE(packet, s32, stream->index);
    nn_packet_write_av_codec_parameters(packet, stream->codecpar);
    NNWT_PACKET_WRITE_TYPE(packet, AVRational, stream->time_base);
    NNWT_PACKET_WRITE_TYPE(packet, s64, stream->duration);
    NNWT_PACKET_WRITE_TYPE(packet, s64, stream->start_time);
    NNWT_PACKET_WRITE_TYPE(packet, s64, stream->nb_frames);
    nn_packet_write_av_dictionary(packet, stream->metadata);
    if (stream->avg_frame_rate.den == 0) {
        // r_frame_rate.den = 0 handled on the client.
        NNWT_PACKET_WRITE_TYPE(packet, AVRational, stream->r_frame_rate);
    } else {
        NNWT_PACKET_WRITE_TYPE(packet, AVRational, stream->avg_frame_rate);
    }
}

void nn_packet_write_av_packet(struct nn_packet *packet, AVPacket *pkt)
{
    NNWT_PACKET_WRITE_TYPE(packet, s64, pkt->pts);
    NNWT_PACKET_WRITE_TYPE(packet, s64, pkt->dts);
    NNWT_PACKET_WRITE_TYPE(packet, s32, pkt->size);
    NNWT_PACKET_WRITE_DATA(packet, pkt->data, pkt->size);
    NNWT_PACKET_WRITE_TYPE(packet, s32, pkt->stream_index);
    NNWT_PACKET_WRITE_TYPE(packet, s32, pkt->flags);
    NNWT_PACKET_WRITE_TYPE(packet, s32, pkt->side_data_elems);
    for (s32 i = 0; i < pkt->side_data_elems; i++) {
        // Possibly upcast size_t.
        s64 size = (s64)pkt->side_data[i].size;
        NNWT_PACKET_WRITE_TYPE(packet, s64, size);
        NNWT_PACKET_WRITE_DATA(packet, pkt->side_data[i].data, pkt->side_data[i].size);
        NNWT_PACKET_WRITE_TYPE(packet, enum AVPacketSideDataType, pkt->side_data[i].type);
    }
    NNWT_PACKET_WRITE_TYPE(packet, s64, pkt->duration);
    NNWT_PACKET_WRITE_TYPE(packet, s64, pkt->pos);
#ifndef CAMU_OLD_FFMPEG
    NNWT_PACKET_WRITE_TYPE(packet, AVRational, pkt->time_base);
#endif
}

// @TODO: coded_side_data
//  https://github.com/FFmpeg/FFmpeg/blob/f62d8789116fef8d8cd7f20062b16ccd93f7ce23/doc/APIchanges#L537
//  https://ffmpeg.org/doxygen/8.0/structAVCodecParameters.html#ad54da9241deabb3601e6e0e8fa832c19
// #if LIBAVFORMAT_VERSION_INT >= AV_VERSION_INT(60, 30, 100)
void nn_packet_read_av_codec_parameters(struct nn_packet *packet, AVCodecParameters *codecpar)
{
    NNWT_PACKET_READ_TYPE(packet, enum AVMediaType, codecpar->codec_type);
    NNWT_PACKET_READ_TYPE(packet, enum AVCodecID, codecpar->codec_id);
    NNWT_PACKET_READ_TYPE(packet, u32, codecpar->codec_tag);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->extradata_size);
    codecpar->extradata = (u8 *)av_malloc(codecpar->extradata_size);
    u8 *extradata;
    NNWT_PACKET_READ_DATA(packet, codecpar->extradata_size, extradata);
    al_memcpy(codecpar->extradata, extradata, codecpar->extradata_size);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->format);
    NNWT_PACKET_READ_TYPE(packet, s64, codecpar->bit_rate);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->bits_per_coded_sample);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->bits_per_raw_sample);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->profile);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->level);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->width);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->height);
    NNWT_PACKET_READ_TYPE(packet, AVRational, codecpar->sample_aspect_ratio);
#if !defined CAMU_OLD_FFMPEG && LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(60, 11, 100)
    NNWT_PACKET_READ_TYPE(packet, AVRational, codecpar->framerate);
#endif
    NNWT_PACKET_READ_TYPE(packet, enum AVFieldOrder, codecpar->field_order);
    NNWT_PACKET_READ_TYPE(packet, enum AVColorRange, codecpar->color_range);
    NNWT_PACKET_READ_TYPE(packet, enum AVColorPrimaries, codecpar->color_primaries);
    NNWT_PACKET_READ_TYPE(packet, enum AVColorTransferCharacteristic, codecpar->color_trc);
    NNWT_PACKET_READ_TYPE(packet, enum AVColorSpace, codecpar->color_space);
    NNWT_PACKET_READ_TYPE(packet, enum AVChromaLocation, codecpar->chroma_location);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->video_delay);
#ifndef CAMU_OLD_FFMPEG
    enum AVChannelOrder order;
    NNWT_PACKET_READ_TYPE(packet, enum AVChannelOrder, order);
    if (order == AV_CHANNEL_ORDER_NATIVE) {
        u64 ch_layout_mask;
        NNWT_PACKET_READ_TYPE(packet, u64, ch_layout_mask);
        av_channel_layout_from_mask(&codecpar->ch_layout, ch_layout_mask);
    } else {
        s32 nb_channels;
        NNWT_PACKET_READ_TYPE(packet, s32, nb_channels);
        av_channel_layout_default(&codecpar->ch_layout, nb_channels);
    }
#else
    NNWT_PACKET_READ_TYPE(packet, u64, codecpar->channel_layout);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->channels);
#endif
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->sample_rate);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->block_align);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->frame_size);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->initial_padding);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->trailing_padding);
    NNWT_PACKET_READ_TYPE(packet, s32, codecpar->seek_preroll);
}

enum AVCodecID nn_packet_read_av_codec_id(struct nn_packet *packet)
{
    enum AVCodecID id;
    NNWT_PACKET_READ_TYPE(packet, enum AVCodecID, id);
    return id;
}

void nn_packet_read_av_dictionary(struct nn_packet *packet, AVDictionary **dict)
{
    s32 count;
    NNWT_PACKET_READ_TYPE(packet, s32, count);
    for (s32 i = 0; i < count; i++) {
        char *key, *value;
        u32 len;
        NNWT_PACKET_READ_TYPE(packet, u32, len);
        NNWT_PACKET_READ_DATA(packet, len, key);
        // Key and value MUST be allocated with av_malloc functions.
        key = av_strndup(key, len);
        NNWT_PACKET_READ_TYPE(packet, u32, len);
        NNWT_PACKET_READ_DATA(packet, len, value);
        value = av_strndup(value, len);
        av_dict_set(dict, key, value, AV_DICT_DONT_STRDUP_KEY | AV_DICT_DONT_STRDUP_VAL);
    }
}

AVStream *nn_packet_read_av_stream(AVFormatContext *format_context, const AVCodec *codec, struct nn_packet *packet)
{
    AVStream *stream = avformat_new_stream(format_context, codec);
    NNWT_PACKET_READ_TYPE(packet, s32, stream->index);
    nn_packet_read_av_codec_parameters(packet, stream->codecpar);
    NNWT_PACKET_READ_TYPE(packet, AVRational, stream->time_base);
    NNWT_PACKET_READ_TYPE(packet, s64, stream->duration);
    NNWT_PACKET_READ_TYPE(packet, s64, stream->start_time);
    NNWT_PACKET_READ_TYPE(packet, s64, stream->nb_frames);
    stream->metadata = NULL;
    nn_packet_read_av_dictionary(packet, &stream->metadata);
    NNWT_PACKET_READ_TYPE(packet, AVRational, stream->avg_frame_rate);
    return stream;
}

void nn_packet_read_av_packet(struct nn_packet *packet, AVPacket *pkt)
{
    NNWT_PACKET_READ_TYPE(packet, s64, pkt->pts);
    NNWT_PACKET_READ_TYPE(packet, s64, pkt->dts);
    NNWT_PACKET_READ_TYPE(packet, s32, pkt->size);
    NNWT_PACKET_READ_DATA(packet, pkt->size, pkt->data);
    NNWT_PACKET_READ_TYPE(packet, s32, pkt->stream_index);
    NNWT_PACKET_READ_TYPE(packet, s32, pkt->flags);
    s32 side_data_elems;
    NNWT_PACKET_READ_TYPE(packet, s32, side_data_elems);
    pkt->side_data_elems = 0;
    for (s32 i = 0; i < side_data_elems; i++) {
        u8 *data;
        enum AVPacketSideDataType type;
        s64 storage_for_size;
        NNWT_PACKET_READ_TYPE(packet, s64, storage_for_size);
        al_assert(sizeof(size_t) > 4 || storage_for_size <= INT32_MAX);
        size_t size = (size_t)storage_for_size;
        NNWT_PACKET_READ_DATA(packet, size, data);
        NNWT_PACKET_READ_TYPE(packet, enum AVPacketSideDataType, type);
        u8 *side_data = av_packet_new_side_data(pkt, type, size);
        al_memcpy(side_data, data, size);
    }
    NNWT_PACKET_READ_TYPE(packet, s64, pkt->duration);
    NNWT_PACKET_READ_TYPE(packet, s64, pkt->pos);
#ifndef CAMU_OLD_FFMPEG
    NNWT_PACKET_READ_TYPE(packet, AVRational, pkt->time_base);
#endif
}