#include "config.h"
#include "cuda_helper.hpp"
#include "video_encoder.h"

#include <nvEncodeAPI.h>

bool check_nvenc_api_call(NVENCSTATUS api_ret, unsigned int line_number,
                          const char *file_name, const char *api_call_str) {
    if (api_ret == NV_ENC_SUCCESS) [[likely]] return true;
    SPDLOG_ERROR("NvEnc api call {} failed at {}:{} with error 0x{:x}.",
                 api_call_str, file_name, line_number, api_ret);
    RET_ERROR;
}

#define NVENC_API_CHECK(api_call) \
    if (!check_nvenc_api_call( \
        api_call, __LINE__, __FILE__, #api_call)) [[unlikely]] \
        return false

struct video_encoder::impl {

    static constexpr auto frame_buffer_type = NV_ENC_BUFFER_FORMAT_ARGB;

    NV_ENCODE_API_FUNCTION_LIST api = {NV_ENCODE_API_FUNCTION_LIST_VER};

    CUcontext cuda_ctx = nullptr;
    void *encoder = nullptr;
    NV_ENC_OUTPUT_PTR output_buf = nullptr;

    int frame_width = image_width * 2, frame_height = image_height;
    int frame_pitch = frame_width * 4; // ARGB image
    int frame_rate = default_camera_fps;
    int frame_bitrate = default_video_stream_bitrate;

    // frame related
    void *frame_ptr = nullptr, **output_ptr = nullptr;
    size_t *output_size = nullptr;
    NV_ENC_REGISTERED_PTR frame_reg_ptr = nullptr;

    bool idr_flag = false;

    bool initialize() {
        NVENC_API_CHECK(NvEncodeAPICreateInstance(&api));
        CUDA_API_CHECK(cuCtxGetCurrent(&cuda_ctx));
        return true;
    }

    bool start_encode() {
        // constant params
        auto codec_guid = NV_ENC_CODEC_HEVC_GUID;
        auto preset_guid = NV_ENC_PRESET_P3_GUID;
//        auto tuning_info = NV_ENC_TUNING_INFO_LOW_LATENCY;
        auto tuning_info = NV_ENC_TUNING_INFO_ULTRA_LOW_LATENCY;

        // create encoder
        NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS session_params = {NV_ENC_OPEN_ENCODE_SESSION_EX_PARAMS_VER};
        session_params.deviceType = NV_ENC_DEVICE_TYPE_CUDA;
        session_params.device = cuda_ctx;
        session_params.apiVersion = NVENCAPI_VERSION;
        NVENC_API_CHECK(api.nvEncOpenEncodeSessionEx(&session_params, &encoder));

        // get preset config
        NV_ENC_PRESET_CONFIG preset_config = {NV_ENC_PRESET_CONFIG_VER, {NV_ENC_CONFIG_VER}};
        NVENC_API_CHECK(api.nvEncGetEncodePresetConfigEx(
                encoder, codec_guid, preset_guid, tuning_info, &preset_config));
        auto &encode_config = preset_config.presetCfg;
        encode_config.gopLength = NVENC_INFINITE_GOPLENGTH;
        encode_config.frameIntervalP = 1;
        auto &rc_params = encode_config.rcParams;
        rc_params.rateControlMode = NV_ENC_PARAMS_RC_CBR;
        rc_params.averageBitRate = frame_bitrate;
        rc_params.enableAQ = true;
        rc_params.multiPass = NV_ENC_TWO_PASS_QUARTER_RESOLUTION;
        // TODO; fine tune encoder config

        // start_encode encoder
        NV_ENC_INITIALIZE_PARAMS init_params = {NV_ENC_INITIALIZE_PARAMS_VER};
        init_params.encodeGUID = codec_guid;
        init_params.presetGUID = preset_guid;
        init_params.encodeWidth = frame_width;
        init_params.encodeHeight = frame_height;
        init_params.darWidth = frame_width; // TODO; learn more about this
        init_params.darHeight = frame_height; // TODO; learn more about this
        init_params.frameRateNum = frame_rate;
        init_params.frameRateDen = 1;
        init_params.enablePTD = 1;
        init_params.encodeConfig = &preset_config.presetCfg;
        init_params.maxEncodeWidth = frame_width;
        init_params.maxEncodeHeight = frame_height;
        init_params.tuningInfo = tuning_info;
        init_params.bufferFormat = frame_buffer_type;
        NVENC_API_CHECK(api.nvEncInitializeEncoder(encoder, &init_params));

        // create output buffer
        NV_ENC_CREATE_BITSTREAM_BUFFER buffer_config = {NV_ENC_CREATE_BITSTREAM_BUFFER_VER};
        NVENC_API_CHECK(api.nvEncCreateBitstreamBuffer(encoder, &buffer_config));
        output_buf = buffer_config.bitstreamBuffer;

        SPDLOG_INFO("Video encoder started.");
        return true;
    }

    void cleanup() {
        if (encoder == nullptr) return;

        // notify the end of stream
        NV_ENC_PIC_PARAMS pic_params = {NV_ENC_PIC_PARAMS_VER};
        pic_params.encodePicFlags = NV_ENC_PIC_FLAG_EOS;
        api.nvEncEncodePicture(encoder, &pic_params);

        // releasing resources
        if (frame_ptr != nullptr) {
            unregister_frame_ptr();
            frame_ptr = nullptr;
        }
        api.nvEncDestroyBitstreamBuffer(encoder, output_buf);

        // close encoder
        api.nvEncDestroyEncoder(encoder);
        encoder = nullptr;

        SPDLOG_INFO("Video encoder stopped.");
    }

    bool register_frame_ptr() {
        NV_ENC_REGISTER_RESOURCE reg_params = {NV_ENC_REGISTER_RESOURCE_VER};
        reg_params.resourceType = NV_ENC_INPUT_RESOURCE_TYPE_CUDADEVICEPTR;
        reg_params.width = frame_width;
        reg_params.height = frame_height;
        reg_params.pitch = frame_pitch;
        reg_params.resourceToRegister = (void *) frame_ptr;
        reg_params.bufferFormat = frame_buffer_type;
        reg_params.bufferUsage = NV_ENC_INPUT_IMAGE;
        NVENC_API_CHECK(api.nvEncRegisterResource(encoder, &reg_params));
        frame_reg_ptr = reg_params.registeredResource;
        return true;
    }

    bool unregister_frame_ptr() {
        if (frame_reg_ptr == nullptr) return true;
        NVENC_API_CHECK(api.nvEncUnregisterResource(encoder, frame_reg_ptr));
        frame_reg_ptr = nullptr;
        return true;
    }

    bool encode_frame() {
        // map input resource
        NV_ENC_MAP_INPUT_RESOURCE map_params = {NV_ENC_MAP_INPUT_RESOURCE_VER};
        map_params.registeredResource = frame_reg_ptr;
        NVENC_API_CHECK(api.nvEncMapInputResource(encoder, &map_params));
        assert(map_params.mappedBufferFmt == frame_buffer_type);

        // encode frame
        NV_ENC_PIC_PARAMS pic_params = {NV_ENC_PIC_PARAMS_VER};
        pic_params.inputWidth = frame_width;
        pic_params.inputHeight = frame_height;
        pic_params.inputPitch = frame_pitch;
        if (idr_flag) { // request for IDR frame
            pic_params.encodePicFlags = NV_ENC_PIC_FLAG_FORCEIDR | NV_ENC_PIC_FLAG_OUTPUT_SPSPPS;
            idr_flag = false;
        } else {
            pic_params.encodePicFlags = 0;
        }
        pic_params.inputBuffer = map_params.mappedResource;
        pic_params.outputBitstream = output_buf;
        pic_params.bufferFmt = frame_buffer_type;
        pic_params.pictureStruct = NV_ENC_PIC_STRUCT_FRAME; // TODO; learn more about this
        NVENC_API_CHECK(api.nvEncEncodePicture(encoder, &pic_params));

        // get encoded bitstream
        NV_ENC_LOCK_BITSTREAM lock_config = {NV_ENC_LOCK_BITSTREAM_VER};
        lock_config.doNotWait = false; // block until encode completed.
        lock_config.outputBitstream = output_buf;
        NVENC_API_CHECK(api.nvEncLockBitstream(encoder, &lock_config));

        // copy bitstream
        assert(output_ptr != nullptr);
        assert(output_size != nullptr);
        *output_ptr = malloc(lock_config.bitstreamSizeInBytes);
        *output_size = lock_config.bitstreamSizeInBytes;
        memcpy(*output_ptr, lock_config.bitstreamBufferPtr, lock_config.bitstreamSizeInBytes);

        // cleanup
        NVENC_API_CHECK(api.nvEncUnlockBitstream(encoder, output_buf));
        NVENC_API_CHECK(api.nvEncUnmapInputResource(encoder, map_params.mappedResource));

        return true;
    }

};

video_encoder::video_encoder()
        : pimpl(std::make_unique<impl>()) {}

video_encoder::~video_encoder() {
    stop_encode();
};

bool video_encoder::initialize() {
    return pimpl->initialize();
}

bool video_encoder::start_encode(int width, int height, int fps, int bitrate) {
    pimpl->frame_width = width;
    pimpl->frame_height = height;
    pimpl->frame_pitch = width * 4; // ARGB image
    pimpl->frame_rate = fps;
    pimpl->frame_bitrate = bitrate;
    return pimpl->start_encode();
}

void video_encoder::stop_encode() {
    pimpl->cleanup();
}

bool video_encoder::encode_frame(void *frame_ptr, void **output_ptr, size_t *output_size) {
    // register frame ptr
    if (pimpl->frame_ptr != frame_ptr) {
        pimpl->unregister_frame_ptr();
        pimpl->frame_ptr = frame_ptr;
        pimpl->register_frame_ptr();
    }

    pimpl->output_ptr = output_ptr;
    pimpl->output_size = output_size;
    return pimpl->encode_frame();
}

bool video_encoder::encode_frame(cudaGraphicsResource *res, void **output_ptr, size_t *output_size) {
    void *pbo_ptr;
    size_t pbo_size;
    CUDA_API_CHECK(cudaGraphicsMapResources(1, &res));
    CUDA_API_CHECK(cudaGraphicsResourceGetMappedPointer(&pbo_ptr, &pbo_size, res));
    assert(pbo_size == pimpl->frame_pitch * pimpl->frame_height);
    CALL_CHECK(encode_frame(pbo_ptr, output_ptr, output_size));
    CUDA_API_CHECK(cudaGraphicsUnmapResources(1, &res));
    return true;
}

bool video_encoder::is_encoding() {
    return pimpl->encoder != nullptr;
}

void video_encoder::refresh() {
    pimpl->idr_flag = true;
}