#include "cuda_helper.hpp"
#include "simple_mq.h"
#include "third_party/scope_guard.hpp"
#include "variable_defs.h"
#include "encoder_nvenc.h"

#include <nvEncodeAPI.h>

#ifdef _MSC_VER

#include <format>

#define fmt std

#else

#include <fmt/chrono.h>

#endif

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, (int) api_ret);
    RET_ERROR_B;
}

#define API_CHECK(api_call) \
    check_nvenc_api_call( \
        api_call, __LINE__, __FILE__, #api_call)

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

namespace video_encoder_impl {
    constexpr auto frame_buffer_type = NV_ENC_BUFFER_FORMAT_ARGB;
    static auto codec_guid = NV_ENC_CODEC_HEVC_GUID;
    static auto preset_guid = NV_ENC_PRESET_P3_GUID;
    constexpr auto tuning_info = NV_ENC_TUNING_INFO_ULTRA_LOW_LATENCY;
    std::unique_ptr<NV_ENCODE_API_FUNCTION_LIST> api;
}

using namespace video_encoder_impl;
using namespace simple_mq_singleton;

struct encoder_nvenc::impl {

    std::unique_ptr<NV_ENC_PRESET_CONFIG> preset_config;
    std::unique_ptr<NV_ENC_INITIALIZE_PARAMS> init_params;
    void *encoder;
    NV_ENC_OUTPUT_PTR output_buf;

    cv::Size frame_size;
    FILE *save_file = nullptr;
    bool save_length;

    void *last_frame_ptr = nullptr;
    NV_ENC_REGISTERED_PTR last_reg_ptr = nullptr;

    ~impl() {
        // 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_CHECK(api->nvEncEncodePicture(encoder, &pic_params));

        // releasing resources
        unregister_frame_ptr();
        API_CHECK(api->nvEncDestroyBitstreamBuffer(encoder, output_buf));

        // close encoder
        API_CHECK(api->nvEncDestroyEncoder(encoder));

        // close save file
        if (save_file != nullptr) {
            fclose(save_file);
        }

        SPDLOG_INFO("Video encoder stopped.");
    }

    static impl *create(const nvenc_config &conf) {
        // initialize api
        if (api == nullptr) [[unlikely]] {
            api = std::make_unique<NV_ENCODE_API_FUNCTION_LIST>(
                    NV_ENCODE_API_FUNCTION_LIST_VER);
            API_CHECK_P(NvEncodeAPICreateInstance(api.get()));
        }

        // get cuda context
//        CUcontext cuda_ctx;
//        CUDA_API_CHECK(cuCtxGetCurrent(&cuda_ctx));
        auto cuda_ctx = mq().query_variable<CUcontext>(CUDA_CONTEXT);

        // create encoder
        auto ret = new impl;
        ret->frame_size = conf.frame_size;
        auto closer = sg::make_scope_guard([&] {
            if (ret->save_file != nullptr) {
                fclose(ret->save_file);
            }
            delete ret;
        });
        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;
        API_CHECK_P(api->nvEncOpenEncodeSessionEx(&session_params, &ret->encoder));

        // get preset config
        ret->preset_config = std::make_unique<NV_ENC_PRESET_CONFIG>();
        auto &preset_config = *ret->preset_config;
        preset_config.version = NV_ENC_PRESET_CONFIG_VER;
        preset_config.presetCfg.version = NV_ENC_CONFIG_VER;
        API_CHECK_P(api->nvEncGetEncodePresetConfigEx(
                ret->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 = conf.bitrate_mbps * 1e6;
        rc_params.enableAQ = true;
        rc_params.multiPass = NV_ENC_TWO_PASS_QUARTER_RESOLUTION;
        // TODO; fine tune encoder config

        // start_encode encoder
        ret->init_params =
                std::make_unique<NV_ENC_INITIALIZE_PARAMS>(NV_ENC_INITIALIZE_PARAMS_VER);
        auto &init_params = *ret->init_params;
        init_params.encodeGUID = codec_guid;
        init_params.presetGUID = preset_guid;
        init_params.encodeWidth = conf.frame_size.width;
        init_params.encodeHeight = conf.frame_size.height;
        init_params.darWidth = conf.frame_size.width; // TODO; learn more about this
        init_params.darHeight = conf.frame_size.height; // TODO; learn more about this
        init_params.frameRateNum = conf.frame_rate;
        init_params.frameRateDen = 1;
        init_params.enablePTD = 1;
        init_params.encodeConfig = &preset_config.presetCfg;
        init_params.maxEncodeWidth = conf.frame_size.width;
        init_params.maxEncodeHeight = conf.frame_size.height;
        init_params.tuningInfo = tuning_info;
        init_params.bufferFormat = frame_buffer_type;
        API_CHECK_P(api->nvEncInitializeEncoder(ret->encoder, &init_params));

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

        // create save file
        if (conf.save_file) {
            auto file_name = fmt::format("record_{:%Y_%m_%d_%H_%M_%S}.{}",
                                         std::chrono::system_clock::now(),
                                         conf.save_length ? "dat" : "hevc");
            ret->save_file = fopen(file_name.c_str(), "wb");
            ret->save_length = conf.save_length;
        }

        SPDLOG_INFO("Video encoder started.");
        closer.dismiss();
        return ret;
    }

    void unregister_frame_ptr() {
        if (last_reg_ptr == nullptr) return;
        API_CHECK(api->nvEncUnregisterResource(encoder, last_reg_ptr));
        last_reg_ptr = nullptr;
    }

    void register_frame_ptr(const cv::cuda::GpuMat &img) {
        NV_ENC_REGISTER_RESOURCE reg_params = {NV_ENC_REGISTER_RESOURCE_VER};
        reg_params.resourceType = NV_ENC_INPUT_RESOURCE_TYPE_CUDADEVICEPTR;
        reg_params.width = img.size().width;
        reg_params.height = img.size().height;
        reg_params.pitch = img.step;
        reg_params.resourceToRegister = img.cudaPtr();
        reg_params.bufferFormat = frame_buffer_type;
        reg_params.bufferUsage = NV_ENC_INPUT_IMAGE;
        API_CHECK(api->nvEncRegisterResource(encoder, &reg_params));
        last_reg_ptr = reg_params.registeredResource;
    }

    void encode(const cv::cuda::GpuMat &img, video_nal *out, bool force_idr) {
        // register pointer if needed
        if (img.cudaPtr() != last_frame_ptr) [[unlikely]] {
            assert(img.size() == frame_size);
            unregister_frame_ptr();
            register_frame_ptr(img);
        }

        // map input resource
        NV_ENC_MAP_INPUT_RESOURCE map_params = {
                NV_ENC_MAP_INPUT_RESOURCE_VER};
        map_params.registeredResource = last_reg_ptr;
        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 = img.size().width;
        pic_params.inputHeight = img.size().height;
        pic_params.inputPitch = img.step;
        if (force_idr) { // request for IDR frame
            pic_params.encodePicFlags = NV_ENC_PIC_FLAG_FORCEIDR | NV_ENC_PIC_FLAG_OUTPUT_SPSPPS;
        } 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
        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;
        API_CHECK(api->nvEncLockBitstream(encoder, &lock_config));

        // copy bitstream
        out->create(lock_config.bitstreamBufferPtr,
                    lock_config.bitstreamSizeInBytes,
                    lock_config.pictureType == NV_ENC_PIC_TYPE_IDR);

        // save bitstream
        if (save_file != nullptr) {
            if (save_length) {
                fwrite(&out->length, sizeof(size_t), 1, save_file);
            }
            fwrite(out->ptr, out->length, 1, save_file);
        }

        // cleanup
        API_CHECK(api->nvEncUnlockBitstream(encoder, output_buf));
        API_CHECK(api->nvEncUnmapInputResource(encoder, map_params.mappedResource));
    }

    void change_config(const nvenc_config &conf) {
        assert(conf.frame_size == frame_size);
        assert(conf.save_file == (save_file != nullptr));
        assert(conf.save_length == save_length);
        NV_ENC_RECONFIGURE_PARAMS params = {NV_ENC_RECONFIGURE_PARAMS_VER};
        init_params->frameRateNum = conf.frame_rate;
        init_params->encodeConfig->rcParams.averageBitRate = conf.bitrate_mbps * 1e6;
        params.reInitEncodeParams = *init_params;
        params.resetEncoder = true;
        params.forceIDR = true;
        API_CHECK(api->nvEncReconfigureEncoder(encoder, &params));
    }
};

encoder_nvenc::~encoder_nvenc() = default;

encoder_nvenc *encoder_nvenc::create(const nvenc_config &conf) {
    auto pimpl = impl::create(conf);
    if (pimpl == nullptr) return nullptr;
    auto ret = new encoder_nvenc;
    ret->pimpl.reset(pimpl);
    return ret;
}

void encoder_nvenc::encode(const cv::cuda::GpuMat &img, video_nal *out, bool force_idr) {
    mq().update_variable(ENCODER_BUSY, true);
    pimpl->encode(img, out, force_idr);
    mq().update_variable(ENCODER_BUSY, false);
}

void encoder_nvenc::change_config(const nvenc_config &conf) {
    pimpl->change_config(conf);
}