/**
 * FreeRDP: A Remote Desktop Protocol Implementation
 * H.264 Bitmap Compression
 *
 * Copyright 2015 Marc-André Moreau <marcandre.moreau@gmail.com>
 * Copyright 2014 Mike McDonald <Mike.McDonald@software.dell.com>
 * Copyright 2014 erbth <t.erbesdobler@team103.com>
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <winpr/wlog.h>
#include <freerdp/log.h>
#include <freerdp/codec/h264.h>
#include <libavcodec/avcodec.h>
#include <libavutil/opt.h>

#ifdef WITH_VAAPI
#if LIBAVUTIL_VERSION_INT >= AV_VERSION_INT(55, 9, 0)
#include <libavutil/hwcontext.h>
#else
#pragma warning You have asked for VA-API decoding, but your version of libavutil is too old! Disabling.
#undef WITH_VAAPI
#endif
#endif

/* Fallback support for older libavcodec versions */
#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(54, 59, 100)
#define AV_CODEC_ID_H264 CODEC_ID_H264
#endif

#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(56, 34, 2)
#define AV_CODEC_FLAG_LOOP_FILTER CODEC_FLAG_LOOP_FILTER
#define AV_CODEC_CAP_TRUNCATED CODEC_CAP_TRUNCATED
#define AV_CODEC_FLAG_TRUNCATED CODEC_FLAG_TRUNCATED
#endif

#if LIBAVUTIL_VERSION_MAJOR < 52
#define AV_PIX_FMT_YUV420P PIX_FMT_YUV420P
#endif

/* Ubuntu 14.04 ships without the functions provided by avutil,
 * so define error to string methods here. */
#if !defined(av_err2str)
static inline char* error_string(char* errbuf, size_t errbuf_size, int errnum)
{
	av_strerror(errnum, errbuf, errbuf_size);
	return errbuf;
}

#define av_err2str(errnum) \
	error_string((char[64]){0}, 64, errnum)
#endif

#ifdef WITH_VAAPI
#define VAAPI_DEVICE "/dev/dri/renderD128"
#endif

struct _H264_CONTEXT_LIBAVCODEC
{
	AVCodec* codecDecoder;
	AVCodecContext* codecDecoderContext;
	AVCodec* codecEncoder;
	AVCodecContext* codecEncoderContext;
	AVCodecParserContext* codecParser;
	AVFrame* videoFrame;
	AVPacket packet;
#ifdef WITH_VAAPI
	AVBufferRef* hwctx;
	AVFrame* hwVideoFrame;
	enum AVPixelFormat hw_pix_fmt;
#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(57, 80, 100)
	AVBufferRef* hw_frames_ctx;
#endif
#endif
};
typedef struct _H264_CONTEXT_LIBAVCODEC H264_CONTEXT_LIBAVCODEC;

static void libavcodec_destroy_encoder(H264_CONTEXT* h264)
{
	H264_CONTEXT_LIBAVCODEC* sys;

	if (!h264 || !h264->subsystem)
		return;

	sys = (H264_CONTEXT_LIBAVCODEC*)h264->pSystemData;

	if (sys->codecEncoderContext)
	{
		avcodec_close(sys->codecEncoderContext);
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(55, 69, 100)
		avcodec_free_context(&sys->codecEncoderContext);
#else
		av_free(sys->codecEncoderContext);
#endif
	}

	sys->codecEncoder = NULL;
	sys->codecEncoderContext = NULL;
}

static BOOL libavcodec_create_encoder(H264_CONTEXT* h264)
{
	BOOL recreate = FALSE;
	H264_CONTEXT_LIBAVCODEC* sys;

	if (!h264 || !h264->subsystem)
		return FALSE;

	sys = (H264_CONTEXT_LIBAVCODEC*)h264->pSystemData;
	recreate = !sys->codecEncoder || !sys->codecEncoderContext;

	if (sys->codecEncoderContext)
	{
		if ((sys->codecEncoderContext->width != h264->width) ||
		    (sys->codecEncoderContext->height != h264->height))
			recreate = TRUE;
	}

	if (!recreate)
		return TRUE;

	libavcodec_destroy_encoder(h264);
	sys->codecEncoder = avcodec_find_encoder(AV_CODEC_ID_H264);

	if (!sys->codecEncoder)
		goto EXCEPTION;

	sys->codecEncoderContext = avcodec_alloc_context3(sys->codecEncoder);

	if (!sys->codecEncoderContext)
		goto EXCEPTION;

	switch (h264->RateControlMode)
	{
		case H264_RATECONTROL_VBR:
			sys->codecEncoderContext->bit_rate = h264->BitRate;
			break;

		case H264_RATECONTROL_CQP:
			/* TODO: sys->codecEncoderContext-> = h264->QP; */
			break;

		default:
			break;
	}

	sys->codecEncoderContext->width = h264->width;
	sys->codecEncoderContext->height = h264->height;
	sys->codecEncoderContext->delay = 0;
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(56, 13, 100)
	sys->codecEncoderContext->framerate = (AVRational)
	{
		h264->FrameRate, 1
	};
#endif
	sys->codecEncoderContext->time_base = (AVRational)
	{
		1, h264->FrameRate
	};
	av_opt_set(sys->codecEncoderContext, "preset", "medium", AV_OPT_SEARCH_CHILDREN);
	av_opt_set(sys->codecEncoderContext, "tune", "zerolatency", AV_OPT_SEARCH_CHILDREN);
	sys->codecEncoderContext->flags |= AV_CODEC_FLAG_LOOP_FILTER;
	sys->codecEncoderContext->pix_fmt = AV_PIX_FMT_YUV420P;

	if (avcodec_open2(sys->codecEncoderContext, sys->codecEncoder, NULL) < 0)
		goto EXCEPTION;

	return TRUE;
EXCEPTION:
	libavcodec_destroy_encoder(h264);
	return FALSE;
}

static int libavcodec_decompress(H264_CONTEXT* h264, const BYTE* pSrcData,
                                 UINT32 SrcSize)
{
	int status;
	int gotFrame = 0;
	H264_CONTEXT_LIBAVCODEC* sys = (H264_CONTEXT_LIBAVCODEC*) h264->pSystemData;
	BYTE** pYUVData = h264->pYUVData;
	UINT32* iStride = h264->iStride;
	av_init_packet(&sys->packet);
	sys->packet.data = (BYTE*)pSrcData;
	sys->packet.size = SrcSize;
	/* avcodec_decode_video2 is deprecated with libavcodec 57.48.101 */
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 48, 101)
	status = avcodec_send_packet(sys->codecDecoderContext, &sys->packet);

	if (status < 0)
	{
		WLog_Print(h264->log, WLOG_ERROR, "Failed to decode video frame (status=%d)", status);
		return -1;
	}

	sys->videoFrame->format = AV_PIX_FMT_YUV420P;

	do
	{
#ifdef WITH_VAAPI
		status = avcodec_receive_frame(sys->codecDecoderContext,
		                               sys->hwctx ? sys->hwVideoFrame : sys->videoFrame);
#else
		status = avcodec_receive_frame(sys->codecDecoderContext, sys->videoFrame);
#endif
	}
	while (status == AVERROR(EAGAIN));

	gotFrame = (status == 0);
#else
#ifdef WITH_VAAPI
	status = avcodec_decode_video2(sys->codecDecoderContext,
	                               sys->hwctx ? sys->hwVideoFrame : sys->videoFrame, &gotFrame,
	                               &sys->packet);
#else
	status = avcodec_decode_video2(sys->codecDecoderContext, sys->videoFrame, &gotFrame,
	                               &sys->packet);
#endif
#endif

	if (status < 0)
	{
		WLog_Print(h264->log, WLOG_ERROR, "Failed to decode video frame (status=%d)", status);
		return -1;
	}

#ifdef WITH_VAAPI

	if (sys->hwctx)
	{
		if (sys->hwVideoFrame->format == sys->hw_pix_fmt)
		{
			sys->videoFrame->width = sys->hwVideoFrame->width;
			sys->videoFrame->height = sys->hwVideoFrame->height;
			status = av_hwframe_transfer_data(sys->videoFrame, sys->hwVideoFrame, 0);
		}
		else
		{
			status = av_frame_copy(sys->videoFrame, sys->hwVideoFrame);
		}
	}

	gotFrame = (status == 0);

	if (status < 0)
	{
		WLog_Print(h264->log, WLOG_ERROR, "Failed to transfer video frame (status=%d) (%s)", status,
		           av_err2str(status));
		return -1;
	}

#endif
#if 0
	WLog_Print(h264->log, WLOG_INFO,
	           "libavcodec_decompress: frame decoded (status=%d, gotFrame=%d, width=%d, height=%d, Y=[%p,%d], U=[%p,%d], V=[%p,%d])",
	           status, gotFrame, sys->videoFrame->width, sys->videoFrame->height,
	           (void*) sys->videoFrame->data[0], sys->videoFrame->linesize[0],
	           (void*) sys->videoFrame->data[1], sys->videoFrame->linesize[1],
	           (void*) sys->videoFrame->data[2], sys->videoFrame->linesize[2]);
#endif

	if (gotFrame)
	{
		pYUVData[0] = sys->videoFrame->data[0];
		pYUVData[1] = sys->videoFrame->data[1];
		pYUVData[2] = sys->videoFrame->data[2];
		iStride[0] = sys->videoFrame->linesize[0];
		iStride[1] = sys->videoFrame->linesize[1];
		iStride[2] = sys->videoFrame->linesize[2];
		h264->width = sys->videoFrame->width;
		h264->height = sys->videoFrame->height;
	}
	else
		return -2;

	return 1;
}

static int libavcodec_compress(H264_CONTEXT* h264, const BYTE** pSrcYuv, const UINT32* pStride,
                               BYTE** ppDstData, UINT32* pDstSize)
{
	int status;
	int gotFrame = 0;
	H264_CONTEXT_LIBAVCODEC* sys = (H264_CONTEXT_LIBAVCODEC*) h264->pSystemData;

	if (!libavcodec_create_encoder(h264))
		return -1;

#if LIBAVUTIL_VERSION_INT >= AV_VERSION_INT(55, 39, 100)
	av_packet_unref(&sys->packet);
#else
	av_free(sys->packet.data);
#endif
	av_init_packet(&sys->packet);
	sys->packet.data = NULL;
	sys->packet.size = 0;
	sys->videoFrame->format = sys->codecEncoderContext->pix_fmt;
	sys->videoFrame->width  = sys->codecEncoderContext->width;
	sys->videoFrame->height = sys->codecEncoderContext->height;
#if LIBAVUTIL_VERSION_INT >= AV_VERSION_INT(52, 48, 100)
	sys->videoFrame->colorspace = AVCOL_SPC_BT709;
#endif
#if LIBAVUTIL_VERSION_INT >= AV_VERSION_INT(52, 92, 100)
	sys->videoFrame->chroma_location = AVCHROMA_LOC_LEFT;
#endif
	sys->videoFrame->data[0] = (uint8_t*)pSrcYuv[0];
	sys->videoFrame->data[1] = (uint8_t*)pSrcYuv[1];
	sys->videoFrame->data[2] = (uint8_t*)pSrcYuv[2];
	sys->videoFrame->linesize[0] = (int)pStride[0];
	sys->videoFrame->linesize[1] = (int)pStride[1];
	sys->videoFrame->linesize[2] = (int)pStride[2];
	sys->videoFrame->pts++;
	/* avcodec_encode_video2 is deprecated with libavcodec 57.48.101 */
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 48, 101)
	status = avcodec_send_frame(sys->codecEncoderContext, sys->videoFrame);

	if (status < 0)
	{
		WLog_Print(h264->log, WLOG_ERROR, "Failed to encode video frame (%s [%d])",
		           av_err2str(status), status);
		return -1;
	}

	status = avcodec_receive_packet(sys->codecEncoderContext,
	                                &sys->packet);

	if (status < 0)
	{
		WLog_Print(h264->log, WLOG_ERROR, "Failed to encode video frame (%s [%d])",
		           av_err2str(status), status);
		return -1;
	}

	gotFrame = (status == 0);
#elif LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(54, 59, 100)

	do
	{
		status = avcodec_encode_video2(sys->codecEncoderContext,
		                               &sys->packet,
		                               sys->videoFrame, &gotFrame);
	}
	while ((status >= 0) && (gotFrame == 0));

#else
	sys->packet.size = avpicture_get_size(sys->codecDecoderContext->pix_fmt,
	                                      sys->codecDecoderContext->width,
	                                      sys->codecDecoderContext->height);
	sys->packet.data = av_malloc(sys->packet.size);

	if (!sys->packet.data)
		status = -1;
	else
	{
		status = avcodec_encode_video(sys->codecDecoderContext, sys->packet.data,
		                              sys->packet.size, sys->videoFrame);
	}

#endif

	if (status < 0)
	{
		WLog_Print(h264->log, WLOG_ERROR, "Failed to encode video frame (%s [%d])",
		           av_err2str(status), status);
		return -1;
	}

	*ppDstData = sys->packet.data;
	*pDstSize = sys->packet.size;

	if (!gotFrame)
	{
		WLog_Print(h264->log, WLOG_ERROR, "Did not get frame! (%s [%d])", av_err2str(status), status);
		return -2;
	}

	return 1;
}

static void libavcodec_uninit(H264_CONTEXT* h264)
{
	H264_CONTEXT_LIBAVCODEC* sys = (H264_CONTEXT_LIBAVCODEC*) h264->pSystemData;

	if (!sys)
		return;

	if (sys->videoFrame)
	{
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(55, 18, 102)
		av_frame_free(&sys->videoFrame);
#else
		av_free(sys->videoFrame);
#endif
	}

#ifdef WITH_VAAPI

	if (sys->hwVideoFrame)
	{
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(55, 18, 102)
		av_frame_free(&sys->hwVideoFrame);
#else
		av_free(sys->hwVideoFrame);
#endif
	}

	if (sys->hwctx)
	{
		av_buffer_unref(&sys->hwctx);
	}

#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(57, 80, 100)

	if (sys->hw_frames_ctx)
	{
		av_buffer_unref(&sys->hw_frames_ctx);
	}

#endif
#endif

	if (sys->codecParser)
		av_parser_close(sys->codecParser);

	if (sys->codecDecoderContext)
	{
		avcodec_close(sys->codecDecoderContext);
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(55, 69, 100)
		avcodec_free_context(&sys->codecDecoderContext);
#else
		av_free(sys->codecDecoderContext);
#endif
	}

	libavcodec_destroy_encoder(h264);
	free(sys);
	h264->pSystemData = NULL;
}

#ifdef WITH_VAAPI
static enum AVPixelFormat libavcodec_get_format(struct AVCodecContext* ctx,
        const enum AVPixelFormat* fmts)
{
	H264_CONTEXT* h264 = (H264_CONTEXT*) ctx->opaque;
	H264_CONTEXT_LIBAVCODEC* sys = (H264_CONTEXT_LIBAVCODEC*) h264->subsystem;
	const enum AVPixelFormat* p;

	for (p = fmts; *p != AV_PIX_FMT_NONE; p++)
	{
		if (*p == sys->hw_pix_fmt)
		{
#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(57, 80, 100)
			sys->hw_frames_ctx = av_hwframe_ctx_alloc(sys->hwctx);

			if (!sys->hw_frames_ctx)
			{
				return AV_PIX_FMT_NONE;
			}

			sys->codecDecoderContext->pix_fmt = *p;
			AVHWFramesContext* frames = (AVHWFramesContext*) sys->hw_frames_ctx->data;
			frames->format = *p;
			frames->height = sys->codecDecoderContext->coded_height;
			frames->width = sys->codecDecoderContext->coded_width;
			frames->sw_format = (sys->codecDecoderContext->sw_pix_fmt == AV_PIX_FMT_YUV420P10 ?
			                     AV_PIX_FMT_P010 : AV_PIX_FMT_NV12);
			frames->initial_pool_size = 20;

			if (sys->codecDecoderContext->active_thread_type & FF_THREAD_FRAME)
				frames->initial_pool_size += sys->codecDecoderContext->thread_count;

			int err = av_hwframe_ctx_init(sys->hw_frames_ctx);

			if (err < 0)
			{
				WLog_Print(h264->log, WLOG_ERROR, "Could not init hwframes context: %s", av_err2str(err));
				return AV_PIX_FMT_NONE;
			}

			sys->codecDecoderContext->hw_frames_ctx = av_buffer_ref(sys->hw_frames_ctx);
#endif
			return *p;
		}
	}

	return AV_PIX_FMT_NONE;
}
#endif

static BOOL libavcodec_init(H264_CONTEXT* h264)
{
	H264_CONTEXT_LIBAVCODEC* sys;
	sys = (H264_CONTEXT_LIBAVCODEC*) calloc(1, sizeof(H264_CONTEXT_LIBAVCODEC));

	if (!sys)
	{
		goto EXCEPTION;
	}

	h264->pSystemData = (void*) sys;
	avcodec_register_all();

	if (!h264->Compressor)
	{
		sys->codecDecoder = avcodec_find_decoder(AV_CODEC_ID_H264);

		if (!sys->codecDecoder)
		{
			WLog_Print(h264->log, WLOG_ERROR, "Failed to find libav H.264 codec");
			goto EXCEPTION;
		}

		sys->codecDecoderContext = avcodec_alloc_context3(sys->codecDecoder);

		if (!sys->codecDecoderContext)
		{
			WLog_Print(h264->log, WLOG_ERROR, "Failed to allocate libav codec context");
			goto EXCEPTION;
		}

		if (sys->codecDecoder->capabilities & AV_CODEC_CAP_TRUNCATED)
		{
			sys->codecDecoderContext->flags |= AV_CODEC_FLAG_TRUNCATED;
		}

#ifdef WITH_VAAPI

		if (!sys->hwctx)
		{
			int ret = av_hwdevice_ctx_create(&sys->hwctx, AV_HWDEVICE_TYPE_VAAPI, VAAPI_DEVICE, NULL, 0);

			if (ret < 0)
			{
				WLog_Print(h264->log, WLOG_ERROR,
				           "Could not initialize hardware decoder, falling back to software: %s",
				           av_err2str(ret));
				sys->hwctx = NULL;
				goto fail_hwdevice_create;
			}
		}

		sys->codecDecoderContext->get_format = libavcodec_get_format;
		sys->hw_pix_fmt = AV_PIX_FMT_VAAPI;
#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(57, 80, 100)
		sys->codecDecoderContext->hw_device_ctx = av_buffer_ref(sys->hwctx);
#endif
		sys->codecDecoderContext->opaque = (void*) h264;
	fail_hwdevice_create:
#endif

		if (avcodec_open2(sys->codecDecoderContext, sys->codecDecoder, NULL) < 0)
		{
			WLog_Print(h264->log, WLOG_ERROR, "Failed to open libav codec");
			goto EXCEPTION;
		}

		sys->codecParser = av_parser_init(AV_CODEC_ID_H264);

		if (!sys->codecParser)
		{
			WLog_Print(h264->log, WLOG_ERROR, "Failed to initialize libav parser");
			goto EXCEPTION;
		}
	}

#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(55, 18, 102)
	sys->videoFrame = av_frame_alloc();
#ifdef WITH_VAAPI
	sys->hwVideoFrame = av_frame_alloc();
#endif
#else
	sys->videoFrame = avcodec_alloc_frame();
#endif

	if (!sys->videoFrame)
	{
		WLog_Print(h264->log, WLOG_ERROR, "Failed to allocate libav frame");
		goto EXCEPTION;
	}

#ifdef WITH_VAAPI

	if (!sys->hwVideoFrame)
	{
		WLog_Print(h264->log, WLOG_ERROR, "Failed to allocate libav hw frame");
		goto EXCEPTION;
	}

#endif
	sys->videoFrame->pts = 0;
	return TRUE;
EXCEPTION:
	libavcodec_uninit(h264);
	return FALSE;
}

H264_CONTEXT_SUBSYSTEM g_Subsystem_libavcodec =
{
	"libavcodec",
	libavcodec_init,
	libavcodec_uninit,
	libavcodec_decompress,
	libavcodec_compress
};