/*
 * Immersive Audio Model and Formats parsing
 * Copyright (c) 2023 James Almer <jamrial@gmail.com>
 *
 * 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
 */

#include "libavutil/avassert.h"
#include "libavutil/intreadwrite.h"
#include "libavutil/log.h"
#include "libavutil/mem.h"
#include "libavcodec/get_bits.h"
#include "libavcodec/flac.h"
#include "libavcodec/leb.h"
#include "libavcodec/mpeg4audio.h"
#include "libavcodec/put_bits.h"
#include "avio_internal.h"
#include "iamf_parse.h"
#include "isom.h"

static int opus_decoder_config(IAMFCodecConfig *codec_config,
                               AVIOContext *pb, int len)
{
    int ret, left = len - avio_tell(pb);

    if (left < 11 || codec_config->audio_roll_distance >= 0)
        return AVERROR_INVALIDDATA;

    codec_config->extradata = av_malloc(left + 8);
    if (!codec_config->extradata)
        return AVERROR(ENOMEM);

    AV_WN32A(codec_config->extradata,     MKBETAG('O','p','u','s'));
    AV_WN32A(codec_config->extradata + 4, MKBETAG('H','e','a','d'));
    ret = ffio_read_size(pb, codec_config->extradata + 8, left);
    if (ret < 0)
        return ret;

    codec_config->extradata_size = left + 8;
    codec_config->sample_rate = 48000;

    return 0;
}

static int aac_decoder_config(IAMFCodecConfig *codec_config,
                              AVIOContext *pb, int len, void *logctx)
{
    MPEG4AudioConfig cfg = { 0 };
    int object_type_id, codec_id, stream_type;
    int ret, tag, left;

    if (codec_config->audio_roll_distance >= 0)
        return AVERROR_INVALIDDATA;

    ff_mp4_read_descr(logctx, pb, &tag);
    if (tag != MP4DecConfigDescrTag)
        return AVERROR_INVALIDDATA;

    object_type_id = avio_r8(pb);
    if (object_type_id != 0x40)
        return AVERROR_INVALIDDATA;

    stream_type = avio_r8(pb);
    if (((stream_type >> 2) != 5) || ((stream_type >> 1) & 1))
        return AVERROR_INVALIDDATA;

    avio_skip(pb, 3); // buffer size db
    avio_skip(pb, 4); // rc_max_rate
    avio_skip(pb, 4); // avg bitrate

    codec_id = ff_codec_get_id(ff_mp4_obj_type, object_type_id);
    if (codec_id && codec_id != codec_config->codec_id)
        return AVERROR_INVALIDDATA;

    left = ff_mp4_read_descr(logctx, pb, &tag);
    if (tag != MP4DecSpecificDescrTag ||
        !left || left > (len - avio_tell(pb)))
        return AVERROR_INVALIDDATA;

    // We pad extradata here because avpriv_mpeg4audio_get_config2() needs it.
    codec_config->extradata = av_malloc((size_t)left + AV_INPUT_BUFFER_PADDING_SIZE);
    if (!codec_config->extradata)
        return AVERROR(ENOMEM);

    ret = ffio_read_size(pb, codec_config->extradata, left);
    if (ret < 0)
        return ret;
    codec_config->extradata_size = left;
    memset(codec_config->extradata + codec_config->extradata_size, 0,
           AV_INPUT_BUFFER_PADDING_SIZE);

    ret = avpriv_mpeg4audio_get_config2(&cfg, codec_config->extradata,
                                        codec_config->extradata_size, 1, logctx);
    if (ret < 0)
        return ret;

    codec_config->sample_rate = cfg.sample_rate;

    return 0;
}

static int flac_decoder_config(IAMFCodecConfig *codec_config,
                               AVIOContext *pb, int len)
{
    int ret, left;

    if (codec_config->audio_roll_distance)
        return AVERROR_INVALIDDATA;

    avio_skip(pb, 4); // METADATA_BLOCK_HEADER

    left = len - avio_tell(pb);
    if (left < FLAC_STREAMINFO_SIZE)
        return AVERROR_INVALIDDATA;

    codec_config->extradata = av_malloc(left);
    if (!codec_config->extradata)
        return AVERROR(ENOMEM);

    ret = ffio_read_size(pb, codec_config->extradata, left);
    if (ret < 0)
        return ret;

    codec_config->extradata_size = left;
    codec_config->sample_rate = AV_RB24(codec_config->extradata + 10) >> 4;

    return 0;
}

static int ipcm_decoder_config(IAMFCodecConfig *codec_config,
                               AVIOContext *pb, int len)
{
    static const enum AVCodecID sample_fmt[2][3] = {
        { AV_CODEC_ID_PCM_S16BE, AV_CODEC_ID_PCM_S24BE, AV_CODEC_ID_PCM_S32BE },
        { AV_CODEC_ID_PCM_S16LE, AV_CODEC_ID_PCM_S24LE, AV_CODEC_ID_PCM_S32LE },
    };
    int sample_format = avio_r8(pb); // 0 = BE, 1 = LE
    int sample_size = (avio_r8(pb) / 8 - 2); // 16, 24, 32
    if (sample_format > 1 || sample_size > 2U || codec_config->audio_roll_distance)
        return AVERROR_INVALIDDATA;

    codec_config->codec_id = sample_fmt[sample_format][sample_size];
    codec_config->sample_rate = avio_rb32(pb);

    if (len - avio_tell(pb))
        return AVERROR_INVALIDDATA;

    return 0;
}

static int codec_config_obu(void *s, IAMFContext *c, AVIOContext *pb, int len)
{
    IAMFCodecConfig **tmp, *codec_config = NULL;
    AVIOContext pbc;
    uint8_t *buf;
    enum AVCodecID avcodec_id;
    unsigned codec_config_id, nb_samples, codec_id;
    int16_t audio_roll_distance;
    int ret;

    buf = av_malloc(len);
    if (!buf)
        return AVERROR(ENOMEM);

    ret = ffio_read_size(pb, buf, len);
    if (ret < 0)
        goto fail;

    ffio_init_context(&pbc, buf, len, 0, NULL, NULL, NULL, NULL);

    codec_config_id = ffio_read_leb(&pbc);
    codec_id = avio_rb32(&pbc);
    nb_samples = ffio_read_leb(&pbc);
    audio_roll_distance = avio_rb16(&pbc);

    switch (codec_id) {
    case MKBETAG('O','p','u','s'):
        avcodec_id = AV_CODEC_ID_OPUS;
        break;
    case MKBETAG('m','p','4','a'):
        avcodec_id = AV_CODEC_ID_AAC;
        break;
    case MKBETAG('f','L','a','C'):
        avcodec_id = AV_CODEC_ID_FLAC;
        break;
    default:
        avcodec_id = AV_CODEC_ID_NONE;
        break;
    }

    for (int i = 0; i < c->nb_codec_configs; i++)
        if (c->codec_configs[i]->codec_config_id == codec_config_id) {
            ret = AVERROR_INVALIDDATA;
            goto fail;
        }

    tmp = av_realloc_array(c->codec_configs, c->nb_codec_configs + 1, sizeof(*c->codec_configs));
    if (!tmp) {
        ret = AVERROR(ENOMEM);
        goto fail;
    }
    c->codec_configs = tmp;

    codec_config = av_mallocz(sizeof(*codec_config));
    if (!codec_config) {
        ret = AVERROR(ENOMEM);
        goto fail;
    }

    codec_config->codec_config_id = codec_config_id;
    codec_config->codec_id = avcodec_id;
    codec_config->nb_samples = nb_samples;
    codec_config->audio_roll_distance = audio_roll_distance;

    switch (codec_id) {
    case MKBETAG('O','p','u','s'):
        ret = opus_decoder_config(codec_config, &pbc, len);
        break;
    case MKBETAG('m','p','4','a'):
        ret = aac_decoder_config(codec_config, &pbc, len, s);
        break;
    case MKBETAG('f','L','a','C'):
        ret = flac_decoder_config(codec_config, &pbc, len);
        break;
    case MKBETAG('i','p','c','m'):
        ret = ipcm_decoder_config(codec_config, &pbc, len);
        break;
    default:
        break;
    }
    if (ret < 0)
        goto fail;

    if ((codec_config->nb_samples > INT_MAX) || codec_config->nb_samples <= 0 ||
        (-codec_config->audio_roll_distance > INT_MAX / codec_config->nb_samples)) {
        ret = AVERROR_INVALIDDATA;
        goto fail;
    }

    c->codec_configs[c->nb_codec_configs++] = codec_config;

    len -= avio_tell(&pbc);
    if (len)
       av_log(s, AV_LOG_WARNING, "Underread in codec_config_obu. %d bytes left at the end\n", len);

    ret = 0;
fail:
    av_free(buf);
    if (ret < 0) {
        if (codec_config)
            av_free(codec_config->extradata);
        av_free(codec_config);
    }
    return ret;
}

int ff_iamf_parse_obu_header(const uint8_t *buf, int buf_size,
                             unsigned *obu_size, int *start_pos, enum IAMF_OBU_Type *type,
                             unsigned *skip_samples, unsigned *discard_padding)
{
    GetBitContext gb;
    int ret, extension_flag, trimming, start;
    unsigned skip = 0, discard = 0;
    unsigned size;

    ret = init_get_bits8(&gb, buf, FFMIN(buf_size, MAX_IAMF_OBU_HEADER_SIZE));
    if (ret < 0)
        return ret;

    *type          = get_bits(&gb, 5);
    /*redundant      =*/ get_bits1(&gb);
    trimming       = get_bits1(&gb);
    extension_flag = get_bits1(&gb);

    *obu_size = get_leb(&gb);
    if (*obu_size > INT_MAX)
        return AVERROR_INVALIDDATA;

    start = get_bits_count(&gb) / 8;

    if (trimming) {
        discard = get_leb(&gb); // num_samples_to_trim_at_end
        skip = get_leb(&gb); // num_samples_to_trim_at_start
    }

    if (skip_samples)
        *skip_samples = skip;
    if (discard_padding)
        *discard_padding = discard;

    if (extension_flag) {
        unsigned int extension_bytes;
        extension_bytes = get_leb(&gb);
        if (extension_bytes > INT_MAX / 8)
            return AVERROR_INVALIDDATA;
        skip_bits_long(&gb, extension_bytes * 8);
    }

    if (get_bits_left(&gb) < 0)
        return AVERROR_INVALIDDATA;

    size = *obu_size + start;
    if (size > INT_MAX)
        return AVERROR_INVALIDDATA;

    *obu_size -= get_bits_count(&gb) / 8 - start;
    *start_pos = size - *obu_size;

    return size;
}

int ff_iamfdec_read_descriptors(IAMFContext *c, AVIOContext *pb,
                                int max_size, void *log_ctx)
{
    uint8_t header[MAX_IAMF_OBU_HEADER_SIZE + AV_INPUT_BUFFER_PADDING_SIZE];
    int ret;

    while (1) {
        unsigned obu_size;
        enum IAMF_OBU_Type type;
        int start_pos, len, size;

        if ((ret = ffio_ensure_seekback(pb, FFMIN(MAX_IAMF_OBU_HEADER_SIZE, max_size))) < 0)
            return ret;
        size = avio_read(pb, header, FFMIN(MAX_IAMF_OBU_HEADER_SIZE, max_size));
        if (size < 0)
            return size;
        memset(header + size, 0, AV_INPUT_BUFFER_PADDING_SIZE);

        len = ff_iamf_parse_obu_header(header, size, &obu_size, &start_pos, &type, NULL, NULL);
        if (len < 0 || obu_size > max_size) {
            av_log(log_ctx, AV_LOG_ERROR, "Failed to read obu header\n");
            avio_seek(pb, -size, SEEK_CUR);
            return len;
        }

        if (type >= IAMF_OBU_IA_PARAMETER_BLOCK && type < IAMF_OBU_IA_SEQUENCE_HEADER) {
            avio_seek(pb, -size, SEEK_CUR);
            break;
        }

        avio_seek(pb, -(size - start_pos), SEEK_CUR);
        switch (type) {
        case IAMF_OBU_IA_CODEC_CONFIG:
            ret = codec_config_obu(log_ctx, c, pb, obu_size);
            break;
        default: {
            int64_t offset = avio_skip(pb, obu_size);
            if (offset < 0)
                ret = offset;
            break;
        }
        }
        if (ret < 0) {
            av_log(log_ctx, AV_LOG_ERROR, "Failed to read obu type %d\n", type);
            return ret;
        }
        max_size -= obu_size + start_pos;
        if (max_size < 0)
            return AVERROR_INVALIDDATA;
        if (!max_size)
            break;
    }

    return 0;
}

void ff_iamf_uninit_context(IAMFContext *c)
{
    if (!c)
        return;

    for (int i = 0; i < c->nb_codec_configs; i++) {
        av_free(c->codec_configs[i]->extradata);
        av_free(c->codec_configs[i]);
    }
    av_freep(&c->codec_configs);
    c->nb_codec_configs = 0;
}