#define __STDC_CONSTANT_MACROS
#include <iostream>

#define WIDTH 100
#define HEIGHT 100

extern "C" {
    #include <libavdevice/avdevice.h>
    #include <libavformat/avformat.h>
    #include <libavutil/imgutils.h>
    #include <libavutil/pixfmt.h>
}

using namespace std;

void OpenImage(const char* imageFileName, AVFrame* pFrame) {
    AVFormatContext *pFormatCtx = avformat_alloc_context();

    if (avformat_open_input(&pFormatCtx, imageFileName, NULL, NULL)!=0) {
        printf("Can't open image file '%s'\n", imageFileName);
        return;
    }

    AVCodecContext *pCodecCtx;

    pCodecCtx = pFormatCtx->streams[0]->codec;
    pCodecCtx->width = WIDTH;
    pCodecCtx->height = HEIGHT;
    pCodecCtx->pix_fmt = PIX_FMT_RGB32;

    // Find the decoder for the video stream
    AVCodec *pCodec = avcodec_find_decoder(pCodecCtx->codec_id);
    if (!pCodec) {
        printf("Codec not found\n");
        return;
    }

    // Open codec
    if(avcodec_open(pCodecCtx, pCodec)<0) {
        printf("Could not open codec\n");
        return ;
    }

    if (!pFrame) {
        printf("Can't allocate memory for AVFrame\n");
        return ;
    }

    int frameFinished;
    int numBytes;

    // Determine required buffer size and allocate buffer
    numBytes = avpicture_get_size(PIX_FMT_RGB32, pCodecCtx->width, pCodecCtx->height);
    uint8_t *buffer = (uint8_t *) av_malloc(numBytes * sizeof(uint8_t));

    avpicture_fill((AVPicture *) pFrame, buffer, PIX_FMT_RGB32, pCodecCtx->width, pCodecCtx->height);

    // Read frame

    AVPacket packet;

    int framesNumber = 0;
    while (av_read_frame(pFormatCtx, &packet) >= 0) {
        if(packet.stream_index != 0)
            continue;

        int ret = avcodec_decode_video2(pCodecCtx, pFrame, &frameFinished, &packet);
        if (ret > 0) {
            printf("Frame is decoded, size %d\n", ret);
            pFrame->quality = 4;
        } else printf("Error [%d] while decoding frame: %s\n", ret, strerror(AVERROR(ret)));
    }
}

int main() {
    av_register_all();

    char* filename = "out.mpg";

    // Find the video encoder
    AVCodec* codec = avcodec_find_encoder(CODEC_ID_MPEG1VIDEO);
    if (!codec)
    {
        printf("Codec not found\n");
        return -1;
    }

    // Initialize codec.
    AVCodecContext* codecContext = avcodec_alloc_context3(codec);

    codecContext->width = WIDTH;
    codecContext->height = HEIGHT;
    codecContext->bit_rate = 400000;
    codecContext->time_base.num = 1;
    codecContext->time_base.den = 25;
    codecContext->gop_size = 10;
    codecContext->max_b_frames=1;
    codecContext->pix_fmt = PIX_FMT_YUV420P;
    codecContext->codec_id = CODEC_ID_MPEG1VIDEO;

    // Open the codec.
    if (avcodec_open2(codecContext, codec, NULL) < 0) {
        printf("could not open codec\n");
        return -2;
    }

    // Open the output file
    FILE* outputFile = fopen(filename, "wb");
    if (!outputFile) {
        printf("could not open %s\n", filename);
        return -3;
    }

    // Create buffer.
    int outbuf_size = WIDTH * HEIGHT;
    int actualSize = (( ( outbuf_size * 3 ) / 2 ) + 100) * sizeof(uint8_t);

    AVPacket packet;
    bool first = true;
    int out_size, got_output;

    AVFrame *frame = avcodec_alloc_frame();

    if (!frame) {
        printf("Can't allocate memory for AVFrame\n");
        return -4;
    }

    int i;
    int frames=25;
    for (i = 0; i < frames; i++) {
        av_init_packet(&packet);
        packet.data = NULL;
        packet.size = 0;
        fflush(stdout);

        OpenImage("test.png", frame);

        frame->pts = i;
        out_size = avcodec_encode_video2(codecContext, &packet, frame, &got_output);
        if (out_size < 0) {
            return -5;
        }
        if (got_output && got_output > 0) {
            printf("Write frame %3d (size=%5d)\n", i, packet.size);
            fwrite(packet.data, 1, packet.size, outputFile);
            av_free_packet(&packet);
        }
    }

    // Get the delayed frames
    for (got_output = 1; got_output; i++) {
        fflush(stdout);
        out_size = avcodec_encode_video2(codecContext, &packet, NULL, &got_output);
        if (out_size < 0) {
            return -6;
        }
        if (got_output) {
            printf("Write frame %3d (size=%5d)\n", i, packet.size);
            fwrite(packet.data, 1, packet.size, outputFile);
            av_free_packet(&packet);
        }
    }

    // add sequence end code to have a real mpeg file
    uint8_t endcode[] = { 0, 0, 1, 0xb7 };
    fwrite(endcode, 1, sizeof(endcode), outputFile);
    fclose(outputFile);
    avcodec_close(codecContext);
    av_free(codecContext);
    av_freep(&frame->data[0]);
    return 0;
}