// Compile with:
// g++ -I <include dir> -o 3byte_startcode_test 3byte_startcode_test.cpp  -L <lib dir> -lx264 -lavformat -lavcodec -lavutil
// Add a -DDUMP_PACKETS to also see the encoded video packets

#ifndef __STDC_CONSTANT_MACROS
#define __STDC_CONSTANT_MACROS
#endif

extern "C" {

#include <stdint.h>
#include <x264.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/pixfmt.h>

} // extern "C"

#include <cstdio>
#include <cstddef>
#include <cstring>
#include <vector>
#include <iostream>
#include <sstream>
#include <stdexcept>

const unsigned int width = 320;
const unsigned int height = 240;
const unsigned int frame_count = 25;
const char filename[] = "test.avi";

class x264_context :
    public x264_param_t
{
public:
    x264_t* ctx;

    x264_context() : ctx(NULL)
    {
        x264_param_default_preset(this, "veryfast", "ssim");

        this->i_csp = X264_CSP_I420;
        this->vui.b_fullrange = 1;
        this->i_width = width;
        this->i_height = height;
        this->i_keyint_min = 5;
        this->i_keyint_max = 50;
        this->b_repeat_headers = 0;
        this->b_annexb = 1;

        // time_base the base time unit in which we provide timestamps - 1 ms
        this->i_timebase_num = 1;
        this->i_timebase_den = 1000;
        this->b_vfr_input = 1;

        this->i_frame_reference = 3;

        this->i_bframe = 0;
        this->i_bframe_pyramid = 0;
        this->b_cabac = 1;
        this->analyse.b_transform_8x8 = 1;
        this->analyse.i_weighted_pred = 1;

        // Set rate control parameters
        this->rc.i_rc_method = 1;
        this->rc.f_rf_constant = 23.0f;

        this->rc.f_qcompress = 0.7f;

        // Do not use any lookahead to reduce latency
        this->rc.i_lookahead = 0;
        this->i_sync_lookahead = 0;

        this->i_level_idc = -1;

        this->b_intra_refresh = 0;
        this->rc.b_mb_tree = 0;

        this->i_slice_max_size = 0;

        ctx = x264_encoder_open(this);
        if (!ctx)
            throw std::runtime_error("Failed to open libx264 encoder context");
    }

    ~x264_context()
    {
        x264_encoder_close(ctx);
    }
};

class raw_image :
    public x264_picture_t
{
public:
    std::vector< uint8_t > data;

    raw_image()
    {
        x264_picture_init(this);

        this->i_type = X264_TYPE_AUTO;

        this->img.i_csp = X264_CSP_I420;
        this->img.i_plane = 3;

        std::size_t luma_size = width * height;
        std::size_t chroma_size = luma_size / 4;
        data.resize(luma_size + 2 * chroma_size);

        this->img.plane[0] = &data[0];
        this->img.plane[1] = this->img.plane[0] + luma_size;
        this->img.plane[2] = this->img.plane[1] + chroma_size;

        this->img.i_stride[0] = width;
        this->img.i_stride[1] = this->img.i_stride[2] = width / 2;
    }
};

class av_format_context
{
public:
    AVFormatContext* ctx;
    bool file_header_written;

    av_format_context() : ctx(avformat_alloc_context()), file_header_written(false)
    {
        if (!ctx)
            throw std::runtime_error("Failed to allocate AVFormatContext");

        ctx->oformat = av_guess_format("avi", filename, NULL);
        if (!ctx->oformat)
            throw std::runtime_error("Could not guess output format");

        ctx->oformat->video_codec = CODEC_ID_NONE;
        ctx->video_codec_id = CODEC_ID_NONE;
        ctx->oformat->audio_codec = CODEC_ID_NONE;
        ctx->audio_codec_id = CODEC_ID_NONE;
    }

    ~av_format_context()
    {
        if (ctx)
        {
            for (unsigned int i = 0; i < ctx->nb_streams; ++i)
            {
                if (ctx->streams[i]->codec->internal) // this field is only set when the codec is opened
                    avcodec_close(ctx->streams[i]->codec);
            }

            if (ctx->iformat)
            {
                // Close input stream
                avformat_close_input(&ctx);
            }
            else
            {
                // Close output stream
                if (ctx->oformat && (ctx->oformat->flags & AVFMT_NOFILE) == 0 && (ctx->oformat->flags & AVFMT_FLAG_CUSTOM_IO) == 0 && ctx->pb)
                {
                    avio_close(ctx->pb);
                    ctx->pb = NULL;
                }
                avformat_free_context(ctx);
            }
        }
    }

    void open()
    {
        if (avio_open(&ctx->pb, filename, AVIO_FLAG_WRITE) < 0)
            throw std::runtime_error("Failed to open output media file");
    }

    void close()
    {
        if (file_header_written)
            av_write_trailer(ctx);
        if (ctx->pb)
        {
            avio_close(ctx->pb);
            ctx->pb = NULL;
        }
        file_header_written = false;
    }

    void write_file_header(const uint8_t* extradata, std::size_t extradata_size)
    {
        AVCodec* video_codec = avcodec_find_encoder(AV_CODEC_ID_H264);
        if (!video_codec)
            throw std::runtime_error("ffmpeg does not support h.264");

        ctx->oformat->video_codec = AV_CODEC_ID_H264;
        ctx->video_codec_id = AV_CODEC_ID_H264;

        AVStream* av_stream = avformat_new_stream(ctx, NULL);
        if (!av_stream)
            throw std::runtime_error("av_new_stream failed");
        av_stream->id = 0;

        av_stream->r_frame_rate.num = av_stream->r_frame_rate.den = 0;
        av_stream->time_base.num = 1;
        av_stream->time_base.den = 1000;

        AVCodecContext* av_video_codec_ctx = av_stream->codec;
        av_video_codec_ctx->codec_id = AV_CODEC_ID_H264;
        av_video_codec_ctx->width = width;
        av_video_codec_ctx->height = height;
        av_video_codec_ctx->pix_fmt = AV_PIX_FMT_YUV420P;

        if (extradata_size > 0)
        {
            uint8_t* data = (uint8_t*)av_malloc(extradata_size + FF_INPUT_BUFFER_PADDING_SIZE);
            if (!data)
                throw std::bad_alloc();
            std::memcpy(data, extradata, extradata_size);
            std::memset(data + extradata_size, 0, FF_INPUT_BUFFER_PADDING_SIZE);
            av_video_codec_ctx->extradata = data;
            av_video_codec_ctx->extradata_size = extradata_size;
        }

        av_video_codec_ctx->codec_type = AVMEDIA_TYPE_VIDEO;
        av_video_codec_ctx->flags |= CODEC_FLAG_GLOBAL_HEADER;
        av_video_codec_ctx->time_base.num = 1;
        av_video_codec_ctx->time_base.den = 1000;

        if (avformat_write_header(ctx, NULL) < 0)
            throw std::runtime_error("error writing header");

        file_header_written = true;
    }

    void write_frame(const uint8_t* payload, std::size_t size, bool is_keyframe, uint64_t pts, uint64_t dts)
    {
#if defined(DUMP_PACKETS)
        std::printf("Frame pts %lu, dts: %lu:", (unsigned long)pts, (unsigned long)dts);
        for (std::size_t i = 0; i < size; ++i)
        {
            std::printf(" %02x", (unsigned int)payload[i]);
        }
        std::printf("\n");
#endif

        AVPacket av_packet;
        av_init_packet(&av_packet);
        av_packet.stream_index = 0;

        // Package the frame and stream it out
        av_packet.flags &= ~AV_PKT_FLAG_KEY;
        if (is_keyframe)
            av_packet.flags |= AV_PKT_FLAG_KEY;

        av_packet.data = const_cast< uint8_t* >(payload);
        av_packet.size = size;

        av_packet.pts = pts;
        av_packet.dts = dts;

        int res = av_interleaved_write_frame(ctx, &av_packet);
        if (res != 0)
        {
            std::ostringstream strm;
            strm << "Failed to write video frame, error: " << res;
            throw std::runtime_error(strm.str());
        }
    }
};

int main(int, char*[])
{
    try
    {
        av_register_all();

        x264_context x264_ctx;
        raw_image picture_in;

        x264_picture_t picture_out = {};
        int nnal = 0;
        x264_nal_t* nal = NULL;

        if (x264_encoder_headers(x264_ctx.ctx, &nal, &nnal) < 0)
            throw std::runtime_error("Failed to encode h.264 parameter sets");

        std::vector< uint8_t > extradata, frame_buffer;
        for (int i = 0; i < nnal; ++i)
        {
            if (nal[i].i_type == NAL_SEI)
            {
                /* only one SEI NAL is expected */
                frame_buffer.insert(frame_buffer.end(), nal[i].p_payload, nal[i].p_payload + nal[i].i_payload);
            }
            else
            {
                extradata.insert(extradata.end(), nal[i].p_payload, nal[i].p_payload + nal[i].i_payload);
            }
        }

        av_format_context file;

        file.open();
        file.write_file_header(&extradata[0], extradata.size());

        uint64_t pts = 0;
        for (unsigned int i = 0; i < frame_count; ++i, pts += 25)
        {
            picture_in.i_pts = picture_in.i_dts = pts;
            picture_in.i_type = X264_TYPE_AUTO;

            int result = x264_encoder_encode(x264_ctx.ctx, &nal, &nnal, &picture_in, &picture_out);
            if (result < 0)
            {
                std::ostringstream strm;
                strm << "Failed to encode frame, error code: " << result;
                throw std::runtime_error(strm.str());
            }
            else if (result == 0)
            {
                continue;
            }

            // Append the first encoded frame to SEI, put the following frames to the start of the buffer
            frame_buffer.insert(frame_buffer.end(), nal[0].p_payload, nal[nnal - 1].p_payload + nal[nnal - 1].i_payload);

            file.write_frame(&frame_buffer[0], frame_buffer.size(), picture_out.b_keyframe != 0, picture_out.i_pts, picture_out.i_dts);

            frame_buffer.clear();
        }

        file.close();
    }
    catch (std::exception& e)
    {
        std::cout << "Failure: " << e.what() << std::endl;
        return 1;
    }

    return 0;
}