#include "FFSaveImg.h"

extern "C"
{
	#include <libavcodec/avcodec.h> 
	#include <libavdevice/avdevice.h> 
	#include <libavformat/avformat.h> 
	#include <libavfilter/avfilter.h> 
	#include <libavutil/avutil.h> 
    #include <libavutil/pixdesc.h> 
	#include <libswscale/swscale.h>
    #include <libavutil/imgutils.h>
}

int saveJpg(AVFrame *pFrame, const char *out_file) {
    
    int width = pFrame->width;
    int height = pFrame->height;
    AVCodecContext *pCodecCtx = NULL;
    
    AVFormatContext *pFormatCtx = avformat_alloc_context();
    // 设置输出文件格式
    pFormatCtx->oformat = av_guess_format("mjpeg", NULL, NULL);

    // 创建并初始化输出AVIOContext
    if (avio_open(&pFormatCtx->pb, out_file, AVIO_FLAG_READ_WRITE) < 0) {
        printf("Couldn't open output file.");
        return -1;
    }

    // 构建一个新stream
    AVStream *video_st = avformat_new_stream(pFormatCtx, 0);
    if (video_st == NULL) {
        return -1;
    }

    pCodecCtx = video_st->codec;
	pCodecCtx->codec_id = pFormatCtx->oformat->video_codec;
	pCodecCtx->codec_type = AVMEDIA_TYPE_VIDEO;
	pCodecCtx->pix_fmt = AV_PIX_FMT_YUVJ420P;
	pCodecCtx->width = width;  
	pCodecCtx->height = height;

	pCodecCtx->time_base = (AVRational) {1, 25};

    AVCodec *pCodec = avcodec_find_encoder(pCodecCtx->codec_id);

    if (!pCodec) {
        printf("Could not find encoder\n");
        return -1;
    }

    if (avcodec_open2(pCodecCtx, pCodec, NULL) < 0) {
        printf("Could not open codec.");
        return -1;
    }

    int ret = avformat_write_header(pFormatCtx, NULL);
    if (ret < 0) {
        printf("write_header fail\n");
        return -1;
    }

    AVPacket* pkt = av_packet_alloc();
    av_init_packet(pkt);

    // 编码数据
    ret = avcodec_send_frame(pCodecCtx, pFrame);
    if (ret < 0) {
        printf("Could not avcodec_send_frame.");
        return -1;
    }

    // 得到编码后数据
    ret = avcodec_receive_packet(pCodecCtx, pkt);
    if (ret < 0) {
        printf("Could not avcodec_receive_packet");
        return -1;
    }

    ret = av_write_frame(pFormatCtx, pkt);

    if (ret < 0) {
        printf("Could not av_write_frame");
        return -1;
    }

    av_frame_free(&pFrame);
    av_packet_free(&pkt);

    //Write Trailer
    av_write_trailer(pFormatCtx);


    avcodec_close(pCodecCtx);
    avio_close(pFormatCtx->pb);
    avformat_free_context(pFormatCtx);

    return 0;
}

AVFrame* convert2yuv(int width, int height, unsigned char * rgbbuf, AVPixelFormat src_pix_fmt){

    AVPixelFormat dst_pix_fmt = AV_PIX_FMT_YUV420P;

    AVFrame *pFrameYUV = av_frame_alloc();
    uint8_t *out_buffer = new uint8_t[avpicture_get_size(dst_pix_fmt, width, height)];

    avpicture_fill((AVPicture *)pFrameYUV, out_buffer, dst_pix_fmt, width, height);

    AVFrame *rgbFrame = av_frame_alloc();

    avpicture_fill((AVPicture *)rgbFrame, rgbbuf, src_pix_fmt, width, height);

    SwsContext *sws_ctx = sws_getContext(
        width, height, src_pix_fmt,
        width, height, dst_pix_fmt,
        SWS_BILINEAR, NULL, NULL, NULL);

    sws_scale(sws_ctx, rgbFrame->data, rgbFrame->linesize, 0, height, pFrameYUV->data, pFrameYUV->linesize);
    sws_freeContext(sws_ctx);
    av_frame_free(&rgbFrame);

    pFrameYUV->width = width;
    pFrameYUV->height = height;
    pFrameYUV->format = dst_pix_fmt;

    return pFrameYUV;
}

void saveJpg(int width, int height, unsigned char * pData, string file_name){
    AVFrame *pFrame = convert2yuv(width, height, pData, AV_PIX_FMT_BGR24);
    saveJpg(pFrame, file_name.c_str());
}