/**

 * FreeRDP: A Remote Desktop Protocol Implementation

 * Video Redirection Virtual Channel - FFmpeg Decoder

 *

 * Copyright 2010-2011 Vic Lee

 *

 * 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 <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <winpr/crt.h>

#include <freerdp/channels/log.h>

#include <freerdp/client/tsmf.h>

#include <libavcodec/avcodec.h>

#include <libavutil/common.h>

#include "tsmf_constants.h"

#include "tsmf_decoder.h"

/* Compatibility with older FFmpeg */

#if LIBAVUTIL_VERSION_MAJOR < 50

#define AVMEDIA_TYPE_VIDEO 0

#define AVMEDIA_TYPE_AUDIO 1

#endif

#if LIBAVCODEC_VERSION_MAJOR < 54

#define MAX_AUDIO_FRAME_SIZE AVCODEC_MAX_AUDIO_FRAME_SIZE

#else

#define MAX_AUDIO_FRAME_SIZE 192000

#endif

#if LIBAVCODEC_VERSION_MAJOR < 55

#define AV_CODEC_ID_VC1 CODEC_ID_VC1

#define AV_CODEC_ID_WMAV2 CODEC_ID_WMAV2

#define AV_CODEC_ID_WMAPRO CODEC_ID_WMAPRO

#define AV_CODEC_ID_MP3 CODEC_ID_MP3

#define AV_CODEC_ID_MP2 CODEC_ID_MP2

#define AV_CODEC_ID_MPEG2VIDEO CODEC_ID_MPEG2VIDEO

#define AV_CODEC_ID_WMV3 CODEC_ID_WMV3

#define AV_CODEC_ID_AAC CODEC_ID_AAC

#define AV_CODEC_ID_H264 CODEC_ID_H264

#define AV_CODEC_ID_AC3 CODEC_ID_AC3

#endif

#if LIBAVCODEC_VERSION_INT < AV_VERSION_INT(56, 34, 2)

#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

typedef struct _TSMFFFmpegDecoder

{

	ITSMFDecoder iface;

	int media_type;

#if LIBAVCODEC_VERSION_MAJOR < 55

	enum CodecID codec_id;

#else

	enum AVCodecID codec_id;

#endif

	AVCodecContext* codec_context;

	AVCodec* codec;

	AVFrame* frame;

	int prepared;

	BYTE* decoded_data;

	UINT32 decoded_size;

	UINT32 decoded_size_max;

} TSMFFFmpegDecoder;

static BOOL tsmf_ffmpeg_init_context(ITSMFDecoder* decoder)

{

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	mdecoder->codec_context = avcodec_alloc_context3(NULL);

	if (!mdecoder->codec_context)

	{

		WLog_ERR(TAG, "avcodec_alloc_context failed.");

		return FALSE;

	}

	return TRUE;

}

static BOOL tsmf_ffmpeg_init_video_stream(ITSMFDecoder* decoder, const TS_AM_MEDIA_TYPE* media_type)

{

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	mdecoder->codec_context->width = media_type->Width;

	mdecoder->codec_context->height = media_type->Height;

	mdecoder->codec_context->bit_rate = media_type->BitRate;

	mdecoder->codec_context->time_base.den = media_type->SamplesPerSecond.Numerator;

	mdecoder->codec_context->time_base.num = media_type->SamplesPerSecond.Denominator;

#if LIBAVCODEC_VERSION_MAJOR < 55

	mdecoder->frame = avcodec_alloc_frame();

#else

	mdecoder->frame = av_frame_alloc();

#endif

	return TRUE;

}

static BOOL tsmf_ffmpeg_init_audio_stream(ITSMFDecoder* decoder, const TS_AM_MEDIA_TYPE* media_type)

{

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	mdecoder->codec_context->sample_rate = media_type->SamplesPerSecond.Numerator;

	mdecoder->codec_context->bit_rate = media_type->BitRate;

	mdecoder->codec_context->channels = media_type->Channels;

	mdecoder->codec_context->block_align = media_type->BlockAlign;

#if LIBAVCODEC_VERSION_MAJOR < 55

#ifdef AV_CPU_FLAG_SSE2

	mdecoder->codec_context->dsp_mask = AV_CPU_FLAG_SSE2 | AV_CPU_FLAG_MMX2;

#else

#if LIBAVCODEC_VERSION_MAJOR < 53

	mdecoder->codec_context->dsp_mask = FF_MM_SSE2 | FF_MM_MMXEXT;

#else

	mdecoder->codec_context->dsp_mask = FF_MM_SSE2 | FF_MM_MMX2;

#endif

#endif

#else /* LIBAVCODEC_VERSION_MAJOR < 55 */

#ifdef AV_CPU_FLAG_SSE2

	av_set_cpu_flags_mask(AV_CPU_FLAG_SSE2 | AV_CPU_FLAG_MMXEXT);

#else

	av_set_cpu_flags_mask(FF_MM_SSE2 | FF_MM_MMX2);

#endif

#endif /* LIBAVCODEC_VERSION_MAJOR < 55 */

	return TRUE;

}

static BOOL tsmf_ffmpeg_init_stream(ITSMFDecoder* decoder, const TS_AM_MEDIA_TYPE* media_type)

{

	BYTE* p;

	UINT32 size;

	const BYTE* s;

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	mdecoder->codec = avcodec_find_decoder(mdecoder->codec_id);

	if (!mdecoder->codec)

	{

		WLog_ERR(TAG, "avcodec_find_decoder failed.");

		return FALSE;

	}

	mdecoder->codec_context->codec_id = mdecoder->codec_id;

	mdecoder->codec_context->codec_type = mdecoder->media_type;

	switch (mdecoder->media_type)

	{

		case AVMEDIA_TYPE_VIDEO:

			if (!tsmf_ffmpeg_init_video_stream(decoder, media_type))

				return FALSE;

			break;

		case AVMEDIA_TYPE_AUDIO:

			if (!tsmf_ffmpeg_init_audio_stream(decoder, media_type))

				return FALSE;

			break;

		default:

			WLog_ERR(TAG, "unknown media_type %d", mdecoder->media_type);

			break;

	}

	if (media_type->ExtraData)

	{

		/* Add a padding to avoid invalid memory read in some codec */

		mdecoder->codec_context->extradata_size = media_type->ExtraDataSize + 8;

		mdecoder->codec_context->extradata = calloc(1, mdecoder->codec_context->extradata_size);

		if (!mdecoder->codec_context->extradata)

			return FALSE;

		if (media_type->SubType == TSMF_SUB_TYPE_AVC1 &&

		    media_type->FormatType == TSMF_FORMAT_TYPE_MPEG2VIDEOINFO)

		{

			/* The extradata format that FFmpeg uses is following CodecPrivate in Matroska.

			   See http://haali.su/mkv/codecs.pdf */

			p = mdecoder->codec_context->extradata;

			*p++ = 1;                         /* Reserved? */

			*p++ = media_type->ExtraData[8];  /* Profile */

			*p++ = 0;                         /* Profile */

			*p++ = media_type->ExtraData[12]; /* Level */

			*p++ = 0xff;                      /* Flag? */

			*p++ = 0xe0 | 0x01;               /* Reserved | #sps */

			s = media_type->ExtraData + 20;

			size = ((UINT32)(*s)) * 256 + ((UINT32)(*(s + 1)));

			memcpy(p, s, size + 2);

			s += size + 2;

			p += size + 2;

			*p++ = 1; /* #pps */

			size = ((UINT32)(*s)) * 256 + ((UINT32)(*(s + 1)));

			memcpy(p, s, size + 2);

		}

		else

		{

			memcpy(mdecoder->codec_context->extradata, media_type->ExtraData,

			       media_type->ExtraDataSize);

			memset(mdecoder->codec_context->extradata + media_type->ExtraDataSize, 0, 8);

		}

	}

	if (mdecoder->codec->capabilities & AV_CODEC_CAP_TRUNCATED)

		mdecoder->codec_context->flags |= AV_CODEC_FLAG_TRUNCATED;

	return TRUE;

}

static BOOL tsmf_ffmpeg_prepare(ITSMFDecoder* decoder)

{

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	if (avcodec_open2(mdecoder->codec_context, mdecoder->codec, NULL) < 0)

	{

		WLog_ERR(TAG, "avcodec_open2 failed.");

		return FALSE;

	}

	mdecoder->prepared = 1;

	return TRUE;

}

static BOOL tsmf_ffmpeg_set_format(ITSMFDecoder* decoder, TS_AM_MEDIA_TYPE* media_type)

{

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	switch (media_type->MajorType)

	{

		case TSMF_MAJOR_TYPE_VIDEO:

			mdecoder->media_type = AVMEDIA_TYPE_VIDEO;

			break;

		case TSMF_MAJOR_TYPE_AUDIO:

			mdecoder->media_type = AVMEDIA_TYPE_AUDIO;

			break;

		default:

			return FALSE;

	}

	switch (media_type->SubType)

	{

		case TSMF_SUB_TYPE_WVC1:

			mdecoder->codec_id = AV_CODEC_ID_VC1;

			break;

		case TSMF_SUB_TYPE_WMA2:

			mdecoder->codec_id = AV_CODEC_ID_WMAV2;

			break;

		case TSMF_SUB_TYPE_WMA9:

			mdecoder->codec_id = AV_CODEC_ID_WMAPRO;

			break;

		case TSMF_SUB_TYPE_MP3:

			mdecoder->codec_id = AV_CODEC_ID_MP3;

			break;

		case TSMF_SUB_TYPE_MP2A:

			mdecoder->codec_id = AV_CODEC_ID_MP2;

			break;

		case TSMF_SUB_TYPE_MP2V:

			mdecoder->codec_id = AV_CODEC_ID_MPEG2VIDEO;

			break;

		case TSMF_SUB_TYPE_WMV3:

			mdecoder->codec_id = AV_CODEC_ID_WMV3;

			break;

		case TSMF_SUB_TYPE_AAC:

			mdecoder->codec_id = AV_CODEC_ID_AAC;

			/* For AAC the pFormat is a HEAACWAVEINFO struct, and the codec data

			   is at the end of it. See

			   http://msdn.microsoft.com/en-us/library/dd757806.aspx */

			if (media_type->ExtraData)

			{

				media_type->ExtraData += 12;

				media_type->ExtraDataSize -= 12;

			}

			break;

		case TSMF_SUB_TYPE_H264:

		case TSMF_SUB_TYPE_AVC1:

			mdecoder->codec_id = AV_CODEC_ID_H264;

			break;

		case TSMF_SUB_TYPE_AC3:

			mdecoder->codec_id = AV_CODEC_ID_AC3;

			break;

		default:

			return FALSE;

	}

	if (!tsmf_ffmpeg_init_context(decoder))

		return FALSE;

	if (!tsmf_ffmpeg_init_stream(decoder, media_type))

		return FALSE;

	if (!tsmf_ffmpeg_prepare(decoder))

		return FALSE;

	return TRUE;

}

static BOOL tsmf_ffmpeg_decode_video(ITSMFDecoder* decoder, const BYTE* data, UINT32 data_size,

                                     UINT32 extensions)

{

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	int decoded;

	int len;

	AVFrame* frame;

	BOOL ret = TRUE;

#if LIBAVCODEC_VERSION_MAJOR < 52 || \

    (LIBAVCODEC_VERSION_MAJOR == 52 && LIBAVCODEC_VERSION_MINOR <= 20)

	len = avcodec_decode_video(mdecoder->codec_context, mdecoder->frame, &decoded, data, data_size);

#else

	{

		AVPacket pkt;

		av_init_packet(&pkt);

		pkt.data = (BYTE*)data;

		pkt.size = data_size;

		if (extensions & TSMM_SAMPLE_EXT_CLEANPOINT)

			pkt.flags |= AV_PKT_FLAG_KEY;

		len = avcodec_decode_video2(mdecoder->codec_context, mdecoder->frame, &decoded, &pkt);

	}

#endif

	if (len < 0)

	{

		WLog_ERR(TAG, "data_size %" PRIu32 ", avcodec_decode_video failed (%d)", data_size, len);

		ret = FALSE;

	}

	else if (!decoded)

	{

		WLog_ERR(TAG, "data_size %" PRIu32 ", no frame is decoded.", data_size);

		ret = FALSE;

	}

	else

	{

		DEBUG_TSMF("linesize[0] %d linesize[1] %d linesize[2] %d linesize[3] %d "

		           "pix_fmt %d width %d height %d",

		           mdecoder->frame->linesize[0], mdecoder->frame->linesize[1],

		           mdecoder->frame->linesize[2], mdecoder->frame->linesize[3],

		           mdecoder->codec_context->pix_fmt, mdecoder->codec_context->width,

		           mdecoder->codec_context->height);

		mdecoder->decoded_size =

		    avpicture_get_size(mdecoder->codec_context->pix_fmt, mdecoder->codec_context->width,

		                       mdecoder->codec_context->height);

		mdecoder->decoded_data = calloc(1, mdecoder->decoded_size);

		if (!mdecoder->decoded_data)

			return FALSE;

#if LIBAVCODEC_VERSION_MAJOR < 55

		frame = avcodec_alloc_frame();

#else

		frame = av_frame_alloc();

#endif

		avpicture_fill((AVPicture*)frame, mdecoder->decoded_data, mdecoder->codec_context->pix_fmt,

		               mdecoder->codec_context->width, mdecoder->codec_context->height);

		av_picture_copy((AVPicture*)frame, (AVPicture*)mdecoder->frame,

		                mdecoder->codec_context->pix_fmt, mdecoder->codec_context->width,

		                mdecoder->codec_context->height);

		av_free(frame);

	}

	return ret;

}

static BOOL tsmf_ffmpeg_decode_audio(ITSMFDecoder* decoder, const BYTE* data, UINT32 data_size,

                                     UINT32 extensions)

{

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	int len;

	int frame_size;

	UINT32 src_size;

	const BYTE* src;

	BYTE* dst;

	int dst_offset;

#if 0

	WLog_DBG(TAG, ("tsmf_ffmpeg_decode_audio: data_size %"PRIu32"", data_size));

	int i;

	for (i = 0; i < data_size; i++)

	{

		WLog_DBG(TAG, ("%02"PRIX8"", data[i]));

		if (i % 16 == 15)

			WLog_DBG(TAG, ("\n"));

	}

#endif

	if (mdecoder->decoded_size_max == 0)

		mdecoder->decoded_size_max = MAX_AUDIO_FRAME_SIZE + 16;

	mdecoder->decoded_data = calloc(1, mdecoder->decoded_size_max);

	if (!mdecoder->decoded_data)

		return FALSE;

	/* align the memory for SSE2 needs */

	dst = (BYTE*)(((uintptr_t)mdecoder->decoded_data + 15) & ~0x0F);

	dst_offset = dst - mdecoder->decoded_data;

	src = data;

	src_size = data_size;

	while (src_size > 0)

	{

		/* Ensure enough space for decoding */

		if (mdecoder->decoded_size_max - mdecoder->decoded_size < MAX_AUDIO_FRAME_SIZE)

		{

			BYTE* tmp_data;

			tmp_data = realloc(mdecoder->decoded_data, mdecoder->decoded_size_max * 2 + 16);

			if (!tmp_data)

				return FALSE;

			mdecoder->decoded_size_max = mdecoder->decoded_size_max * 2 + 16;

			mdecoder->decoded_data = tmp_data;

			dst = (BYTE*)(((uintptr_t)mdecoder->decoded_data + 15) & ~0x0F);

			if (dst - mdecoder->decoded_data != dst_offset)

			{

				/* re-align the memory if the alignment has changed after realloc */

				memmove(dst, mdecoder->decoded_data + dst_offset, mdecoder->decoded_size);

				dst_offset = dst - mdecoder->decoded_data;

			}

			dst += mdecoder->decoded_size;

		}

		frame_size = mdecoder->decoded_size_max - mdecoder->decoded_size;

#if LIBAVCODEC_VERSION_MAJOR < 52 || \

    (LIBAVCODEC_VERSION_MAJOR == 52 && LIBAVCODEC_VERSION_MINOR <= 20)

		len = avcodec_decode_audio2(mdecoder->codec_context, (int16_t*)dst, &frame_size, src,

		                            src_size);

#else

		{

#if LIBAVCODEC_VERSION_MAJOR < 55

			AVFrame* decoded_frame = avcodec_alloc_frame();

#else

			AVFrame* decoded_frame = av_frame_alloc();

#endif

			int got_frame = 0;

			AVPacket pkt;

			av_init_packet(&pkt);

			pkt.data = (BYTE*)src;

			pkt.size = src_size;

			len = avcodec_decode_audio4(mdecoder->codec_context, decoded_frame, &got_frame, &pkt);

			if (len >= 0 && got_frame)

			{

				frame_size = av_samples_get_buffer_size(NULL, mdecoder->codec_context->channels,

				                                        decoded_frame->nb_samples,

				                                        mdecoder->codec_context->sample_fmt, 1);

				memcpy(dst, decoded_frame->data[0], frame_size);

			}

			else

			{

				frame_size = 0;

			}

			av_free(decoded_frame);

		}

#endif

		if (len > 0)

		{

			src += len;

			src_size -= len;

		}

		if (frame_size > 0)

		{

			mdecoder->decoded_size += frame_size;

			dst += frame_size;

		}

	}

	if (mdecoder->decoded_size == 0)

	{

		free(mdecoder->decoded_data);

		mdecoder->decoded_data = NULL;

	}

	else if (dst_offset)

	{

		/* move the aligned decoded data to original place */

		memmove(mdecoder->decoded_data, mdecoder->decoded_data + dst_offset,

		        mdecoder->decoded_size);

	}

	DEBUG_TSMF("data_size %" PRIu32 " decoded_size %" PRIu32 "", data_size, mdecoder->decoded_size);

	return TRUE;

}

static BOOL tsmf_ffmpeg_decode(ITSMFDecoder* decoder, const BYTE* data, UINT32 data_size,

                               UINT32 extensions)

{

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	if (mdecoder->decoded_data)

	{

		free(mdecoder->decoded_data);

		mdecoder->decoded_data = NULL;

	}

	mdecoder->decoded_size = 0;

	switch (mdecoder->media_type)

	{

		case AVMEDIA_TYPE_VIDEO:

			return tsmf_ffmpeg_decode_video(decoder, data, data_size, extensions);

		case AVMEDIA_TYPE_AUDIO:

			return tsmf_ffmpeg_decode_audio(decoder, data, data_size, extensions);

		default:

			WLog_ERR(TAG, "unknown media type.");

			return FALSE;

	}

}

static BYTE* tsmf_ffmpeg_get_decoded_data(ITSMFDecoder* decoder, UINT32* size)

{

	BYTE* buf;

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	*size = mdecoder->decoded_size;

	buf = mdecoder->decoded_data;

	mdecoder->decoded_data = NULL;

	mdecoder->decoded_size = 0;

	return buf;

}

static UINT32 tsmf_ffmpeg_get_decoded_format(ITSMFDecoder* decoder)

{

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	switch (mdecoder->codec_context->pix_fmt)

	{

		case AV_PIX_FMT_YUV420P:

			return RDP_PIXFMT_I420;

		default:

			WLog_ERR(TAG, "unsupported pixel format %u", mdecoder->codec_context->pix_fmt);

			return (UINT32)-1;

	}

}

static BOOL tsmf_ffmpeg_get_decoded_dimension(ITSMFDecoder* decoder, UINT32* width, UINT32* height)

{

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	if (mdecoder->codec_context->width > 0 && mdecoder->codec_context->height > 0)

	{

		*width = mdecoder->codec_context->width;

		*height = mdecoder->codec_context->height;

		return TRUE;

	}

	else

	{

		return FALSE;

	}

}

static void tsmf_ffmpeg_free(ITSMFDecoder* decoder)

{

	TSMFFFmpegDecoder* mdecoder = (TSMFFFmpegDecoder*)decoder;

	if (mdecoder->frame)

		av_free(mdecoder->frame);

	free(mdecoder->decoded_data);

	if (mdecoder->codec_context)

	{

		if (mdecoder->prepared)

			avcodec_close(mdecoder->codec_context);

		free(mdecoder->codec_context->extradata);

		av_free(mdecoder->codec_context);

	}

	free(decoder);

}

static INIT_ONCE g_Initialized = INIT_ONCE_STATIC_INIT;

static BOOL CALLBACK InitializeAvCodecs(PINIT_ONCE once, PVOID param, PVOID* context)

{

	avcodec_register_all();

	return TRUE;

}

#ifdef BUILTIN_CHANNELS

#define freerdp_tsmf_client_subsystem_entry ffmpeg_freerdp_tsmf_client_decoder_subsystem_entry

#else

#define freerdp_tsmf_client_subsystem_entry FREERDP_API freerdp_tsmf_client_decoder_subsystem_entry

#endif

ITSMFDecoder* freerdp_tsmf_client_subsystem_entry(void)

{

	TSMFFFmpegDecoder* decoder;

	InitOnceExecuteOnce(&g_Initialized, InitializeAvCodecs, NULL, NULL);

	WLog_DBG(TAG, "TSMFDecoderEntry FFMPEG");

	decoder = (TSMFFFmpegDecoder*)calloc(1, sizeof(TSMFFFmpegDecoder));

	if (!decoder)

		return NULL;

	decoder->iface.SetFormat = tsmf_ffmpeg_set_format;

	decoder->iface.Decode = tsmf_ffmpeg_decode;

	decoder->iface.GetDecodedData = tsmf_ffmpeg_get_decoded_data;

	decoder->iface.GetDecodedFormat = tsmf_ffmpeg_get_decoded_format;

	decoder->iface.GetDecodedDimension = tsmf_ffmpeg_get_decoded_dimension;

	decoder->iface.Free = tsmf_ffmpeg_free;

	return (ITSMFDecoder*)decoder;

}