#include <stdio.h>
#include <errno.h>
#include <stdlib.h>

extern "C" {
#include <libavutil/hwcontext.h>
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/buffer.h>
#include <libavfilter/buffersink.h>
#include <libavfilter/buffersrc.h>
#include <libavutil/opt.h>
#include <libavutil/time.h>
#include <libavcodec/packet.h>
#include <libavfilter/avfilter.h>
#include <libavutil/avutil.h>
#include <libavutil/error.h>
#include <libavutil/frame.h>
#include <libavutil/log.h>
#include <libavutil/pixfmt.h>
#include <libavutil/rational.h>
}

#include <string>
#include <spdlog/spdlog.h>

static AVFormatContext *dFmtCtx = NULL;
static AVFormatContext *eFmtCtx = NULL;

static AVBufferRef *hwDeviceCtx = NULL;

static AVCodecContext *dvCodecCtx = NULL;
static AVCodecContext *evCodecCtx = NULL;

static AVCodecContext *daCodecCtx = NULL;
static AVCodecContext *eaCodecCtx = NULL;

static const AVCodec *dvCodec = NULL;
static const AVCodec *evCodec = NULL;

static const AVCodec *daCodec = NULL;
static const AVCodec *eaCodec = NULL;

AVFilterContext *bufferSinkCtx = NULL;
AVFilterContext *bufferSrcCtx = NULL;
AVFilterGraph *filterGraph = NULL;

static AVStream *evStream = NULL;

AVFrame *frame = NULL;
AVFrame *filterFrame = NULL;

int videoStream = -1;
int audioStream = -1;

int eInitialized = 0;
uint64_t frameCounter = av_gettime();

static const std::string av_make_error_string(int errnum)
{
    char errbuf[AV_ERROR_MAX_STRING_SIZE];
    av_strerror(errnum, errbuf, AV_ERROR_MAX_STRING_SIZE);
    return (std::string)errbuf;
}

static void ffmpeg_log(void *avcl, int level, const char *fmt, va_list args)
{
    char logBuff[1024];
    int prefix = 1;

    if (level == AV_LOG_ERROR || level == AV_LOG_PANIC || level == AV_LOG_FATAL) {
        int rc = av_log_format_line2(avcl, level, fmt, args, logBuff, 1024, &prefix);
        if (rc > 0) {
            spdlog::error("FFMPEG internal error: {}", logBuff);
        }
    }
}

static void initVideoDevEnc()
{
    dFmtCtx = NULL;
    eFmtCtx = NULL;

    evCodecCtx = NULL;
    dvCodecCtx = NULL;

    eaCodecCtx = NULL;
    daCodecCtx = NULL;

    daCodec = NULL;
    dvCodec = NULL;

    eaCodec = NULL;
    evCodec = NULL;

    videoStream = -1;
    audioStream = -1;

}

static enum AVPixelFormat get_cuda_format(AVCodecContext *ctx,
                                          const enum AVPixelFormat *pixFmts)
{
    const enum AVPixelFormat *p;

    for (p = pixFmts; *p != AV_PIX_FMT_NONE; ++p) {
        if (*p == AV_PIX_FMT_CUDA) {
            return *p;
        }
    }

    fprintf(stderr, "Unable to decode this file using CUDA.\n");
    return AV_PIX_FMT_NONE;
}

static int init_filters(const char *filtersDescr) {

    int ret = 0;

    const AVFilter *bufferSrc  = avfilter_get_by_name("buffer");
    const AVFilter *bufferSink = avfilter_get_by_name("buffersink");
    AVFilterInOut *outputs = avfilter_inout_alloc();
    AVFilterInOut *inputs  = avfilter_inout_alloc();
    enum AVPixelFormat pixFmts[] = { AV_PIX_FMT_CUDA, AV_PIX_FMT_YUV420P, AV_PIX_FMT_NONE };

    AVRational time_base = dFmtCtx->streams[videoStream]->time_base;

    filterGraph = avfilter_graph_alloc();

    if (!outputs || !inputs || !filterGraph) {
        ret = AVERROR(ENOMEM);

        avfilter_inout_free(&inputs);
        avfilter_inout_free(&outputs);

        return ret;
    }

    spdlog::debug("Init filter with input format: {} {}", dvCodecCtx->pix_fmt, dvCodecCtx->sw_pix_fmt);

    char args[512];
    snprintf(args, sizeof(args),
            "video_size=%dx%d:pix_fmt=%d:time_base=%d/%d:pixel_aspect=%d/%d",
            dvCodecCtx->width, dvCodecCtx->height, AV_PIX_FMT_CUDA,
            time_base.num, time_base.den,
            dvCodecCtx->sample_aspect_ratio.num, dvCodecCtx->sample_aspect_ratio.den);

    ret = avfilter_graph_create_filter(&bufferSrcCtx, bufferSrc, "in",
                                       args, NULL, filterGraph);

    bufferSrcCtx->hw_device_ctx = av_buffer_ref(hwDeviceCtx);

    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "Cannot create buffer source\n");

        avfilter_inout_free(&inputs);
        avfilter_inout_free(&outputs);

        return ret;
    }

     /* buffer video sink: to terminate the filter chain. */
    ret = avfilter_graph_create_filter(&bufferSinkCtx, bufferSink, "out",
                                       NULL, NULL, filterGraph);
    if (ret < 0) {
        av_log(NULL, AV_LOG_ERROR, "Cannot create buffer sink\n");

        avfilter_inout_free(&inputs);
        avfilter_inout_free(&outputs);

        return ret;
    }

    bufferSinkCtx->hw_device_ctx = av_buffer_ref(hwDeviceCtx);

    ret = av_opt_set_int_list(bufferSinkCtx, "pix_fmts", pixFmts,
                              AV_PIX_FMT_NONE, AV_OPT_SEARCH_CHILDREN);
    if (ret < 0) {
        spdlog::error("Cannot set output pixel format");

        avfilter_inout_free(&inputs);
        avfilter_inout_free(&outputs);

        return ret;
    }

    outputs->name       = av_strdup("in");
    outputs->filter_ctx = bufferSrcCtx;
    outputs->pad_idx    = 0;
    outputs->next       = NULL;

    inputs->name       = av_strdup("out");
    inputs->filter_ctx = bufferSinkCtx;
    inputs->pad_idx    = 0;
    inputs->next       = NULL;

    if ((ret = avfilter_graph_parse_ptr(filterGraph, filtersDescr,
                                    &inputs, &outputs, NULL)) < 0){
        avfilter_inout_free(&inputs);
        avfilter_inout_free(&outputs);

        return ret;
    }

    if ((ret = avfilter_graph_config(filterGraph, NULL)) < 0) {

        avfilter_inout_free(&inputs);
        avfilter_inout_free(&outputs);

        return ret;
    }

    avfilter_inout_free(&inputs);
    avfilter_inout_free(&outputs);

    return ret;
}

static int open_input(const char *input, const int videoStreamRequest)
{
    int ret;

    if ((ret = avformat_open_input(&dFmtCtx, input, NULL, NULL)) < 0) {
        fprintf(stderr, "Cannot open input '%s', Error code: %s\n", input, av_make_error_string(ret).c_str());
        return ret;
    }

    if ((ret = avformat_find_stream_info(dFmtCtx, NULL)) < 0) {
        fprintf(stderr, "Cannot find input stream information. Error code: %s", av_make_error_string(ret).c_str());
        return ret;
    }

    spdlog::debug("Start find video stream");
    spdlog::debug("Input programs: %d", dFmtCtx->nb_programs);

    for (int i = 0; i < dFmtCtx->nb_programs; ++i) {
        AVProgram *program = dFmtCtx->programs[i];

        spdlog::debug("AVProgram program number: %d", program->program_num);

        for (int j = 0; j < program->nb_stream_indexes; ++j) {
            spdlog::debug("-- AVProgram program {} Stream Index: {}", program->program_num, program->stream_index[j]);
        }
    }

    for (int i = 0; /*(audioStream == -1
                    || videoStream == -1)
                    &&*/ i < dFmtCtx->nb_streams; ++i) {
        const AVStream *stream = dFmtCtx->streams[i];

        if (stream->codecpar->codec_type == AVMEDIA_TYPE_VIDEO && stream->id == videoStreamRequest) {
            spdlog::debug("Found Video stream with ID: {0:x}", stream->id);
            spdlog::debug("-- Stream Index: {}", stream->index);
            spdlog::debug("-- Side data count: {}", stream->nb_side_data);
            spdlog::debug("-- Pixel format: {}", stream->codecpar->format);
            spdlog::debug("-- Metadata count: {}", av_dict_count(stream->metadata));

            /* AVDictionaryEntry * av_dict_get */

            videoStream = i;
        } else if (stream->codecpar->codec_type == AVMEDIA_TYPE_AUDIO) {

            spdlog::debug("Found Audio stream with ID: {0:x}", stream->id);
            spdlog::debug("-- Stream Index: {}", stream->index);
            spdlog::debug("-- Side data count: {}", stream->nb_side_data);
            spdlog::debug("-- Metadata count: {}", av_dict_count(stream->metadata));

            audioStream = i;
        }
    }

    fprintf(stdout, "Found Video stream index '%d' and audio stream index '%d'\n", videoStream, audioStream);

    if (videoStream == -1) {
        fprintf(stderr, "Cannot find a video stream in the input.\n");
        return -1;
    }

    dvCodec = avcodec_find_decoder_by_name("h264");
    if (dvCodec == NULL) {
        fprintf(stderr, "Cannot find codec for decompress content.\n");

        return -1;
    }

    printf("Found video stream\n");

    if (!(dvCodecCtx = avcodec_alloc_context3(dvCodec))) {
        return AVERROR(ENOMEM);
    }

    AVStream *video = dFmtCtx->streams[videoStream];

    spdlog::debug("Copy parameters to code context: {}", video->codecpar->codec_id);
    spdlog::info("Codec to be used to decoding: {}", dvCodec->long_name);

    if (!(dvCodecCtx = avcodec_alloc_context3(dvCodec)))
        return AVERROR(ENOMEM);


    if ((ret = avcodec_parameters_to_context(dvCodecCtx, video->codecpar)) < 0) {
        fprintf(stderr, "Parameters to context error. Error code: %s\n", av_make_error_string(ret).c_str());
        return ret;
    }

    spdlog::debug("Set Decoding context hardware device.");

    dvCodecCtx->hw_device_ctx = av_buffer_ref(hwDeviceCtx);
    if (!dvCodecCtx->hw_device_ctx) {
        fprintf(stderr, "A hardware device reference create failed.\n");
        return AVERROR(ENOMEM);
    }
    dvCodecCtx->get_format = get_cuda_format;

    if ((ret = avcodec_open2(dvCodecCtx, dvCodec, NULL)) < 0) {
        spdlog::error("Failed to open codec for decoding. Description: {}", av_make_error_string(ret));
    }

    return ret;
}

static int encode_write (AVPacket *packet, AVFrame *frame)
{
    int ret = 0;

    av_packet_unref(packet);

    if ((ret = avcodec_send_frame(evCodecCtx, frame)) < 0) {
        fprintf(stderr, "Error during encoding. Error: %s\n", av_make_error_string(ret).c_str());

        return (ret == AVERROR_EOF) ? 0 : ret;
    }

    while (1) {
        ret = avcodec_receive_packet(evCodecCtx, packet);
        if (ret) {
            break;
        }

        AVStream *encStream = eFmtCtx->streams[0];

        packet->pts = av_rescale_q(frameCounter, evCodecCtx->time_base, evStream->time_base);
        packet->dts = AV_NOPTS_VALUE;
        packet->stream_index = 0;

        frameCounter++;

        av_packet_rescale_ts(packet, dFmtCtx->streams[videoStream]->time_base, encStream->time_base);

        ret = av_write_frame(eFmtCtx, packet);
        if (ret < 0) {
            fprintf(stderr, "Error writing data to output. Error: %s\n", av_make_error_string(ret).c_str());
            return -1;
        }
    }

    if (ret == AVERROR_EOF) {
        return 0;
    }

    ret = ((ret == AVERROR(EAGAIN)) ? 0 : -1);

    return ret;
}

static int v_dec_enc (AVPacket *packet)
{
    int ret = avcodec_send_packet(dvCodecCtx, packet);
    if (ret < 0) {
        fprintf(stderr, "Can't send packet to decoder process: %s\n", av_make_error_string(ret).c_str());
        return ret;
    }


    while (ret >= 0) {

        ret = avcodec_receive_frame(dvCodecCtx, frame);

        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            av_frame_unref(frame);
            av_frame_unref(filterFrame);

            return 0;
        } else if (ret < 0) {
            fprintf(stderr, "Error while decoding. Error code: %s\n", av_make_error_string(ret).c_str());
            av_frame_unref(frame);
            av_frame_unref(filterFrame);

            return ret;
        }

         /* push the decoded frame into the filtergraph */
        if (av_buffersrc_add_frame_flags(bufferSrcCtx, frame, AV_BUFFERSRC_FLAG_KEEP_REF) < 0) {
            av_log(NULL, AV_LOG_ERROR, "Error while feeding the filtergraph\n");
            break;
        }

        if (!eInitialized) {
            evCodecCtx->hw_frames_ctx = av_buffer_ref(dvCodecCtx->hw_frames_ctx);
            if (!evCodecCtx->hw_frames_ctx) {
                fprintf(stderr, "Can't initialize hardware encoding on encoder context\n");
                ret = AVERROR(ENOMEM);

                av_frame_unref(frame);
                av_frame_unref(filterFrame);

                return ret;
            }

            evCodecCtx->time_base = av_inv_q(dvCodecCtx->framerate);
            evCodecCtx->pix_fmt = AV_PIX_FMT_CUDA;

            //evCodecCtx->width = dvCodecCtx->width;
            //evCodecCtx->height = dvCodecCtx->height;

            evCodecCtx->width = 100;
            evCodecCtx->height = 100;


            if ((ret = avcodec_open2(evCodecCtx, evCodec, NULL)) < 0) {
                fprintf(stderr, "Failed to open encode codec. Error codes: %s\n", av_make_error_string(ret).c_str());

                av_frame_unref(frame);
                av_frame_unref(filterFrame);

                return ret;
            }

            if (!(evStream = avformat_new_stream(eFmtCtx, evCodec))) {
                fprintf(stderr, "Failed to create video out stream\n");

                av_frame_unref(frame);
                av_frame_unref(filterFrame);

                return AVERROR(ENOMEM);
            }

            evStream->time_base = evCodecCtx->time_base;
            ret = avcodec_parameters_from_context(evStream->codecpar, evCodecCtx);
            if (ret < 0) {
                fprintf(stderr, "Failed to copy the stream parameters. Error: %s\n", av_make_error_string(ret).c_str());

                av_frame_unref(frame);
                av_frame_unref(filterFrame);

                return ret;
            }

            if ((ret = avformat_write_header(eFmtCtx, NULL)) < 0) {
                fprintf(stderr, "Error while writing stream header. Error: %s\n", av_make_error_string(ret).c_str());

                av_frame_unref(frame);
                av_frame_unref(filterFrame);

                return ret;
            }

            eInitialized = 1;
        }

        /* pull filtered frames from the filtergraph */
        while (1) {
            ret = av_buffersink_get_frame(bufferSinkCtx, filterFrame);

            if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
                ret = 0;
                break;
            }

            if (ret < 0) {
                break;
            }

            spdlog::debug("Frames format: {} {} {}", AV_PIX_FMT_CUDA, frame->format, filterFrame->format);

            if ((ret = encode_write(packet, frame)) < 0) {
                fprintf(stderr, "Error during encoding and write\n");
            }

            av_frame_unref(filterFrame);
        }

        av_frame_unref(frame);
        av_frame_unref(filterFrame);

        if (ret < 0) {
            return ret;
        }
    }

    return ret;
}

int main(int argc, char **argv)
{
    spdlog::set_level(spdlog::level::debug);

    av_log_set_callback(ffmpeg_log);

    int rc = av_hwdevice_ctx_create(&hwDeviceCtx, AV_HWDEVICE_TYPE_CUDA, NULL, NULL, 0);
    if (rc < 0) {
        fprintf(stderr, "Failed to create a CUDA device. Error code: %s\n", av_make_error_string(rc).c_str());

        return EXIT_FAILURE;
    }

    spdlog::debug("Hardware Context created");


    if (!(frame = av_frame_alloc()) || !(filterFrame = av_frame_alloc())) {
        return AVERROR(ENOMEM);
    }

    AVPacket *packet = av_packet_alloc();
    if (!packet) {
        fprintf(stderr, "Can't create new packet for decoding\n");

        return EXIT_FAILURE;
    }

    rc = open_input(argv[1], atoi(argv[2]));

     if ((rc = init_filters("hwupload_cuda,yadif_cuda,scale_cuda=w=100:h=100")) < 0 ) {
        return EXIT_FAILURE;
    }

    if (!(evCodec = avcodec_find_encoder_by_name("h264_nvenc"))) {
        fprintf(stderr, "Could not find encoder h264_nvenc\n");

        return EXIT_FAILURE;
    }

    if ((rc = avformat_alloc_output_context2(&eFmtCtx, NULL, "mpegts", NULL)) < 0) {
        fprintf(stderr, "Failed to get output format. Error code: %s\n", av_make_error_string(rc).c_str());

        return EXIT_FAILURE;
    }

    if (!(evCodecCtx = avcodec_alloc_context3(evCodec))) {
        return EXIT_FAILURE;
    }

    evCodecCtx->time_base = dvCodecCtx->time_base;
    evCodecCtx->framerate = dvCodecCtx->framerate;

    fprintf(stdout, "Starting to decode-encode\n");

    int ret = avio_open(&eFmtCtx->pb, argv[3], AVIO_FLAG_WRITE);
    if (ret < 0) {
        fprintf(stderr, "Cannot open output stream '%s': Error code: %s\n", argv[3], av_make_error_string(ret).c_str());

        return EXIT_FAILURE;
    }

    spdlog::info("Start decoding packets");

    /* read all packets and only transcoding video */
    while (ret >= 0) {
        if ((ret = av_read_frame(dFmtCtx, packet)) < 0) {
            break;
        }

        if (videoStream == packet->stream_index) {
            ret = v_dec_enc(packet);
        }

        av_packet_unref(packet);
    }

    fprintf(stdout, "Finished to decode-encode\n");

    av_packet_unref(packet);

    av_frame_free(&frame);
    av_frame_free(&filterFrame);

    return EXIT_SUCCESS;
}