#include <algorithm>
#include <iterator>
#include <math.h>

extern "C"
{
#include "libavcodec/avcodec.h"
#include "libavdevice/avdevice.h"
#include "libavfilter/avfilter.h"
#include "libavformat/avformat.h"
#include "libavutil/avutil.h"
#include "libavutil/imgutils.h"
#include "libswscale/swscale.h"
#include "libswresample/swresample.h"

#define RES_NOT_MUL_OF_TWO 1
#define COULD_NOT_FIND_VID_CODEC 2
#define CONTEXT_CREATION_ERROR 3
#define COULD_NOT_OPEN_VID_CODEC 4
#define COULD_NOT_OPEN_FILE 5
#define COULD_NOT_ALLOCATE_FRAME 6
#define COULD_NOT_ALLOCATE_PIC_BUF 7
#define ERROR_ENCODING_FRAME_SEND 8
#define ERROR_ENCODING_FRAME_RECEIVE 9
#define COULD_NOT_FIND_AUD_CODEC 10
#define COULD_NOT_OPEN_AUD_CODEC 11
#define COULD_NOT_ALL_RESMPL_CONTEXT 12
#define FAILED_TO_INIT_RESMPL_CONTEXT 13
#define COULD_NOT_ALLOC_SAMPLES 14
#define COULD_NOT_CONVERT_AUD 15
#define ERROR_ENCODING_SAMPLES_SEND 16
#define ERROR_ENCODING_SAMPLES_RECEIVE 17
#define ENCODED_VIDEO 18
#define ENCODED_AUDIO 19
#define ENCODED_AUDIO_AND_VIDEO 20

    AVCodec *aud_codec;
    AVCodecContext *aud_codec_context = NULL;
    AVFormatContext *outctx;
    AVStream *audio_st;
    AVFrame *aud_frame;

    int aud_frame_counter;

    int initialize_encoding_audio(const char *filename)
    {
        int ret;
        AVCodecID aud_codec_id = AV_CODEC_ID_AAC;
        AVSampleFormat sample_fmt = AV_SAMPLE_FMT_FLTP;

        avcodec_register_all();
        av_register_all();

        // Fixup audio codec
        aud_codec = avcodec_find_encoder(aud_codec_id);
        //avcodec_register(aud_codec);

        if (!aud_codec)
            return COULD_NOT_FIND_AUD_CODEC;

        aud_codec_context = avcodec_alloc_context3(aud_codec);
        if (!aud_codec_context)
            return CONTEXT_CREATION_ERROR;

        /* select other audio parameters supported by the encoder */
        aud_codec_context->bit_rate = 128000;
        aud_codec_context->sample_rate = 48000;
        aud_codec_context->sample_fmt = sample_fmt;
        aud_codec_context->channel_layout = AV_CH_LAYOUT_STEREO;
        aud_codec_context->channels = av_get_channel_layout_nb_channels(aud_codec_context->channel_layout);
        //aud_codec_context->profile = FF_PROFILE_AAC_MAIN;

        aud_codec_context->codec = aud_codec;
        aud_codec_context->codec_id = aud_codec_id;

        ret = avcodec_open2(aud_codec_context, aud_codec, NULL);

        if (ret < 0)
            return COULD_NOT_OPEN_AUD_CODEC;

        outctx = avformat_alloc_context();
        ret = avformat_alloc_output_context2(&outctx, NULL, "mp4", filename);

        outctx->audio_codec = aud_codec;
        outctx->audio_codec_id = aud_codec_id;

        audio_st = avformat_new_stream(outctx, aud_codec);
        outctx->streams = &audio_st;

        audio_st->codecpar->codec_id = aud_codec_id;
        audio_st->codecpar->codec_type = AVMEDIA_TYPE_AUDIO;
        audio_st->codecpar->codec_tag = aud_codec_context->codec_tag;
        //audio_st->codecpar->bits_per_coded_sample = aud_codec_context->bits_per_raw_sample;
        audio_st->codecpar->bits_per_raw_sample = aud_codec_context->bits_per_raw_sample;
        audio_st->codecpar->bit_rate = aud_codec_context->bit_rate;
        audio_st->codecpar->sample_rate = aud_codec_context->sample_rate;
        audio_st->codecpar->channels = aud_codec_context->channels;
        audio_st->codecpar->channel_layout = aud_codec_context->channel_layout;
        audio_st->codecpar->format = sample_fmt;
        audio_st->codecpar->frame_size = aud_codec_context->frame_size;
        audio_st->codecpar->block_align = aud_codec_context->block_align;
        audio_st->codecpar->initial_padding = aud_codec_context->initial_padding;

        av_dump_format(outctx, 0, filename, 1);

        if (!(outctx->oformat->flags & AVFMT_NOFILE))
        {
            if (avio_open(&outctx->pb, filename, AVIO_FLAG_WRITE) < 0)
                return COULD_NOT_OPEN_FILE;
        }

        ret = avformat_write_header(outctx, NULL);

        aud_frame = av_frame_alloc();
        aud_frame->nb_samples = aud_codec_context->frame_size;
        //aud_frame->sample_rate = aud_codec_context->sample_rate;
        aud_frame->format = aud_codec_context->sample_fmt;
        aud_frame->channel_layout = aud_codec_context->channel_layout;

        int buffer_size = av_samples_get_buffer_size(NULL, aud_codec_context->channels, aud_codec_context->frame_size,
            aud_codec_context->sample_fmt, 0);

        av_frame_get_buffer(aud_frame, buffer_size / 2);

        if (!aud_frame)
            return COULD_NOT_ALLOCATE_FRAME;

        aud_frame_counter = 0;

        return 0;
    }

    int encode_audio_samples(uint8_t **aud_samples)
    {
        int ret;

        double audio_time = audio_st ? av_stream_get_end_pts(audio_st) * av_q2d(audio_st->time_base) : 0.0;

        aud_frame->data[0] = aud_samples[0];
        aud_frame->data[1] = aud_samples[1];

        aud_frame->extended_data = aud_frame->data;

        aud_frame->pts = aud_frame_counter++;

        ret = avcodec_send_frame(aud_codec_context, aud_frame);
        if (ret < 0)
            return ERROR_ENCODING_SAMPLES_SEND;

        AVPacket pkt;
        av_init_packet(&pkt);
        pkt.data = NULL;
        pkt.size = 0;

        while (true)
        {
            ret = avcodec_receive_packet(aud_codec_context, &pkt);
            if (!ret)
            {
                av_packet_rescale_ts(&pkt, aud_codec_context->time_base, audio_st->time_base);

                pkt.stream_index = audio_st->index;
                av_write_frame(outctx, &pkt);
                av_packet_unref(&pkt);
            }
            if (ret == AVERROR(EAGAIN))
                break;
            else if (ret < 0)
                return ERROR_ENCODING_SAMPLES_RECEIVE;
            else
                break;
        }

        return 0;
    }

    int finish_audio_encoding()
    {
        AVPacket pkt;
        av_init_packet(&pkt);
        pkt.data = NULL;
        pkt.size = 0;

        fflush(stdout);

        int ret = avcodec_send_frame(aud_codec_context, NULL);
        if (ret < 0)
            return ERROR_ENCODING_FRAME_SEND;

        while (true)
        {
            ret = avcodec_receive_packet(aud_codec_context, &pkt);
            if (!ret)
            {
                av_packet_rescale_ts(&pkt, aud_codec_context->time_base, audio_st->time_base);

                pkt.stream_index = audio_st->index;
                av_write_frame(outctx, &pkt);
                av_packet_unref(&pkt);
            }
            if (ret == -AVERROR(AVERROR_EOF))
                break;
            else if (ret < 0)
                return ERROR_ENCODING_FRAME_RECEIVE;
        }

        av_write_trailer(outctx);
    }

    void cleanup()
    {
        if (aud_frame)
        {
            av_frame_free(&aud_frame);
        }
        if (outctx)
        {
            for (int i = 0; i < outctx->nb_streams; i++)
                av_freep(&outctx->streams[i]);

            avio_close(outctx->pb);
            av_free(outctx);
        }

        if (aud_codec_context)
        {
            avcodec_close(aud_codec_context);
            av_free(aud_codec_context);
        }
    }

    void get_audio_frame(float_t *left_samples, float_t *right_samples, int frame_size, float* t)
    {
        int j, i;
        float v;
        for (j = 0; j < frame_size; j++)
        {
            v = sin(*t) * 0.5;
            left_samples[j] = v;
            right_samples[j] = v;

            *t += (M_PI_2 * 220) / frame_size;
        }
    }

    int main()
    {
        int frame_rate = 30;
        float t = 0;

        initialize_encoding_audio("audio.m4a");

        int sec = 5;

        float_t** aud_samples;
        int src_samples_linesize;
        int src_nb_samples = 1024;
        int src_channels = 2;

        int ret = av_samples_alloc_array_and_samples((uint8_t***)&aud_samples, &src_samples_linesize, src_channels,
            src_nb_samples, AV_SAMPLE_FMT_FLTP, 0);

        for (size_t i = 0; i < frame_rate * sec; i++)
        {
            get_audio_frame(aud_samples[0], aud_samples[1], src_nb_samples, &t);

            encode_audio_samples((uint8_t **)aud_samples);
        }

        finish_audio_encoding();
        cleanup();

        return 0;
    }
}