#include <iostream>
#include <vector>

extern "C" {
    #include <libavutil/hwcontext.h>
    #include <libavutil/hwcontext_cuda.h>
    #include <libavformat/avformat.h>
    #include <libavcodec/avcodec.h>
    #include <libavutil/error.h>
    #include <libavutil/opt.h>
    #include <libavutil/hwcontext.h>
    #include <libavutil/hwcontext_cuda.h>
}


// Cuda
#include <cuda.h>
#include <cudaGL.h>

int screenWidth = 1280;
int screenHeight = 720;

CUcontext* m_cuContext;
CUgraphicsResource* cuInpTexRes;
CUDA_MEMCPY2D_st* m_memCpyStruct;

// ffmpeg
AVCodecContext* c;
AVFrame* frame;
AVBufferRef* m_avBufferRefFrame;
AVBufferRef* m_avBufferRefDevice;
AVHWDeviceContext* hwDevContext;
AVCUDADeviceContext* cudaDevCtx;
AVHWFramesContext* frameCtxPtr;

//test
FILE* outfile;

std::string getDeviceName() {
    CUdevice cuDevice = 0;
    cuDeviceGet(&cuDevice, 0);
    char szDeviceName[80];
    cuDeviceGetName(szDeviceName, sizeof(szDeviceName), cuDevice);
    return szDeviceName;
}

int initFFmpeg() {
    CUresult res;
    CUcontext oldCtx;
    // Get device name
    cuInit(0);
    CUdevice cuDevice = 0;
    cuDeviceGet(&cuDevice, 0);
    char szDeviceName[80];
    cuDeviceGetName(szDeviceName, sizeof(szDeviceName), cuDevice);

    int ret = av_hwdevice_ctx_create(&m_avBufferRefDevice, AV_HWDEVICE_TYPE_CUDA, szDeviceName, NULL, NULL);

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

    hwDevContext = (AVHWDeviceContext*)(m_avBufferRefDevice->data);
    cudaDevCtx = (AVCUDADeviceContext*)(hwDevContext->hwctx);
    m_cuContext = &(cudaDevCtx->cuda_ctx);

    m_avBufferRefFrame = av_hwframe_ctx_alloc(m_avBufferRefDevice);
    //Setup some values before initialization
    frameCtxPtr = (AVHWFramesContext*)(m_avBufferRefFrame->data);
    frameCtxPtr->width = screenWidth;
    frameCtxPtr->height = screenHeight;
    frameCtxPtr->sw_format = AV_PIX_FMT_RGB0;
    frameCtxPtr->format = AV_PIX_FMT_CUDA;
    frameCtxPtr->device_ref = m_avBufferRefDevice;
    frameCtxPtr->device_ctx = (AVHWDeviceContext*)m_avBufferRefDevice->data;

    ret = av_hwframe_ctx_init(m_avBufferRefFrame);
    if (ret < 0) {
        return -1;
    }

    res = cuCtxPopCurrent_v2(&oldCtx); // THIS IS ALLOWED TO FAIL

    AVCodec* codec = avcodec_find_encoder_by_name("h264_nvenc");
    c = avcodec_alloc_context3(codec);

    //Assign some hardware accel specific data to AvCodecContext
    c->hw_device_ctx = m_avBufferRefDevice;//This must be done BEFORE avcodec_open2()
    c->pix_fmt = AV_PIX_FMT_CUDA; //Since this is a cuda buffer, although its really opengl with a cuda ptr
    c->hw_frames_ctx = m_avBufferRefFrame;
    c->codec_type = AVMEDIA_TYPE_VIDEO;
    c->sw_pix_fmt = AV_PIX_FMT_RGB0;

    /* resolution must be a multiple of two */
    c->width = screenWidth;
    c->height = screenHeight;
    c->time_base = { 1, 30 };
    c->framerate = { 30, 1 };
    c->bit_rate = 8 * 1000 * 1000;

    ret = avcodec_open2(c, codec, NULL);
    if (ret < 0) {
        return -1;
    }

    std::cout << "freeing context" << std::endl;
    avcodec_free_context(&c);
    std::cout << "Context freed" << std::endl;
    av_buffer_unref(&m_avBufferRefDevice);
    av_buffer_unref(&m_avBufferRefFrame);
    av_frame_free(&frame);
}


int main(int argc, char *argv[]) {
    if(initFFmpeg() == -1) {
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}