#include "NvJpegEncoder.h"

#include <nvjpeg.h>

#include <fstream>
#include <vector>
#include <iostream>



int saveJpeg(const char * filepath, unsigned char* d_srcBGR, int width, int height)
{
	nvjpegHandle_t nvjpeg_handle;
	nvjpegEncoderState_t encoder_state;
	nvjpegEncoderParams_t encoder_params;

	cudaEvent_t ev_start, ev_end;
	cudaEventCreate(&ev_start);
	cudaEventCreate(&ev_end);
	

	nvjpegImage_t input;
	nvjpegInputFormat_t input_format = NVJPEG_INPUT_BGRI;
	int image_width = width;
	int image_height = height;

	// int channel_size = image_width * image_height;
	// for (int i = 0; i < 3; i++)
	// {
	// 	input.pitch[i] = image_width;
	// 	(cudaMalloc((void**)&(input.channel[i]), channel_size));
	// 	(cudaMemset(input.channel[i], 50 * 40 * i, channel_size));
	// }

    input.channel[0] = d_srcBGR;
    input.pitch[0] = image_width * 3;

	nvjpegBackend_t backend = NVJPEG_BACKEND_DEFAULT;

	nvjpegCreate(backend, nullptr, &nvjpeg_handle);
	
	nvjpegEncoderParamsCreate(nvjpeg_handle, &encoder_params, NULL);
	nvjpegEncoderStateCreate(nvjpeg_handle, &encoder_state, NULL);

	// set params
	nvjpegEncoderParamsSetEncoding(encoder_params, nvjpegJpegEncoding_t::NVJPEG_ENCODING_PROGRESSIVE_DCT_HUFFMAN, NULL);
	nvjpegEncoderParamsSetOptimizedHuffman(encoder_params, 1, NULL);
	nvjpegEncoderParamsSetQuality(encoder_params, 70, NULL);
	nvjpegEncoderParamsSetSamplingFactors(encoder_params, nvjpegChromaSubsampling_t::NVJPEG_CSS_420, NULL);

	cudaEventRecord(ev_start);
	nvjpegEncodeImage(nvjpeg_handle, encoder_state, encoder_params, &input, input_format, image_width, image_height, NULL);
	cudaEventRecord(ev_end);

	std::vector<unsigned char> obuffer;
	size_t length;
	nvjpegEncodeRetrieveBitstream(
		nvjpeg_handle,
		encoder_state,
		NULL,
		&length,
		NULL);

	obuffer.resize(length);
	nvjpegEncodeRetrieveBitstream(
		nvjpeg_handle,
		encoder_state,
		obuffer.data(),
		&length,
		NULL);

	std::ofstream outputFile(filepath, std::ios::out | std::ios::binary);
	outputFile.write(reinterpret_cast<const char *>(obuffer.data()), static_cast<int>(length));

	cudaEventSynchronize(ev_end);

	float ms;
	cudaEventElapsedTime(&ms, ev_start, ev_end);
	std::cout << "time spend " << ms << " ms" << std::endl;
	
	return 0;
}