#include "encoder_aac.h"
#include "headers_ffmpeg.h"

#include "ring_buffer.h"

#include "error_define.h"
#include "log_helper.h"

namespace am {
encoder_aac::encoder_aac()
{
    ffmpeg_register_all();

    _ring_buffer = new ring_buffer<AVFrame>(1024 * 1024 * 10);

    _inited = false;
    _running = false;

    _encoder = NULL;
    _encoder_ctx = NULL;
    _frame = NULL;
    _buff = NULL;
    _buff_size = 0;

    _cond_notify = false;

#ifdef SAVE_AAC
    _aac_io_ctx = NULL;
    _aac_stream = NULL;
    _aac_fmt_ctx = NULL;
#endif
}

encoder_aac::~encoder_aac()
{
    stop();

    cleanup();

    delete _ring_buffer;
}

int encoder_aac::init(int nb_channels, int sample_rate, AVSampleFormat fmt, int bit_rate)
{
    int err = AE_NO;
    int ret = 0;

    if (_inited == true)
        return err;

    do {
        _encoder = (AVCodec*)avcodec_find_encoder(AV_CODEC_ID_AAC);
        if (!_encoder) {
            err = AE_FFMPEG_FIND_ENCODER_FAILED;
            break;
        }

        _encoder_ctx = avcodec_alloc_context3(_encoder);
        if (!_encoder_ctx) {
            err = AE_FFMPEG_ALLOC_CONTEXT_FAILED;
            break;
        }

        ffmpeg_set_channels(_encoder_ctx, nb_channels);
        ffmpeg_set_channel_layout(_encoder_ctx, ffmpeg_get_default_channel_layout(nb_channels));
        _encoder_ctx->sample_rate = sample_rate;
        _encoder_ctx->sample_fmt = fmt;
        _encoder_ctx->bit_rate = bit_rate;

        _encoder_ctx->strict_std_compliance = FF_COMPLIANCE_EXPERIMENTAL;
        _encoder_ctx->time_base.den = sample_rate;
        _encoder_ctx->time_base.num = 1;
        _encoder_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

        ret = avcodec_open2(_encoder_ctx, _encoder, NULL);
        if (ret < 0) {
            err = AE_FFMPEG_OPEN_CODEC_FAILED;
            break;
        }

        _buff_size = av_samples_get_buffer_size(NULL, nb_channels, _encoder_ctx->frame_size, fmt, 1);
        _buff = (uint8_t *) av_malloc(_buff_size);

        _frame = av_frame_alloc();
        if (!_frame) {
            err = AE_FFMPEG_ALLOC_FRAME_FAILED;
            break;
        }

        ffmpeg_set_frame_channels(_frame, nb_channels);
        _frame->nb_samples = _encoder_ctx->frame_size;
        ffmpeg_set_frame_channel_layout(_frame, ffmpeg_get_default_channel_layout(nb_channels));
        _frame->format = fmt;
        _frame->sample_rate = sample_rate;

        ret = avcodec_fill_audio_frame(_frame, nb_channels, fmt, _buff, _buff_size, 0);

#ifdef SAVE_AAC
        ret = avio_open(&_aac_io_ctx, "save.aac", AVIO_FLAG_READ_WRITE);
        if (ret < 0) {
            err = AE_FFMPEG_OPEN_IO_FAILED;
            break;
        }

        _aac_fmt_ctx = avformat_alloc_context();
        if (!_aac_fmt_ctx) {
            err = AE_FFMPEG_ALLOC_CONTEXT_FAILED;
            break;
        }

        _aac_fmt_ctx->pb = _aac_io_ctx;
        _aac_fmt_ctx->oformat = av_guess_format(NULL, "save.aac", NULL);
        _aac_fmt_ctx->url = av_strdup("save.aac");

        _aac_stream = avformat_new_stream(_aac_fmt_ctx, NULL);
        if (!_aac_stream) {
            err = AE_FFMPEG_CREATE_STREAM_FAILED;
            break;
        }

        ret = avcodec_parameters_from_context(_aac_stream->codecpar, _encoder_ctx);
        if (ret < 0) {
            err = AE_FFMPEG_COPY_PARAMS_FAILED;
            break;
        }

        if (_aac_fmt_ctx->oformat->flags | AV_CODEC_FLAG_GLOBAL_HEADER) {
            _aac_stream->codec->extradata_size
                = _encoder_ctx->extradata_size; // +AV_INPUT_BUFFER_PADDING_SIZE;
            _aac_stream->codec->extradata = (uint8_t *) av_memdup(_encoder_ctx->extradata,
                                                                  _encoder_ctx->extradata_size);
        }
#endif

        _inited = true;

    } while (0);

    if (err != AE_NO) {
        al_debug("%s,error:%d", err2str(err), ret);
        cleanup();
    }

    return err;
}

int encoder_aac::get_extradata_size()
{
    return _encoder_ctx->extradata_size;
}

const uint8_t *encoder_aac::get_extradata()
{
    return (const uint8_t *) _encoder_ctx->extradata;
}

int encoder_aac::get_nb_samples()
{
    return _encoder_ctx->frame_size;
}

int encoder_aac::start()
{
    int error = AE_NO;

    if (_running == true) {
        return error;
    }

    if (_inited == false) {
        return AE_NEED_INIT;
    }

    _running = true;
    _thread = std::thread(std::bind(&encoder_aac::encode_loop, this));

    return error;
}

void encoder_aac::stop()
{
    _running = false;

    _cond_notify = true;
    _cond_var.notify_all();

    if (_thread.joinable())
        _thread.join();
}

int encoder_aac::put(const uint8_t *data, int data_len, AVFrame *frame)
{
    std::unique_lock<std::mutex> lock(_mutex);

    AVFrame frame_cp;
    memcpy(&frame_cp, frame, sizeof(AVFrame));

    _ring_buffer->put(data, data_len, frame_cp);

    _cond_notify = true;
    _cond_var.notify_all();

    return 0;
}

const AVRational &encoder_aac::get_time_base()
{
    return _encoder_ctx->time_base;
}

AVCodecID encoder_aac::get_codec_id()
{
    if (_inited == false)
        return AV_CODEC_ID_NONE;

    return _encoder->id;
}

int encoder_aac::encode(AVFrame *frame, AVPacket *packet)
{
    int ret = avcodec_send_frame(_encoder_ctx, frame);
    if (ret < 0) {
        return AE_FFMPEG_ENCODE_FRAME_FAILED;
    }

    while (ret == 0) {
        av_init_packet(packet);

        ret = avcodec_receive_packet(_encoder_ctx, packet);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            break;
        }

        if (ret < 0) {
            return AE_FFMPEG_READ_PACKET_FAILED;
        }

        if (ret == 0 && _on_data)
            _on_data(packet);

#ifdef SAVE_AAC
        av_packet_rescale_ts(packet, _encoder_ctx->time_base, _aac_stream->time_base);
        packet->stream_index = _aac_stream->index;
        av_write_frame(_aac_fmt_ctx, packet);
#endif

        av_packet_unref(packet);
    }

    return AE_NO;
}

void encoder_aac::encode_loop()
{
    int len = 0;
    int error = AE_NO;

    AVPacket *packet = av_packet_alloc();
    AVFrame pcm_frame;

#ifdef SAVE_AAC
    avformat_write_header(_aac_fmt_ctx, NULL);
#endif

    while (_running) {
        std::unique_lock<std::mutex> lock(_mutex);
        while (!_cond_notify && _running)
            _cond_var.wait_for(lock, std::chrono::milliseconds(300));

        while ((len = _ring_buffer->get(_buff, _buff_size, pcm_frame))) {
            _frame->pts = pcm_frame.pts;
#if FFMPEG_VERSION_MAJOR >= 7
            // In FFmpeg 7+, pkt_pts is removed, use pts instead
            // _frame->pts is already set above
#else
            _frame->pkt_pts = pcm_frame.pkt_pts;
#endif
            _frame->pkt_dts = pcm_frame.pkt_dts;

            if ((error = encode(_frame, packet)) != AE_NO) {
                if (_on_error)
                    _on_error(error);

                al_fatal("read aac packet failed:%d", error);

                break;
            }
        }

        _cond_notify = false;
    }

    //flush pcm data in encoder
    encode(NULL, packet);

    av_packet_free(&packet);
#ifdef SAVE_AAC
    av_write_trailer(_aac_fmt_ctx);
#endif
}

void encoder_aac::cleanup()
{
    if (_encoder) {
        avcodec_close(_encoder_ctx);
    }

    _encoder = NULL;

    if (_encoder_ctx) {
        avcodec_free_context(&_encoder_ctx);
    }

    if (_frame)
        av_free(_frame);
    _frame = NULL;

    if (_buff)
        av_free(_buff);

    _buff = NULL;

    _encoder_ctx = NULL;

#ifdef SAVE_AAC
    if (_aac_fmt_ctx) {
        avio_closep(&_aac_fmt_ctx->pb);
        avformat_free_context(_aac_fmt_ctx);
    }
#endif
}
} // namespace am