/*
 * Depacketization for RTP Payload Format For AV1 (v1.0)
 * https://aomediacodec.github.io/av1-rtp-spec/
 * Copyright (c) 2024 Axis Communications
 *
 * This file is part of FFmpeg.
 *
 * FFmpeg is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * FFmpeg is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with FFmpeg; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 */

/**
 * @file
 * @brief AV1 / RTP depacketization code (RTP Payload Format For AV1 (v1.0))
 * @author Chris Hodges <chris.hodges@axis.com>
 * @note The process will restore TDs and put back size fields into headers.
 *       It will also try to keep complete OBUs and remove partial OBUs
 *       caused by packet drops and thus keep the stream syntactically intact.
 */

#include "libavutil/avstring.h"
#include "libavutil/mem.h"
#include "avformat.h"

#include "rtpdec.h"
#include "libavcodec/av1.h"
#include "rtp_av1.h"

// enable tracing of packet data
//#define RTPDEC_AV1_VERBOSE_TRACE

/**
 * RTP/AV1 specific private data.
 */
struct PayloadContext {
    uint32_t timestamp; ///< last received timestamp for frame
    uint8_t profile;    ///< profile (main/high/professional)
    uint8_t level_idx;  ///< level (0-31)
    uint8_t tier;       ///< main tier or high tier
    uint16_t prev_seq;  ///< sequence number of previous packet
    unsigned int frag_obu_size;     ///< current total size of fragmented OBU
    unsigned int frag_pkt_leb_pos;  ///< offset in buffer where OBU LEB starts
    unsigned int frag_lebs_res;     ///< number of bytes reserved for LEB
    unsigned int frag_header_size;  ///< size of OBU header (1 or 2)
    int needs_td;                   ///< indicates that a TD should be output
    int drop_fragment;              ///< drop all fragments until next frame
    int keyframe_seen;              ///< keyframe was seen
    int wait_for_keyframe;          ///< message about waiting for keyframe has been issued
};

static int sdp_parse_fmtp_config_av1(AVFormatContext *s,
                                     AVStream *stream,
                                     PayloadContext *av1_data,
                                     const char *attr, const char *value) {
    if (!strcmp(attr, "profile")) {
        av1_data->profile = atoi(value);
        av_log(s, AV_LOG_DEBUG, "RTP AV1 profile: %u\n", av1_data->profile);
    } else if (!strcmp(attr, "level-idx")) {
        av1_data->level_idx = atoi(value);
        av_log(s, AV_LOG_DEBUG, "RTP AV1 level: %u\n", av1_data->profile);
    } else if (!strcmp(attr, "tier")) {
        av1_data->tier = atoi(value);
        av_log(s, AV_LOG_DEBUG, "RTP AV1 tier: %u\n", av1_data->tier);
    }
    return 0;
}

// return 0 on complete packet, -1 on partial packet
static int av1_handle_packet(AVFormatContext *ctx, PayloadContext *data,
                             AVStream *st, AVPacket *pkt, uint32_t *timestamp,
                             const uint8_t *buf, int len, uint16_t seq,
                             int flags) {
    uint8_t aggr_hdr;
    int result = 0;
    int is_frag_cont;
    int is_last_fragmented;
    int is_first_pkt;
    unsigned int num_obus;
    unsigned int obu_cnt = 1;
    unsigned int rem_pkt_size = len;
    unsigned int pktpos;
    const uint8_t *buf_ptr = buf;
    uint16_t expected_seq = data->prev_seq + 1;
    int16_t seq_diff = seq - expected_seq;

    data->prev_seq = seq;

    if (!len) {
        av_log(ctx, AV_LOG_ERROR, "Empty AV1 RTP packet\n");
        return AVERROR_INVALIDDATA;
    }
    if (len < 2) {
        av_log(ctx, AV_LOG_ERROR, "AV1 RTP packet too short\n");
        return AVERROR_INVALIDDATA;
    }

    /* The payload structure is supposed to be straight-forward, but there are a
     * couple of edge cases which need to be tackled and make things a bit more
     * complex.
     * These are mainly due to:
     * - To reconstruct the OBU size for fragmented packets and place it the OBU
     *   header, the final size will not be known until the last fragment has
     *   been parsed. However, the number LEBs in the header is variable
     *   depending on the length of the payload.
     * - We are increasing the out-packet size while we are getting fragmented
     *   OBUs. If an RTP packet gets dropped, we would create corrupted OBUs.
     *   In this case we decide to drop the whole frame.
     */

#ifdef RTPDEC_AV1_VERBOSE_TRACE
    av_log(ctx, AV_LOG_TRACE, "RTP Packet %d in (%x), len=%d:\n",
           seq, flags, len);
    av_hex_dump_log(ctx, AV_LOG_TRACE, buf, FFMIN(len, 64));
    av_log(ctx, AV_LOG_TRACE, "... end at offset %x:\n", FFMAX(len - 64, 0));
    av_hex_dump_log(ctx, AV_LOG_TRACE, buf + FFMAX(len - 64, 0), FFMIN(len - 64, 64));
#endif

    /* 8 bit aggregate header: Z Y W W N - - - */
    aggr_hdr = *buf_ptr++;
    rem_pkt_size--;

    /* Z: MUST be set to 1 if the first OBU element is an OBU fragment that is a
     * continuation of an OBU fragment from the previous packet, and MUST be set
     * to 0 otherwise */
    is_frag_cont = (aggr_hdr >> AV1B_AGGR_HDR_FRAG_CONT) & 1;

    /* Y: MUST be set to 1 if the last OBU element is an OBU fragment that will
     * continue in the next packet, and MUST be set to 0 otherwise */
    is_last_fragmented = (aggr_hdr >> AV1B_AGGR_HDR_LAST_FRAG) & 1;

    /* W: two bit field that describes the number of OBU elements in the packet.
     * This field MUST be set equal to 0 or equal to the number of OBU elements
     * contained in the packet.
     * If set to 0, each OBU element MUST be preceded by a length field.
     * If not set to 0 (i.e., W = 1, 2 or 3) the last OBU element MUST NOT be
     * preceded by a length field (it's derived from RTP packet size minus other
     * known lengths). */
    num_obus = (aggr_hdr >> AV1S_AGGR_HDR_NUM_OBUS) & AV1M_AGGR_HDR_NUM_OBUS;

    /* N: MUST be set to 1 if the packet is the first packet of a coded video
     * sequence, and MUST be set to 0 otherwise.*/
    is_first_pkt = (aggr_hdr >> AV1B_AGGR_HDR_FIRST_PKT) & 1;

    if (is_frag_cont) {
        if (data->drop_fragment) {
            return AVERROR_INVALIDDATA;
        }
        if (is_first_pkt) {
            av_log(ctx, AV_LOG_ERROR, "Illegal aggregation header in first AV1 RTP packet\n");
            return AVERROR_INVALIDDATA;
        }
        if (seq_diff) {
            av_log(ctx, AV_LOG_WARNING, "AV1 RTP frag packet sequence mismatch (%d != %d), dropping temporal unit\n",
                   seq, expected_seq);
            goto drop_fragment;
        }
        if (!pkt->size || !data->frag_obu_size) {
            av_log(ctx, AV_LOG_WARNING, "Unexpected fragment continuation in AV1 RTP packet\n");
            goto drop_fragment; // avoid repeated output for the same fragment
        }
    } else {
        if (!is_first_pkt && !data->keyframe_seen) {
            if (!data->wait_for_keyframe) {
                data->wait_for_keyframe = 1;
                av_log(ctx, AV_LOG_WARNING, "AV1 RTP packet before keyframe, dropping and waiting for next keyframe\n");
            }
            goto drop_fragment;
        }
        if (seq_diff && !is_first_pkt) {
            av_log(ctx, AV_LOG_WARNING, "AV1 RTP unfrag packet sequence mismatch (%d != %d), dropping temporal unit\n",
                   seq, expected_seq);
            goto drop_fragment;
        }
        data->drop_fragment = 0;
        if (!data->needs_td && ((data->timestamp != *timestamp) || is_first_pkt)) {
            av_log(ctx, AV_LOG_TRACE, "Timestamp changed to %u (or first pkt %d), forcing TD\n", *timestamp, is_first_pkt);
            data->needs_td = 1;
            data->frag_obu_size = 0; // new temporal unit might have been caused by dropped packets
        }
        if (data->frag_obu_size) {
            data->frag_obu_size = 0; // make sure we recover
            av_log(ctx, AV_LOG_ERROR, "Missing fragment continuation in AV1 RTP packet\n");
            return AVERROR_INVALIDDATA;
        }
        // update the timestamp in the frame packet with the one from the RTP packet
        data->timestamp = *timestamp;
    }
    pktpos = pkt->size;

#ifdef RTPDEC_AV1_VERBOSE_TRACE
    av_log(ctx, AV_LOG_TRACE, "Input buffer size %d, aggr head 0x%02x fc %d, lf %d, no %d, fp %d\n",
           len, aggr_hdr, is_frag_cont, is_last_fragmented, num_obus, is_first_pkt);
#endif

    if (is_first_pkt) {
        pkt->flags |= AV_PKT_FLAG_KEY;
        data->keyframe_seen = 1;
        data->wait_for_keyframe = 0;
    }

    // loop over OBU elements
    while (rem_pkt_size) {
        uint32_t obu_size;
        int num_lebs;
        int needs_size_field;
        int output_size;
        unsigned int obu_payload_size;
        uint8_t obu_hdr;

        obu_size = rem_pkt_size;
        if (!num_obus || obu_cnt < num_obus) {
            // read out explicit OBU element size (which almost corresponds to the original OBU size)
            num_lebs = parse_leb(ctx, buf_ptr, rem_pkt_size, &obu_size);
            if (!num_lebs) {
                return AVERROR_INVALIDDATA;
            }
            rem_pkt_size -= num_lebs;
            buf_ptr += num_lebs;
        }
        // read first byte (which is the header byte only for non-fragmented elements)
        obu_hdr = *buf_ptr;
        if (obu_size > rem_pkt_size) {
            av_log(ctx, AV_LOG_ERROR, "AV1 OBU size %u larger than remaining pkt size %d\n", obu_size, rem_pkt_size);
            return AVERROR_INVALIDDATA;
        }

        if (!obu_size) {
            av_log(ctx, AV_LOG_ERROR, "Unreasonable AV1 OBU size %u\n", obu_size);
            return AVERROR_INVALIDDATA;
        }

        if (!is_frag_cont) {
            uint8_t obu_type = (obu_hdr >> AV1S_OBU_TYPE) & AV1M_OBU_TYPE;
            if (obu_hdr & AV1F_OBU_FORBIDDEN) {
                av_log(ctx, AV_LOG_ERROR, "Forbidden bit set in AV1 OBU header (0x%02x)\n", obu_hdr);
                return AVERROR_INVALIDDATA;
            }
            // ignore and remove OBUs according to spec
            if ((obu_type == AV1_OBU_TEMPORAL_DELIMITER) ||
                (obu_type == AV1_OBU_TILE_LIST)) {
                pktpos += obu_size;
                rem_pkt_size -= obu_size;
                // TODO: This probably breaks if the OBU_TILE_LIST is fragmented
                // into the next RTP packet, so at least check and fail here
                if (rem_pkt_size == 0 && is_last_fragmented) {
                    avpriv_report_missing_feature(ctx, "AV1 OBU_TILE_LIST (should not be there!) to be ignored but is fragmented\n");
                    return AVERROR_PATCHWELCOME;
                }
                obu_cnt++;
                continue;
            }
        }

        // If we need to add a size field, out size will be different
        output_size = obu_size;
        // Spec says the OBUs should have their size fields removed,
        // but this is not mandatory
        if (is_frag_cont || (obu_hdr & AV1F_OBU_HAS_SIZE_FIELD)) {
            needs_size_field = 0;
        } else {
            needs_size_field = 1;
            // (re)calculate number of LEB bytes needed (if it was implicit, there were no LEBs)
            output_size += calc_leb_size(obu_size - (1 + ((obu_hdr & AV1F_OBU_EXTENSION_FLAG) ? 1 : 0)));
        }

        if (!is_frag_cont && (obu_cnt == 1)) {
            if (data->needs_td) {
                output_size += 2; // for Temporal Delimiter (TD)
            }
            if (pkt->data) {
                if ((result = av_grow_packet(pkt, output_size)) < 0)
                    return result;
            } else {
                if ((result = av_new_packet(pkt, output_size) < 0))
                    return result;
            }

            if (data->needs_td) {
                // restore TD
                pkt->data[pktpos++] = 0x12;
                pkt->data[pktpos++] = 0x00;
            }
            data->needs_td = 0;
        } else {
            if ((result = av_grow_packet(pkt, output_size)) < 0)
                return result;
        }

        obu_payload_size = obu_size;
        // do we need to restore the OBU size field?
        if (needs_size_field) {
            // set obu_has_size_field in header byte
            pkt->data[pktpos++] = *buf_ptr++ | AV1F_OBU_HAS_SIZE_FIELD;
            data->frag_header_size = 1;
            obu_payload_size--;

            // copy extension byte, if available
            if (obu_hdr & AV1F_OBU_EXTENSION_FLAG) {
                /* TODO we cannot handle the edge case where last element is a
                 * fragment of exactly one byte AND the header has the extension
                 * flag set. Note that it would be more efficient to not send a
                 * fragment of one byte and instead drop the size field of the
                 * prior element */
                if (!obu_payload_size) {
                    av_log(ctx, AV_LOG_ERROR, "AV1 OBU too short for extension byte (0x%02x)\n",
                           obu_hdr);
                    return AVERROR_INVALIDDATA;
                }
                pkt->data[pktpos++] = *buf_ptr++;
                data->frag_header_size = 2;
                obu_payload_size--;
            }

            // remember start position of LEB for possibly fragmented packet to
            // fixup OBU size later
            data->frag_pkt_leb_pos = pktpos;
            // write intermediate OBU size field
            num_lebs = write_leb(pkt->data + pktpos, obu_payload_size);
            data->frag_lebs_res = num_lebs;
            pktpos += num_lebs;
        } else if (!is_frag_cont) {
            data->frag_lebs_res = 0;
        }
        // copy verbatim or without above header size patch
        memcpy(pkt->data + pktpos, buf_ptr, obu_payload_size);
        pktpos += obu_payload_size;
        buf_ptr += obu_payload_size;
        rem_pkt_size -= obu_size;

        // if we were handling a fragmented packet and this was the last
        // fragment, correct OBU size field
        if (data->frag_obu_size && (rem_pkt_size || !is_last_fragmented)) {
            if (data->frag_lebs_res) {
                uint32_t final_obu_size = data->frag_obu_size + obu_size - data->frag_header_size;
                uint8_t *lebptr = pkt->data + data->frag_pkt_leb_pos;
                num_lebs = calc_leb_size(final_obu_size);

                // check if we had allocated enough LEB bytes in header,
                // otherwise make some extra space
                if (num_lebs > data->frag_lebs_res) {
                    int extra_bytes = num_lebs - data->frag_lebs_res;
                    if ((result = av_grow_packet(pkt, extra_bytes)) < 0)
                        return result;
                    // update pointer in case buffer address changed
                    lebptr = pkt->data + data->frag_pkt_leb_pos;
                    // move existing data for OBU back a bit
                    memmove(lebptr + extra_bytes, lebptr,
                            pkt->size - extra_bytes - data->frag_pkt_leb_pos);
                    // move pktpos further down for following OBUs in same packet.
                    pktpos += extra_bytes;
                }

                // update OBU size field
                write_leb(lebptr, final_obu_size);
            }
            data->frag_obu_size = 0; // signal end of fragment
        } else if (is_last_fragmented && !rem_pkt_size) {
            data->frag_obu_size += obu_size;
            // fragment not yet finished!
            result = -1;
        }
        is_frag_cont = 0;

        if (!rem_pkt_size && num_obus && (num_obus != obu_cnt)) {
            av_log(ctx, AV_LOG_WARNING, "AV1 aggregation header indicated %u OBU elements, was %u\n",
                   num_obus, obu_cnt);
        }
        obu_cnt++;
    }

    if (flags & RTP_FLAG_MARKER) {
        av_log(ctx, AV_LOG_TRACE, "TD on next packet due to marker\n");
        data->needs_td = 1;
    } else {
        // fragment may be complete, but temporal unit is not yet finished
        result = -1;
    }

    if (!is_last_fragmented) {
        data->frag_obu_size = 0;
        data->frag_pkt_leb_pos = 0;
    }

#ifdef RTPDEC_AV1_VERBOSE_TRACE
    if (!result) {
        av_log(ctx, AV_LOG_TRACE, "AV1 out pkt-size: %d\n", pkt->size);
        av_hex_dump_log(ctx, AV_LOG_TRACE, pkt->data, FFMIN(pkt->size, 64));
        av_log(ctx, AV_LOG_TRACE, "... end at offset %x:\n", FFMAX(pkt->size - 64, 0));
        av_hex_dump_log(ctx, AV_LOG_TRACE, pkt->data + FFMAX(pkt->size - 64, 0), FFMIN(pkt->size, 64));
    }
#endif
    pkt->stream_index = st->index;

    return result;

drop_fragment:
    data->keyframe_seen = 0;
    data->drop_fragment = 1;
    data->frag_obu_size = 0;
    data->needs_td = 1;
    if (pkt->size) {
        av_log(ctx, AV_LOG_TRACE, "Dumping current AV1 frame packet\n");
        // we can't seem to deallocate the fragmented packet, but we can shrink it to 0
        av_shrink_packet(pkt, 0);
    }
    return AVERROR_INVALIDDATA;
}

static void av1_close_context(PayloadContext *data) {
}

static int av1_need_keyframe(PayloadContext *data)
{
    return !data->keyframe_seen;
}

static int parse_av1_sdp_line(AVFormatContext *s, int st_index,
                              PayloadContext *av1_data, const char *line) {
    AVStream * stream;
    const char *p = line;
    int result = 0;

    if (st_index < 0)
        return 0;

    stream = s->streams[st_index];

    /* Optional parameters are profile, level-idx, and tier.
     * See Section 7.2.1 of https://aomediacodec.github.io/av1-rtp-spec/ */
    if (av_strstart(p, "fmtp:", &p)) {
        result = ff_parse_fmtp(s, stream, av1_data, p, sdp_parse_fmtp_config_av1);
        av_log(s, AV_LOG_DEBUG, "RTP AV1 Profile: %u, Level: %u, Tier: %u\n",
               av1_data->profile, av1_data->level_idx, av1_data->tier);
    }

    return result;
}

const RTPDynamicProtocolHandler ff_av1_dynamic_handler = {
        .enc_name         = "AV1",
        .codec_type       = AVMEDIA_TYPE_VIDEO,
        .codec_id         = AV_CODEC_ID_AV1,
        .need_parsing     = AVSTREAM_PARSE_FULL,
        .priv_data_size   = sizeof(PayloadContext),
        .parse_sdp_a_line = parse_av1_sdp_line,
        .close            = av1_close_context,
        .parse_packet     = av1_handle_packet,
        .need_keyframe    = av1_need_keyframe,
};