/*
#   live_aac_encoder.c: encode using libavformat
#   Copyright (C) 2015-2022 Stephen Fairchild (s-fairchild@users.sf.net)
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 2 of the License, or
#   (at your option) any later version.
#
#   This program 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 General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program in the file entitled COPYING.
#   If not, see <http://www.gnu.org/licenses/>.
*/

#include "../config.h"

#ifdef HAVE_AVCODEC

#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

#include <libavcodec/avcodec.h>

#include <libavutil/channel_layout.h>
#include <libavutil/common.h>
#include <libavutil/frame.h>
#include <libavutil/samplefmt.h>

#include "main.h"
#include "sourceclient.h"
#include "live_aac_encoder2.h"

#ifndef AV_PROFILE_AAC_LOW /// ok to remove in year 2027
#define AV_PROFILE_AAC_LOW FF_PROFILE_AAC_LOW
#endif

#ifndef AV_PROFILE_AAC_HE_V2 /// ok to remove in year 2027
#define AV_PROFILE_AAC_HE_V2 FF_PROFILE_AAC_HE_V2
#endif

static const AVProfile profiles[] = {
    { AV_PROFILE_AAC_LOW,   "LC"       },
    { AV_PROFILE_AAC_MAIN,  "Main" },
    { AV_PROFILE_AAC_LTP,   "LTP"  },
    { AV_PROFILE_AAC_HE,    "HE-AAC"   },
    { AV_PROFILE_AAC_HE_V2, "HE-AACv2" },
    { AV_PROFILE_UNKNOWN }
};

typedef struct State {
    struct encoder *encoder;
    AVCodecContext *c;
    enum packet_flags packet_flags;
    int encoded_samples;
    AVFrame *frame;
    AVPacket *pkt;
    size_t buf_size;
    uint8_t *buf;
    int av_profile;
    int sfi;
    int aot;
    char *metadata;
} State;

#if 0
static void show_profiles(const AVProfile *p)
{
    if (!p || p->profile == AV_PROFILE_UNKNOWN) {
        fprintf(stderr, "profiles list is empty\n");
        return;
    }
    fputs("available profiles listed below\n", stderr);
    while (p->profile != AV_PROFILE_UNKNOWN) {
        printf("profile: %s\n", p->name);
        ++p;
    }
}
#endif

static void packetize_metadata(State * const s)
{
    struct encoder *e = s->encoder;
    size_t l = 4;

    pthread_mutex_lock(&e->metadata_mutex);

    l += strlen(e->custom_meta);
    l += strlen(e->artist);
    l += strlen(e->title);
    l += strlen(e->album);

    if ((s->metadata = realloc(s->metadata, l)))
        snprintf(s->metadata, l, "%s\n%s\n%s\n%s", e->custom_meta, e->artist, e->title, e->album);
    else
        fprintf(stderr, "malloc failure\n");

    e->new_metadata = FALSE;
    pthread_mutex_unlock(&e->metadata_mutex);
}

static int sr2adts_sfi(long sr)
{
    // get adts representation of the encoded data sample rate
    const long table[] = {96000, 88200, 64000, 48000, 44100, 32000, 24000,
                          22050, 16000, 12000, 11025, 8000, 7350, 0};

    for (int i = 0; table[i]; i++) {
        if (table[i] == sr)
            return i;
    }
    return -1;
}

static int write_packet(State *s, uint8_t *buf, const int buf_size, enum packet_flags pf)
{
    struct encoder *e = s->encoder;
    struct encoder_op_packet_rb_feed packet;

    packet.header.bit_rate = e->bitrate;
    packet.header.sample_rate = e->target_samplerate;
    packet.header.n_channels = e->n_channels;
    packet.header.flags = pf;
    packet.header.data_size = buf_size;
    packet.header.serial = e->oggserial;
    packet.header.timestamp = e->timestamp = s->encoded_samples / (double)e->target_samplerate;
    packet.data = buf;
    encoder_write_packet_all(e, &packet);

    return 1;
}

static int write_packet_with_adts(State *s, const uint8_t *data, const int dsiz)
{
    struct encoder *e = s->encoder;
    const int hsiz = 7;  // header size
    const int csiz = hsiz + dsiz;  // combined header and data size;
    int ret;

    if (dsiz > 16376) {
        // adts header frame size limited to 13 bits
        fprintf(stderr, "write_packet_with_adts: frame size too large for adts header\n");
        e->encoder_state = ES_STOPPING;
        return 0;
    }

    if (csiz > s->buf_size)
        if (!(s->buf = realloc(s->buf, s->buf_size = csiz))) {
            fprintf(stderr, "write_packet_with_adts: malloc failure\n");
            e->encoder_state = ES_STOPPING;
            return 0;
        }

    // Copy ADTS header into buffer
    {   // ADTS header - n.b. buffer fullness set to all ones as per the FAAC encoder's behaviour
        uint8_t *b = s->buf;
        *b++ = 0xff;    // syncword
        *b++ = 0xf1;    // syncword, mpeg 4, layer 0, unprotected
        *b++ = ((s->aot - 1) << 6) | (s->sfi << 2); // low complexity, sample frequency index, private unset, in-band channel config
        *b++ = e->n_channels << 6 | csiz >> 11; // in-band channel config, originality false, home false, copyright id bits unset, frame length
        *b++ = (csiz >> 3) & 0xFF;             // frame length
        *b++ = ((csiz << 5) & 0xFF) | 0x1f;    // frame length, buffer fullness all 1's
        *b++ = 0xFC;                        // buffer fullness all 1's, 1 data block per frame
    }

    // Copy AAC data packet into buffer
    memcpy(s->buf + hsiz, data, dsiz);
    ret = write_packet(s, s->buf, csiz, PF_AAC | s->packet_flags);
    s->packet_flags &= ~PF_INITIAL;

    return ret;
}

static int encode(State *s, AVCodecContext *ctx, AVFrame *frame, AVPacket *pkt)
{
    int ret = 0;

    if (avcodec_send_frame(ctx, frame) < 0) {
        fprintf(stderr, "live_aac_encoder2: error sending frame to encoder\n");
        return 0;
    }

    /* read all the available output packets (in general there may be any
    * number of them */
    while (ret >= 0) {
        ret = avcodec_receive_packet(ctx, pkt);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
            return 1;
        if (ret < 0) {
            fprintf(stderr, "live_aac_encoder2: error encoding audio frame\n");
            return 0;
        }

        write_packet_with_adts(s, pkt->data, pkt->size);
        av_packet_unref(pkt);
    }
    return 1;
}

static int set_sample_fmt(AVCodecContext *c, const AVCodec *codec)
{
#ifdef HAVE_AVCODEC_GET_SUPPORTED_CONFIG
    const void *sfs;
    int nsfs;
    enum AVSampleFormat sf;
    if (!avcodec_get_supported_config(c, codec, AV_CODEC_CONFIG_SAMPLE_FORMAT,
                                      0, &sfs, &nsfs)) {
        for (int i = 0; i < nsfs; ++i) {
            sf = ((enum AVSampleFormat *)sfs)[i];
            if (sf == AV_SAMPLE_FMT_FLTP) { // || sf == AV_SAMPLE_FMT_FLT) {
                c->sample_fmt = sf;
                return 1;
            }
        }
    }
    return 0;
#else
    if (codec->sample_fmts) {
        c->sample_fmt = codec->sample_fmts[0];
        for (int i = 0; codec->sample_fmts[i]; i++) {
            if (codec->sample_fmts[i] == AV_SAMPLE_FMT_FLTP) {
                c->sample_fmt = AV_SAMPLE_FMT_FLTP;
                break;
            }
            if (codec->sample_fmts[i] == AV_SAMPLE_FMT_FLT) {
                c->sample_fmt = AV_SAMPLE_FMT_FLT;
                break;
            }
        }
    } else {
        c->sample_fmt = AV_SAMPLE_FMT_FLTP;
    }
    return 1;
#endif /* HAVE_AVCODEC_GET_SUPPORTED_CONFIG */
}

static void live_aac_encoder2_main(struct encoder *encoder)
{
    State * const s = encoder->encoder_private;

    if (encoder->encoder_state == ES_STARTING) {
        const AVCodec *codec;
        AVCodecContext *c = NULL;
        int ch;

        if (!encoder->run_request_f)
            goto bailout;

        if (!(codec = avcodec_find_encoder(AV_CODEC_ID_AAC))) {
            fprintf(stderr, "live_aac_encoder2_main: aac encoder not found\n");
            goto bailout;
        }
        if (!(s->c = c = avcodec_alloc_context3(codec))) {
            fprintf(stderr, "live_aac_encoder2_main: alloc context failed\n");
            goto bailout;
        }
        c->bit_rate = encoder->bitrate;
        c->sample_rate = encoder->target_samplerate;
        c->strict_std_compliance = -2;
        if (!set_sample_fmt(c, codec)) {
            fprintf(stderr, "live_aac_encoder2_main: fltp unsupported\n");
            goto bailout;
        }
        c->ch_layout.order = AV_CHANNEL_ORDER_NATIVE;
        c->ch_layout.nb_channels = ch = encoder->n_channels;
        switch (ch) {
            case 1:  c->ch_layout.u.mask = AV_CH_LAYOUT_MONO;
                     break;
            case 2:  c->ch_layout.u.mask = AV_CH_LAYOUT_STEREO;
                     break;
            default: fprintf(stderr, "unsupported number of channels: %d\n", ch);
                     goto bailout;
        }
        c->profile = s->av_profile;
        // show_profiles(codec->profiles);
        if (avcodec_open2(c, codec, NULL) < 0) {
            fprintf(stderr, "live_aac_encoder_main: codec open failed\n");
            goto bailout;
        }
        if (!(s->pkt = av_packet_alloc())) {
            fprintf(stderr, "live_aac_encoder_main: av_packet_alloc failed\n");
            goto bailout;
        }
        if (!(s->frame = av_frame_alloc())) {
            fprintf(stderr, "live_aac_encoder_main: av_frame_alloc failed\n");
            goto bailout;
        }
        s->frame->nb_samples = c->frame_size;
        s->frame->format = c->sample_fmt;
        if (av_channel_layout_copy(&s->frame->ch_layout, &c->ch_layout) < 0) {
            fprintf(stderr, "live_aac_encoder_main: av_channel_layout_copy failed\n");
            goto bailout;
        }
        if (av_frame_get_buffer(s->frame, 0) < 0) {
            fprintf(stderr, "live_acc_encoder_main: could not allocate audio data buffers\n");
            goto bailout;
        }

        ++encoder->oggserial;
        s->packet_flags = PF_INITIAL;
        s->encoded_samples = 0;
        if (encoder->run_request_f)
            encoder->encoder_state = ES_RUNNING;
        else
            encoder->encoder_state = ES_STOPPING;
        return;
    }
    if (encoder->encoder_state == ES_RUNNING) {
        struct encoder_ip_data *id;

        if (encoder->flush || !encoder->run_request_f) {
            encoder->flush = FALSE;
            encoder->encoder_state = ES_STOPPING;
            s->packet_flags |= PF_FINAL;
            encode(s, s->c, NULL, s->pkt);
            return;
        }
        if (av_frame_make_writable(s->frame) < 0) {
            fprintf(stderr, "live_aac_encoder2: av_frame_make_writable failed\n");
            encoder->run_request_f = FALSE;
            return;
        }
        if ((id = encoder_get_input_data(encoder, s->c->frame_size, s->c->frame_size, (float **)s->frame->data))) {
            if (!(encode(s, s->c, s->frame, s->pkt)))
                encoder->encoder_state = ES_STOPPING;
            s->encoded_samples += id->qty_samples;
            encoder_ip_data_free(id);
        }
        if (encoder->new_metadata && encoder->use_metadata) {
            packetize_metadata(s);
            if (s->metadata)
                write_packet(s, (unsigned char *)s->metadata, strlen(s->metadata) + 1, PF_METADATA);
        }
        return;
    }
    if (encoder->encoder_state == ES_STOPPING) {
        if (encoder->run_request_f) {
            av_frame_free(&s->frame);
            av_packet_free(&s->pkt);
            avcodec_free_context(&s->c);
            encoder->encoder_state = ES_STARTING;
            return;
        }
        goto bailout;
    }
    // if here we have a weird encoder_state and bailout is appropriate
bailout:
    fprintf(stderr, "live_aac_encoder2_main: performing cleanup\n");
    if (s->buf)
        free(s->buf);
    if (s->frame)
        av_frame_free(&s->frame);
    if (s->pkt)
        av_packet_free(&s->pkt);
    if (s->c)
        avcodec_free_context(&s->c);
    if (s->metadata)
        free(s->metadata);
    free(s);
    encoder->run_request_f = FALSE;
    encoder->run_encoder = NULL;
    encoder->flush = FALSE;
    encoder->encoder_private = NULL;
    encoder->encoder_state = ES_STOPPED;
    fprintf(stderr, "live_aac_encoder2_main: finished cleanup\n");
}

int live_aac_encoder2_init(struct encoder *encoder, struct encoder_vars *ev)
{
    State *s;

    if (!(s = calloc(1, sizeof (State)))) {
        fprintf(stderr, "malloc failure\n");
        return FAILED;
    }

    s->av_profile = AV_PROFILE_UNKNOWN;

    for (const AVProfile *p = profiles; p->profile != AV_PROFILE_UNKNOWN; ++p) {
        if (!strcmp(ev->profile, p->name)) {
            s->av_profile = p->profile;
            break;
        }
    }
    if (s->av_profile == AV_PROFILE_UNKNOWN) {
        fprintf(stderr, "chosen aac profile '%s' missing from list\n", ev->profile);
        return FAILED;
    }

    switch (s->av_profile) {
        case AV_PROFILE_AAC_MAIN:
            s->aot = 1;
            break;
        case AV_PROFILE_AAC_LOW:
            s->aot = 2;
            break;
        case AV_PROFILE_AAC_LTP:
            s->aot = 4;
            break;
        default:
            s->aot = 2;  // captures AAC_HE and AAC_HEv2 aka aac+
    }

    s->sfi = sr2adts_sfi(encoder->target_samplerate);
    s->encoder = encoder;
    encoder->encoder_private = s;
    encoder->run_encoder = live_aac_encoder2_main;
    return SUCCEEDED;
}

void live_aac_encoder2_profiles(FILE *fp)
{
    /* try to initiate an encoder with each of the profiles we support to see
     * what works */

    const AVProfile *p = profiles;

    struct encoder e = {
        .target_samplerate = 48000,
        .bitrate = 64000
    };

    struct encoder_vars ev = { 0 };

    while (p->profile != FF_PROFILE_UNKNOWN) {
        e.n_channels = (p->profile == FF_PROFILE_AAC_HE) ? 1 : 2;
        e.run_request_f = 1;
        e.encoder_state = ES_STARTING;
        ev.profile = (char *)p->name;
        if (live_aac_encoder2_init(&e, &ev) == SUCCEEDED) {
            e.run_encoder(&e);
            if (e.encoder_state == ES_RUNNING) {
                fprintf(g.out, "|%s|", p->name);
                e.encoder_state = ES_STOPPING;
                e.run_request_f = 0;
                e.run_encoder(&e);
            }
        }
    ++p;
    }
}

#endif /* HAVE_AVCODEC */
