/********************
 功能：將YUV編碼成JPG
 FFmpeg：4.0.2
 ********************/
#include <iostream>
extern "C"
{
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/avutil.h>
#include <libswscale/swscale.h>
#include <libavutil/imgutils.h>
 

};
using namespace std;

int EncodeYUVToJPEG(const char* InputFileName, const char* OutputFileName, int in_w, int in_h)
{
    FILE *in_file = NULL;                        
    fopen_s(&in_file, InputFileName, "rb");      //Input File

    //av_register_all();// 4.0.2已棄用
    /*
    //Method 1
    AVFormatContext *pFormatCtx = avformat_alloc_context();
    //Guess format
    AVOutputFormat *fmt = av_guess_format("mjpeg", NULL, NULL);
    pFormatCtx->oformat = fmt;
    //Output URL
    if (avio_open(&pFormatCtx->pb, out_file, AVIO_FLAG_READ_WRITE) < 0)
    {
        printf("Couldn't open output file.");
        return -1;
    }
    */
 

    //Method 2. More simple
    AVFormatContext *pFormatCtx = avformat_alloc_context();
    avformat_alloc_output_context2(&pFormatCtx, NULL, NULL, OutputFileName);
    AVOutputFormat *fmt = pFormatCtx->oformat;

    AVStream *video_st = avformat_new_stream(pFormatCtx, 0);
    if (video_st == NULL)
    {
        return
 
 -1;
    }
    // 獲取編解碼器上下文資訊
    AVCodecContext* pCodecCtx = avcodec_alloc_context3(NULL);
    //pCodecCtx = video_st->codec;//已棄用
    if (avcodec_parameters_to_context(pCodecCtx, video_st->codecpar) < 0)
    {
        cout << "Copy stream failed!" << endl;
        return -1;
    }
    pCodecCtx->codec_id = fmt->video_codec;
    pCodecCtx->codec_type = AVMEDIA_TYPE_VIDEO;
    pCodecCtx->pix_fmt = AV_PIX_FMT_YUVJ420P;
    pCodecCtx->width = in_w;
    pCodecCtx->height = in_h;
    pCodecCtx->time_base.num = 1;
    pCodecCtx->time_base.den = 25;

    //dump some information
    av_dump_format(pFormatCtx, 0, OutputFileName, 1);

    AVCodec *pCodec = avcodec_find_encoder(pCodecCtx->codec_id);
    if (!pCodec)
    {
        printf("Codec not found.");
        return -1;
    }
    if (avcodec_open2(pCodecCtx, pCodec, NULL) < 0)
    {
        printf("Could not open codec.");
        return -1;
    }
    /*4.0.2中已棄用
    AVFrame* pictureFrame = av_frame_alloc();
    size = avpicture_get_size(pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height);
    picture_buf = (uint8_t *)av_malloc(size);
    if (!picture_buf)
    {
    return -1;
    }
    avpicture_fill((AVPicture *)pictureFrame, picture_buf, pCodecCtx->pix_fmt, pCodecCtx->width, pCodecCtx->height);
    */
    AVFrame* pictureFrame = av_frame_alloc();
    pictureFrame->width = pCodecCtx->width;
    pictureFrame->height = pCodecCtx->height;
    pictureFrame->format = AV_PIX_FMT_YUV420P;
    int size = av_image_get_buffer_size(AV_PIX_FMT_YUV420P, pCodecCtx->width, pCodecCtx->height, 1);
    uint8_t *picture_buf = (uint8_t *)av_malloc(size);
    av_image_fill_arrays(pictureFrame->data, pictureFrame->linesize, picture_buf, AV_PIX_FMT_YUV420P, pCodecCtx->width, pCodecCtx->height, 1);

    //Write Header
    avformat_write_header(pFormatCtx, NULL);

    int y_size = pCodecCtx->width * pCodecCtx->height;
    AVPacket *pkt = av_packet_alloc();
    av_new_packet(pkt, y_size * 3);
    //Read YUV
    if (fread(picture_buf, 1, y_size * 3 / 2, in_file) <= 0)
    {
        printf("Could not read input file.");
        return -1;
    }
    pictureFrame->data[0] = picture_buf;                    // Y
    pictureFrame->data[1] = picture_buf + y_size;           // U 
    pictureFrame->data[2] = picture_buf + y_size * 5 / 4;   // V
    /////// Encode ///////////////////// 
    /*4.0.2中已棄用
    int ret = 0;
    int got_picture = 0;
    ret = avcodec_encode_video2(pCodecCtx, &pkt, pictureFrame, &got_picture);
    if (ret < 0) 
    {
        printf("Encode Error.\n");
        return -1;
    }
    if (got_picture == 1) 
    {
        pkt.stream_index = video_st->index;
        ret = av_write_frame(pFormatCtx, &pkt);
    }
    av_free_packet(&pkt);
    */
    int ret = avcodec_send_frame(pCodecCtx, pictureFrame);
    while (ret >= 0)
    {
        pkt->stream_index = video_st->index;
        ret = avcodec_receive_packet(pCodecCtx, pkt);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
        {
            return -1;
        }
        else if (ret < 0)
        {
            fprintf(stderr, "Error during encoding\n");
            exit(1);
        }
        av_write_frame(pFormatCtx, pkt);
    }
    av_packet_unref(pkt);
    /////// Encode //////////////// 
    //Write Trailer
    av_write_trailer(pFormatCtx);

    printf("Encode Successful.\n");

    if (video_st)
    {
        //avcodec_close(video_st->codec);
        av_free(pictureFrame);
        av_free(picture_buf);
    }
    //avio_close(pFormatCtx->pb);//Method 1
    avformat_free_context(pFormatCtx);
    fclose(in_file);
    return 0;
}
int main()
{
    const char* OutputFileName = "../JPG/test.jpg";
    const char* InputFileName = "../YUV/test.yuv";
    int in_w = 352, in_h = 288;
    EncodeYUVToJPEG(InputFileName, OutputFileName, in_w, in_h);
    return getchar();
}