#include "encoder_video_nvenc.h"

namespace am {

	encoder_video_nvenc::encoder_video_nvenc()
	{
		av_register_all();

		_encoder = NULL;
		_encoder_ctx = NULL;
		_frame = NULL;
		_buff = NULL;
		_buff_size = 0;
		_y_size = 0;
	}


	encoder_video_nvenc::~encoder_video_nvenc()
	{
		stop();

		cleanup();
	}

	int encoder_video_nvenc::init(int pic_width, int pic_height, int frame_rate, int bit_rate, int qb, int key_pic_sec)
	{
		if (_inited == true)
			return AE_NO;

		int err = AE_NO;
		int ret = 0;

		do {
			_encoder = avcodec_find_encoder_by_name("h264_nvenc");
			if(!_encoder) _encoder = avcodec_find_encoder_by_name("nvenc_h264");

			if (!_encoder) {
				err = AE_FFMPEG_FIND_ENCODER_FAILED;
				break;
			}

			_encoder_ctx = avcodec_alloc_context3(_encoder);
			if (!_encoder_ctx) {
				err = AE_FFMPEG_ALLOC_CONTEXT_FAILED;
				break;
			}

			// ffmpeg -h encoder=h264_nvenc show all encoder options

			const char *rate_control = "cbr"; // cbr | cqp | vbr | lossless
			const char *profile = "baseline"; // baseline | main | high |high444p
			const char *preset = "default";      // default | slow | medium | fast | 
											  // hp | hq | bd | 11 | 11hq | 11hp | lossless | losslesshp

#if 0//USE_CBR
			av_opt_set_int(_encoder_ctx->priv_data, "cbr", true, 0);
			_encoder_ctx->bit_rate = bit_rate - bit_rate * (100 - qb) / 100;
			_encoder_ctx->rc_buffer_size = _encoder_ctx->bit_rate;
			_encoder_ctx->rc_max_rate = _encoder_ctx->bit_rate;
			_encoder_ctx->rc_min_rate = _encoder_ctx->bit_rate;
#else
			_encoder_ctx->bit_rate = bit_rate;
			
			//qb is 0 ~ 100
			qb = max(min(qb, 100), 0);

			//for qmax more larger,quality is more less, max qmax is qmin + 30*(100 - 0)/100 = qmin + 30
			_encoder_ctx->qmin = 30;
			_encoder_ctx->qmax = _encoder_ctx->qmin + 15 * (100 - qb) / 100;
			
#endif

			av_opt_set(_encoder_ctx->priv_data, "profile", profile, 0);
			av_opt_set(_encoder_ctx->priv_data, "preset", preset, 0);

			av_opt_set(_encoder_ctx->priv_data, "level", "auto", 0);
			av_opt_set_int(_encoder_ctx->priv_data, "2pass", false, 0);
			av_opt_set_int(_encoder_ctx->priv_data, "gpu", 0, 0);

			
			_encoder_ctx->width = pic_width;
			_encoder_ctx->height = pic_height;
			_encoder_ctx->time_base.num = 1;
			_encoder_ctx->time_base.den = frame_rate;
			_encoder_ctx->framerate = { frame_rate,1 };
			_encoder_ctx->pix_fmt = AV_PIX_FMT_YUV420P;		
			
			if (key_pic_sec == 0)
				_encoder_ctx->gop_size = 250;
			else
				_encoder_ctx->gop_size = key_pic_sec * _encoder_ctx->time_base.den / _encoder_ctx->time_base.num;

			_encoder_ctx->max_b_frames = 0;//NO B Frame
			_encoder_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

			ret = avcodec_open2(_encoder_ctx, _encoder, nullptr);
			if (ret != 0) {
				err = AE_FFMPEG_OPEN_CODEC_FAILED;
				break;
			}

			_frame = av_frame_alloc();

			_buff_size = av_image_get_buffer_size(_encoder_ctx->pix_fmt, _encoder_ctx->width, _encoder_ctx->height, 1);

			_buff = (uint8_t*)av_malloc(_buff_size);
			if (!_buff) {
				break;
			}

			av_image_fill_arrays(_frame->data, _frame->linesize, _buff, _encoder_ctx->pix_fmt, _encoder_ctx->width, _encoder_ctx->height, 1);

			_frame->format = _encoder_ctx->pix_fmt;
			_frame->width = _encoder_ctx->width;
			_frame->height = _encoder_ctx->height;

			_y_size = _encoder_ctx->width * _encoder_ctx->height;

			_time_base = _encoder_ctx->time_base;

			_inited = true;
		} while (0);

		if (err != AE_NO) {
			al_debug("%s,error:%d %lu", err2str(err), ret, GetLastError());
			cleanup();
		}

		return err;
	}

	int encoder_video_nvenc::get_extradata_size()
	{
		return _encoder_ctx->extradata_size;
	}

	const uint8_t * encoder_video_nvenc::get_extradata()
	{
		return (const uint8_t*)_encoder_ctx->extradata;
	}

	AVCodecID encoder_video_nvenc::get_codec_id()
	{
		if (_inited == false) return AV_CODEC_ID_NONE;

		return _encoder->id;
	}

	void encoder_video_nvenc::cleanup()
	{
		if (_frame)
			av_free(_frame);
		_frame = NULL;

		if (_buff)
			av_free(_buff);

		_buff = NULL;

		if (_encoder)
			avcodec_close(_encoder_ctx);

		_encoder = NULL;

		if (_encoder_ctx)
			avcodec_free_context(&_encoder_ctx);

		_encoder_ctx = NULL;
	}

	void encoder_video_nvenc::encode_loop()
	{
		AVPacket *packet = av_packet_alloc();
		AVFrame yuv_frame;

		int error = AE_NO;

		while (_running)
		{
			std::unique_lock<std::mutex> lock(_mutex);
			while (!_cond_notify && _running)
				_cond_var.wait_for(lock, std::chrono::milliseconds(300));

			while (_ring_buffer->get(_buff, _buff_size, yuv_frame)) {
				_frame->pkt_dts = yuv_frame.pkt_dts;
				_frame->pkt_dts = yuv_frame.pkt_dts;
				_frame->pts = yuv_frame.pts;

				if ((error = encode(_frame, packet)) != AE_NO) {
					if (_on_error)
						_on_error(error);

					al_fatal("encode 264 packet failed:%d", error);

					break;
				}
			}

			_cond_notify = false;
		}

		//flush frame in encoder
		encode(NULL, packet);

		av_packet_free(&packet);
	}

	int encoder_video_nvenc::encode(AVFrame * frame, AVPacket * packet)
	{
		int ret = avcodec_send_frame(_encoder_ctx, frame);
		if (ret < 0) {
			return AE_FFMPEG_ENCODE_FRAME_FAILED;
		}

		while (ret >= 0) {

			av_init_packet(packet);

			ret = avcodec_receive_packet(_encoder_ctx, packet);
			if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
				break;
			}

			if (ret < 0) {
				return AE_FFMPEG_READ_PACKET_FAILED;
			}

			if (ret == 0 && _on_data)
				_on_data(packet);

			av_packet_unref(packet);
		}

		return AE_NO;
	}

}