/*
 * Immersive Audio Model and Formats demuxer
 * 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 "avformat.h"
#include "iamf_parse.h"
#include "internal.h"

#define RAW_PACKET_SIZE 2048

//return < 0 if we need more data
static int get_score(const uint8_t *buf, int buf_size, enum IAMF_OBU_Type type, int *seq)
{
    if (type == IAMF_OBU_IA_SEQUENCE_HEADER) {
        if (buf_size < 4 || AV_RB32(buf) != MKBETAG('i','a','m','f'))
            return 0;
        *seq = 1;
        return -1;
    }
    if (type >= IAMF_OBU_IA_CODEC_CONFIG && type <= IAMF_OBU_IA_TEMPORAL_DELIMITER)
        return *seq ? -1 : 0;
    if (type >= IAMF_OBU_IA_AUDIO_FRAME && type <= IAMF_OBU_IA_AUDIO_FRAME_ID17)
        return *seq ? AVPROBE_SCORE_EXTENSION + 1 : 0;
    return 0;
}

static int iamf_probe(const AVProbeData *p)
{
    unsigned obu_size;
    enum IAMF_OBU_Type type;
    int seq = 0, cnt = 0, start_pos;
    int ret;

    while (1) {
        int size = ff_iamf_parse_obu_header(p->buf + cnt, p->buf_size - cnt,
                                            &obu_size, &start_pos, &type,
                                            NULL, NULL);
        if (size < 0)
            return 0;

        ret = get_score(p->buf + cnt + start_pos,
                        p->buf_size - cnt - start_pos,
                        type, &seq);
        if (ret >= 0)
            return ret;

        cnt += FFMIN(size, p->buf_size - cnt);
    }
    return 0;
}

static int iamf_read_header(AVFormatContext *s)
{
    IAMFContext *iamf;
    int ret;

    iamf = av_mallocz(sizeof(IAMFContext));
    if (!iamf)
        return AVERROR(ENOMEM);

    ret = ff_iamfdec_read_descriptors(iamf, s->pb, INT_MAX, s);
    if (ret < 0) {
        ff_iamf_uninit_context(iamf);
        av_freep(&iamf);
        return ret;
    }

    if (iamf->nb_codec_configs > 0) {
        AVStream *st = avformat_new_stream(s, NULL);
        if (!st) {
            ff_iamf_uninit_context(iamf);
            av_freep(&iamf);
            return AVERROR(ENOMEM);
        }

        st->codecpar->codec_type = AVMEDIA_TYPE_AUDIO;
        st->codecpar->codec_id = iamf->codec_configs[0]->codec_id;
        st->codecpar->codec_tag = iamf->codec_configs[0]->codec_tag;
        st->codecpar->sample_rate = iamf->codec_configs[0]->sample_rate;
        st->codecpar->channels = 2;
        st->need_parsing = AVSTREAM_PARSE_FULL_RAW;
        st->start_time = 0;

        // read iamf header, current pos is iamf header size
        int64_t iamf_header_size = avio_tell(s->pb);
        st->codecpar->extradata_size = iamf_header_size;
        st->codecpar->extradata = av_malloc(iamf_header_size);
        if (!st->codecpar->extradata) {
            ff_iamf_uninit_context(iamf);
            av_freep(&iamf);
            return AVERROR(ENOMEM);
        }

        avio_seek(s->pb, 0, SEEK_SET);
        ret = avio_read(s->pb, st->codecpar->extradata, iamf_header_size);
        if (ret != iamf_header_size) {
            ff_iamf_uninit_context(iamf);
            av_freep(&iamf);
            return ret < 0 ? ret : AVERROR_INVALIDDATA;
        }

        av_dict_set_int(&st->metadata, "iamf", 1, 0);

        ff_iamf_uninit_context(iamf);
        av_freep(&iamf);
    }

    if (!s->nb_streams)
        return AVERROR_INVALIDDATA;

    return 0;
}

static int iamf_read_packet(AVFormatContext *s, AVPacket *pkt)
{
    int ret, size = RAW_PACKET_SIZE;

    if ((ret = av_new_packet(pkt, size)) < 0)
        return ret;

    pkt->pos= avio_tell(s->pb);
    pkt->stream_index = 0;
    ret = avio_read_partial(s->pb, pkt->data, size);
    if (ret < 0) {
        av_packet_unref(pkt);
        return ret;
    }
    av_shrink_packet(pkt, ret);

    return ret;
}

AVInputFormat ff_iamf_demuxer = {
    .name         = "iamf",
    .long_name    = NULL_IF_CONFIG_SMALL("Raw Immersive Audio Model and Formats"),
    .extensions   = "iamf",
    .flags        = AVFMT_GENERIC_INDEX | AVFMT_NO_BYTE_SEEK | AVFMT_NOTIMESTAMPS | AVFMT_SHOW_IDS,
    .read_probe     = iamf_probe,
    .read_header    = iamf_read_header,
    .read_packet    = iamf_read_packet,
};