#include "RawVideoReader.h"
#include <time.h>
#include <cstdio>
#define AV_OUTPUT_FORMAT "mpegts"
#define AV_OUTPUT_CODEC  "libx264"

#define AV_OUTPUT_BITRATE 40000000
#define AV_OUTPUT_THREADS 16
#define AV_OUTPUT_THREAD_TYPE FF_THREAD_SLICE

extern "C"{
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libswscale/swscale.h>
#include <libavutil/avassert.h>
}
typedef struct {
    AVFormatContext *formatCtx;
    AVCodecContext  *codecCtx;
    AVCodec         *encoder;
    AVStream        *outStream;
    AVFrame         *frame;
    AVProgram       *program;
    AVPacket        packet;
} ff_output_t;

ff_output_t ff_output;

int BytesPerFrame;

AVDictionary    *codecOptions = NULL;
AVInputFormat mpegts_mux;
uint8_t * picture_buf;


void convertFRame(Frame *f, AVFrame * ff_frame) {
    int size =(int) f->height * f->width;
    printf("Alocou\n");

    ff_frame->data[0] =(uint8_t*) f->data;
    ff_frame->data[1] = ff_frame->data[0] + size;
    ff_frame->data[2] = ff_frame->data[1] + size / 4;

    ff_frame->linesize[0] = f->width;
    ff_frame->linesize[1] = f->width / 2;
    ff_frame->linesize[2] = f->width / 2;

    ff_frame->height = (int) f->height;
    ff_frame->width = (int) f->width;
    ff_frame->format = AV_PIX_FMT_YUV420P;
}

void av_packet_rescale_ts(AVPacket *pkt, AVRational src_tb, AVRational dst_tb){
    if (pkt->pts != AV_NOPTS_VALUE)
        pkt->pts = av_rescale_q(pkt->pts, src_tb, dst_tb);
    if (pkt->dts != AV_NOPTS_VALUE)
        pkt->dts = av_rescale_q(pkt->dts, src_tb, dst_tb);
    if (pkt->duration > 0)
        pkt->duration = av_rescale_q(pkt->duration, src_tb, dst_tb);
    if (pkt->convergence_duration > 0)
        pkt->convergence_duration = av_rescale_q(pkt->convergence_duration, src_tb, dst_tb);
}

int main()
{

    char * videoOutputName = "1.ts";
    int width = 4096, heigth = 2304;
    //RawVideoReader reader(4096, 2304, 4, "CineMag005-70frames.rgba"); //RGBA32 - 4BPP
    //RawVideoReader reader(4096, 2304, 1.5, "CineMag005-26frames.nv12"); //NV12 - 1.5BPP
    RawVideoReader reader(4096, 2304, 1.5, "Cinemag005.yuv420p"); //YUV420P - 1.5BPP
    BytesPerFrame =  avpicture_get_size(AV_PIX_FMT_YUV420P, 4096, 2304);

    av_register_all();
    avformat_network_init();

    //Open OUTPUT

    if (avformat_alloc_output_context2(&ff_output.formatCtx, NULL, AV_OUTPUT_FORMAT, videoOutputName) < 0) {
        printf("could not create output context\n");
        return -1;
    }
    ff_output.encoder = avcodec_find_encoder_by_name(AV_OUTPUT_CODEC);

    if (ff_output.encoder == NULL) {
        printf("Codec %s not found..\n", AV_OUTPUT_CODEC);
        return -1;
    }

    ff_output.outStream = avformat_new_stream(ff_output.formatCtx, ff_output.encoder);
    if (ff_output.outStream == NULL) {
        printf("Could not create output stream\n");
        return -1;
    }

    ff_output.outStream->id = ff_output.formatCtx->nb_streams - 1;


    ff_output.codecCtx = avcodec_alloc_context3(ff_output.encoder);

    ff_output.codecCtx->codec_type = AVMEDIA_TYPE_VIDEO;
    ff_output.codecCtx->height = 2304;
    ff_output.codecCtx->width = 4096;
    ff_output.codecCtx->pix_fmt = AV_PIX_FMT_YUV420P;

    //Set custom BIT RATE and THREADs
    ff_output.codecCtx->bit_rate = AV_OUTPUT_BITRATE;
    ff_output.codecCtx->thread_count = AV_OUTPUT_THREADS;
    ff_output.codecCtx->thread_type = AV_OUTPUT_THREAD_TYPE;
    ff_output.codecCtx->bit_rate = AV_OUTPUT_BITRATE;

    ff_output.formatCtx->bit_rate = AV_OUTPUT_BITRATE;
    ff_output.codecCtx->bit_rate_tolerance = 0;
    ff_output.codecCtx->rc_max_rate = 0;
    ff_output.codecCtx->rc_buffer_size = 0;
    ff_output.codecCtx->gop_size = 40;
    ff_output.codecCtx->max_b_frames = 3;
    ff_output.codecCtx->b_frame_strategy = 1;
    ff_output.codecCtx->coder_type = 1;
    ff_output.codecCtx->me_cmp = 1;
    ff_output.codecCtx->me_range = 16;
    ff_output.codecCtx->qmin = 10;
    ff_output.codecCtx->qmax = 51;
    ff_output.codecCtx->scenechange_threshold = 40;
    ff_output.codecCtx->flags |= CODEC_FLAG_LOOP_FILTER;
    ff_output.codecCtx->me_method = ME_HEX;
    ff_output.codecCtx->me_subpel_quality = 5;
    ff_output.codecCtx->i_quant_factor = 0.71;
    ff_output.codecCtx->qcompress = 0.6;
    ff_output.codecCtx->max_qdiff = 4;

    //Set custo timebase for codec and streams
    ff_output.codecCtx->time_base.num = 1;
    ff_output.codecCtx->time_base.den = 24;

    ff_output.formatCtx->streams[0]->codec->width = ff_output.codecCtx->width;
    ff_output.formatCtx->streams[0]->codec->height = ff_output.codecCtx->height;
    ff_output.formatCtx->streams[0]->codec->pix_fmt = AV_PIX_FMT_YUV420P;
    ff_output.formatCtx->streams[0]->codec->bit_rate = AV_OUTPUT_BITRATE;

    ff_output.outStream->avg_frame_rate = {24,1};


    if (avcodec_open2(ff_output.codecCtx, ff_output.encoder, &codecOptions)) {
        printf("Could not open output codec...\n");
        return -1;
    }
    ff_output.formatCtx->iformat = &mpegts_mux;
    av_dump_format(ff_output.formatCtx, 0, videoOutputName, 0);

    if (avio_open(&ff_output.formatCtx->pb, videoOutputName, AVIO_FLAG_WRITE)) {
        printf("avio_open failed %s\n", videoOutputName);
        return -1;
    }
    int inc = 0, gotPacket;

    Frame* f;
    int valor = 0;

    while (1) {

        ff_output.frame = av_frame_alloc();
        f = reader.getFrame(++valor);

        convertFRame(f, ff_output.frame);

        ff_output.frame->pts = inc++;
        ff_output.packet.data = NULL;
        ff_output.packet.size = 0;
        av_init_packet(&ff_output.packet);

        int ret = avcodec_encode_video2(ff_output.codecCtx, &ff_output.packet, ff_output.frame, &gotPacket);

        if (gotPacket) {

            ff_output.packet.stream_index = 0;
            av_packet_rescale_ts(&ff_output.packet,
                                 ff_output.codecCtx->time_base,
                                 ff_output.outStream->time_base);

            if (av_interleaved_write_frame(ff_output.formatCtx, &ff_output.packet) < 0) {
                printf("Unable to write to output stream..\n");
                return 0;
            }

        }
        av_frame_free(&ff_output.frame);
    }

    return 0;
}