// https://aomediacodec.github.io/av1-rtp-spec/
// 7.1. Media Type Definition: video/av1

#include "rtp-packet.h"
#include "rtp-profile.h"
#include "rtp-payload-internal.h"
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <errno.h>

// Timestamp: The RTP timestamp indicates the time when the input frame was sampled, at a clock rate of 90 kHz
#define KHz 90 // 90000Hz

#define N_AV1_HEADER 1

#define OBU_SEQUENCE_HEADER        1
#define OBU_TEMPORAL_DELIMITER     2
#define OBU_FRAME_HEADER           3
#define OBU_TILE_GROUP             4
#define OBU_METADATA               5
#define OBU_FRAME                  6
#define OBU_REDUNDANT_FRAME_HEADER 7
#define OBU_TILE_LIST              8

#define AV1_AGGREGATION_HEADER_Z                                                                                       \
    0x80 // set to 1 if the first OBU element is an OBU fragment that is a continuation of an OBU fragment from the
         // previous packet, 0 otherwise.
#define AV1_AGGREGATION_HEADER_Y                                                                                       \
    0x40 // set to 1 if the last OBU element is an OBU fragment that will continue in the next packet, 0 otherwise.
#define AV1_AGGREGATION_HEADER_N                                                                                       \
    0x08 // set to 1 if the packet is the first packet of a coded video sequence, 0 otherwise. Note: if N equals 1 then
         // Z must equal 0.

struct rtp_encode_av1_t {
    struct rtp_packet_t pkt;
    struct rtp_payload_t handler;
    void *cbparam;
    int size;

    uint8_t *ptr;
    int offset;

    uint8_t aggregation;
};

static inline const uint8_t *leb128(const uint8_t *data, size_t bytes, uint64_t *size)
{
    size_t i;
    for (*size = i = 0; i * 7 < 64 && i < bytes;) {
        *size |= ((uint64_t)(data[i] & 0x7F)) << (i * 7);
        if (0 == (data[i++] & 0x80))
            break;
    }
    return data + i;
}

static inline uint8_t *leb128_write(int64_t size, uint8_t *data, size_t bytes)
{
    size_t i;
    for (i = 0; i * 7 < 64 && i < bytes;) {
        data[i] = (uint8_t)(size & 0x7F);
        size >>= 7;
        data[i++] |= size > 0 ? 0x80 : 0;
        if (0 == size)
            break;
    }
    return data + i;
}

static void *rtp_av1_pack_create(int size, uint8_t pt, uint16_t seq, uint32_t ssrc, struct rtp_payload_t *handler,
                                 void *cbparam)
{
    struct rtp_encode_av1_t *packer;
    packer = (struct rtp_encode_av1_t *)calloc(1, sizeof(*packer));
    if (!packer)
        return NULL;

    memcpy(&packer->handler, handler, sizeof(packer->handler));
    packer->cbparam = cbparam;
    packer->size = size;

    packer->pkt.rtp.v = RTP_VERSION;
    packer->pkt.rtp.pt = pt;
    packer->pkt.rtp.seq = seq;
    packer->pkt.rtp.ssrc = ssrc;
    return packer;
}

static void rtp_av1_pack_destroy(void *pack)
{
    struct rtp_encode_av1_t *packer;
    packer = (struct rtp_encode_av1_t *)pack;
#if defined(_DEBUG) || defined(DEBUG)
    memset(packer, 0xCC, sizeof(*packer));
#endif
    free(packer);
}

static void rtp_av1_pack_get_info(void *pack, uint16_t *seq, uint32_t *timestamp)
{
    struct rtp_encode_av1_t *packer;
    packer = (struct rtp_encode_av1_t *)pack;
    *seq = (uint16_t)packer->pkt.rtp.seq;
    *timestamp = packer->pkt.rtp.timestamp;
}

static int rtp_av1_pack_flush(struct rtp_encode_av1_t *packer, uint8_t aggregation)
{
    int r, n;
    if (!packer->ptr || packer->offset <= RTP_FIXED_HEADER)
        return 0; // nothing to send

    packer->ptr[RTP_FIXED_HEADER] = aggregation;
    packer->pkt.payloadlen = packer->offset - RTP_FIXED_HEADER;
    n = rtp_packet_serialize_header(&packer->pkt, packer->ptr, packer->size);
    if (n != RTP_FIXED_HEADER) {
        assert(0);
        return -1;
    }

    ++packer->pkt.rtp.seq;
    packer->pkt.rtp.m = 0; // clear marker bit
    packer->aggregation &= ~(AV1_AGGREGATION_HEADER_N | AV1_AGGREGATION_HEADER_Z);

    r = packer->handler.packet(packer->cbparam, packer->ptr, n + packer->pkt.payloadlen, packer->pkt.rtp.timestamp, 0);
    packer->handler.free(packer->cbparam, packer->ptr);
    packer->offset = 0;
    packer->ptr = NULL;
    return r;
}

static int rtp_av1_pack_obu(struct rtp_encode_av1_t *packer, const uint8_t *obu, int64_t bytes)
{
    int r;
    int64_t n;
    uint8_t *ptr, *end;

    while (bytes > 0) {
        if (NULL == packer->ptr) {
            packer->ptr = (uint8_t *)packer->handler.alloc(packer->cbparam, packer->size);
            if (!packer->ptr)
                return -ENOMEM;
            packer->offset = RTP_FIXED_HEADER + 1; // RTP Header + AV1 aggregation header
        }

        ptr = packer->ptr + packer->offset;
        end = packer->ptr + packer->size;

        // OBU element size
        assert(packer->size < 0x3FFF); // 14bits
        if (ptr + bytes + ((bytes > 0x7F) ? 2 : 1) > end)
            n = end - ptr - 2;
        else
            n = bytes;

        ptr = leb128_write(n, ptr, end - ptr);
        memcpy(ptr, obu, (size_t)n);
        ptr += n;
        obu += n;
        bytes -= n;
        packer->offset = (int)(ptr - packer->ptr);

        if (packer->size - packer->offset < 8) {
            r = rtp_av1_pack_flush(packer, packer->aggregation | (bytes > 0 ? AV1_AGGREGATION_HEADER_Y : 0));
            if (0 != r)
                return r;
        }

        if (bytes > 0)
            packer->aggregation |= AV1_AGGREGATION_HEADER_Z;
    }

    return 0;
}

/// https://aomediacodec.github.io/av1-spec/av1-spec.pd
/// Annex B: Length delimited bitstream format
/// @param[in] data temporal_unit
/// @param[in] bytes temporal_unit_sizetemporal_unit_size
static int rtp_av1_pack_input_annexb(void *pack, const void *data, int bytes, uint32_t timestamp)
{
    int r;
    //	uint8_t obu_has_size_field;
    uint8_t obu_extension_flag;
    uint8_t temporal_id, temporal_id0;
    uint8_t spatial_id, spatial_id0;
    uint8_t obu_type;
    uint64_t obu_size, frame_size;
    const uint8_t *ptr, *end, *frame_end, *obu_end;
    struct rtp_encode_av1_t *packer;
    packer = (struct rtp_encode_av1_t *)pack;
    packer->pkt.rtp.timestamp = timestamp;
    packer->pkt.rtp.m = 0;
    packer->ptr = NULL; // TODO: ptr memory leak

    temporal_id0 = spatial_id0 = 0;
    ptr = (const uint8_t *)data;
    end = ptr + bytes;
    for (packer->aggregation = AV1_AGGREGATION_HEADER_N; ptr < end; ptr = frame_end) {
        ptr = leb128(ptr, end - ptr, &frame_size);
        frame_end = ptr + frame_size;
        if (frame_end > end) {
            assert(0);
            return -1;
        }

        for (; ptr < frame_end; ptr = obu_end) {
            ptr = leb128(ptr, bytes, &obu_size);
            obu_end = ptr + obu_size;
            if (obu_end > frame_end) {
                assert(0);
                return -1;
            }

            obu_type = (*ptr >> 3) & 0x0F;
            obu_extension_flag = *ptr & 0x04;
            // obu_has_size_field = *ptr & 0x02;
            if (obu_extension_flag) {
                temporal_id = (ptr[1] >> 5) & 0x07;
                spatial_id = (ptr[1] >> 3) & 0x03;

                // If more than one OBU contained in an RTP packet has an OBU extension header
                // then the values of the temporal_id and spatial_id must be the same in all such
                // OBUs in the RTP packet.
                if (temporal_id != temporal_id0 || spatial_id != spatial_id0) {
                    r = rtp_av1_pack_flush(packer, packer->aggregation);
                    if (0 != r)
                        return r;

                    temporal_id0 = temporal_id;
                    spatial_id0 = spatial_id;
                }
            }

            // 5. Packetization rules
            // The temporal delimiter OBU, if present, SHOULD be removed
            // when transmitting, and MUST be ignored by receivers.
            if (OBU_TEMPORAL_DELIMITER == obu_type)
                continue;

            if (0 != rtp_av1_pack_obu(packer, ptr, obu_size))
                return -ENOMEM;
        }
    }

    // The RTP header Marker bit MUST be set equal to 0 if the packet is not the last
    // packet of the temporal unit, it SHOULD be set equal to 1 otherwise.
    // Note: It is possible for a receiver to receive the last packet of a temporal unit
    // without the marker bit being set equal to 1, and a receiver should be able to handle
    // this case. The last packet of a temporal unit is also indicated by the next packet,
    // in RTP sequence number order, having an incremented timestamp.
    packer->pkt.rtp.m = 1;
    return rtp_av1_pack_flush(packer, packer->aggregation);
}

/// http://aomedia.org/av1/specification/syntax/#general-obu-syntax
/// Low overhead bitstream format
static int rtp_av1_pack_input_obu(void *pack, const void *data, int bytes, uint32_t timestamp)
{
    int r;
    size_t i;
    size_t offset;
    uint64_t len;
    uint8_t obu_type;
    const uint8_t *ptr, *raw;
    struct rtp_encode_av1_t *packer;

    packer = (struct rtp_encode_av1_t *)pack;
    packer->pkt.rtp.timestamp = timestamp;
    packer->pkt.rtp.m = 0;
    packer->ptr = NULL; // TODO: ptr memory leak
    packer->aggregation = 0;

    raw = (const uint8_t *)data;
    for (i = r = 0; i < bytes && 0 == r; i += (size_t)len) {
        // http://aomedia.org/av1/specification/syntax/#obu-header-syntax
        obu_type = (raw[i] >> 3) & 0x0F;
        if (raw[i] & 0x04) // obu_extension_flag
        {
            // http://aomedia.org/av1/specification/syntax/#obu-extension-header-syntax
            // temporal_id = (obu[1] >> 5) & 0x07;
            // spatial_id = (obu[1] >> 3) & 0x03;
            offset = 2;
        } else {
            offset = 1;
        }

        if (raw[i] & 0x02) // obu_has_size_field
        {
            ptr = leb128(raw + i + offset, (int)(bytes - i - offset), &len);
            if (ptr + len > raw + bytes)
                return -1;
            len += ptr - raw - i;
        } else {
            len = bytes - i;
        }

        // 5. Packetization rules
        // The temporal delimiter OBU, if present, SHOULD be removed
        // when transmitting, and MUST be ignored by receivers.
        if (OBU_TEMPORAL_DELIMITER == obu_type)
            continue;

        packer->aggregation |= OBU_SEQUENCE_HEADER == obu_type ? AV1_AGGREGATION_HEADER_N : 0;
        r = rtp_av1_pack_obu(packer, raw + i, (size_t)len);
    }

    // The RTP header Marker bit MUST be set equal to 0 if the packet is not the last
    // packet of the temporal unit, it SHOULD be set equal to 1 otherwise.
    // Note: It is possible for a receiver to receive the last packet of a temporal unit
    // without the marker bit being set equal to 1, and a receiver should be able to handle
    // this case. The last packet of a temporal unit is also indicated by the next packet,
    // in RTP sequence number order, having an incremented timestamp.
    packer->pkt.rtp.m = 1;
    return rtp_av1_pack_flush(packer, packer->aggregation);
}

struct rtp_payload_encode_t *rtp_av1_encode()
{
    static struct rtp_payload_encode_t encode = {
        rtp_av1_pack_create,
        rtp_av1_pack_destroy,
        rtp_av1_pack_get_info,
        rtp_av1_pack_input_obu,
    };

    return &encode;
}