/*
 * filter_imageconvert.c -- colorspace and pixel format converter
 * Copyright (C) 2009-2021 Meltytech, LLC
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 */

#include <framework/mlt_filter.h>
#include <framework/mlt_frame.h>
#include <framework/mlt_image.h>
#include <framework/mlt_log.h>
#include <framework/mlt_pool.h>

#include <stdlib.h>

/** This macro converts a YUV value to the RGB color space. */
#define RGB2YUV_601_UNSCALED(r, g, b, y, u, v)\
  y = (299*r + 587*g + 114*b) >> 10;\
  u = ((-169*r - 331*g + 500*b) >> 10) + 128;\
  v = ((500*r - 419*g - 81*b) >> 10) + 128;\
  y = y < 16 ? 16 : y;\
  u = u < 16 ? 16 : u;\
  v = v < 16 ? 16 : v;\
  y = y > 235 ? 235 : y;\
  u = u > 240 ? 240 : u;\
  v = v > 240 ? 240 : v

/** This macro converts a YUV value to the RGB color space. */
#define YUV2RGB_601_UNSCALED( y, u, v, r, g, b ) \
  r = ((1024 * y + 1404 * ( v - 128 ) ) >> 10 ); \
  g = ((1024 * y - 715 * ( v - 128 ) - 345 * ( u - 128 ) ) >> 10 ); \
  b = ((1024 * y + 1774 * ( u - 128 ) ) >> 10 ); \
  r = r < 0 ? 0 : r > 255 ? 255 : r; \
  g = g < 0 ? 0 : g > 255 ? 255 : g; \
  b = b < 0 ? 0 : b > 255 ? 255 : b;

#define SCALED 1
#if SCALED
#define RGB2YUV_601 RGB2YUV_601_SCALED
#define YUV2RGB_601 YUV2RGB_601_SCALED
#else
#define RGB2YUV_601 RGB2YUV_601_UNSCALED
#define YUV2RGB_601 YUV2RGB_601_UNSCALED
#endif

static void convert_yuv422_to_rgba( mlt_image src, mlt_image dst )
{
	int yy, uu, vv;
	int r,g,b;

	mlt_image_set_values( dst, NULL, mlt_image_rgba, src->width, src->height );
	mlt_image_alloc_data( dst );

	for ( int line = 0; line < src->height; line++ )
	{
		uint8_t* pSrc = src->planes[0] + src->strides[0] * line;
		uint8_t* pAlpha = src->planes[3] + src->strides[3] * line;
		uint8_t* pDst = dst->planes[0] + dst->strides[0] * line;
		int total = src->width / 2 + 1;

		if ( pAlpha )
			while ( --total )
			{
				yy = pSrc[0];
				uu = pSrc[1];
				vv = pSrc[3];
				YUV2RGB_601( yy, uu, vv, r, g, b );
				pDst[0] = r;
				pDst[1] = g;
				pDst[2] = b;
				pDst[3] = *pAlpha++;
				yy = pSrc[2];
				YUV2RGB_601( yy, uu, vv, r, g, b );
				pDst[4] = r;
				pDst[5] = g;
				pDst[6] = b;
				pDst[7] = *pAlpha++;
				pSrc += 4;
				pDst += 8;
			}
		else
			while ( --total )
			{
				yy = pSrc[0];
				uu = pSrc[1];
				vv = pSrc[3];
				YUV2RGB_601( yy, uu, vv, r, g, b );
				pDst[0] = r;
				pDst[1] = g;
				pDst[2] = b;
				pDst[3] = 0xff;
				yy = pSrc[2];
				YUV2RGB_601( yy, uu, vv, r, g, b );
				pDst[4] = r;
				pDst[5] = g;
				pDst[6] = b;
				pDst[7] = 0xff;
				pSrc += 4;
				pDst += 8;
			}
	}
}

static void convert_yuv422_to_rgb( mlt_image src, mlt_image dst )
{
	int yy, uu, vv;
	int r,g,b;

	mlt_image_set_values( dst, NULL, mlt_image_rgb, src->width, src->height );
	mlt_image_alloc_data( dst );

	for ( int line = 0; line < src->height; line++ )
	{
		uint8_t* pSrc = src->planes[0] + src->strides[0] * line;
		uint8_t* pDst = dst->planes[0] + dst->strides[0] * line;
		int total = src->width / 2 + 1;
		while ( --total )
		{
			yy = pSrc[0];
			uu = pSrc[1];
			vv = pSrc[3];
			YUV2RGB_601( yy, uu, vv, r, g, b );
			pDst[0] = r;
			pDst[1] = g;
			pDst[2] = b;
			yy = pSrc[2];
			YUV2RGB_601( yy, uu, vv, r, g, b );
			pDst[3] = r;
			pDst[4] = g;
			pDst[5] = b;
			pSrc += 4;
			pDst += 6;
		}
	}
}

static void convert_rgba_to_yuv422( mlt_image src, mlt_image dst )
{
	int y0, y1, u0, u1, v0, v1;
	int r, g, b;

	mlt_image_set_values( dst, NULL, mlt_image_yuv422, src->width, src->height );
	mlt_image_alloc_data( dst );
	mlt_image_alloc_alpha( dst );

	for ( int line = 0; line < src->height; line++ )
	{
		uint8_t* pSrc = src->planes[0] + src->strides[0] * line;
		uint8_t* pDst = dst->planes[0] + dst->strides[0] * line;
		uint8_t* pAlpha = dst->planes[3] + dst->strides[3] * line;
		int j = src->width / 2 + 1;
		while ( --j )
		{
			r = *pSrc++;
			g = *pSrc++;
			b = *pSrc++;
			*pAlpha++ = *pSrc++;
			RGB2YUV_601( r, g, b, y0, u0 , v0 );
			r = *pSrc++;
			g = *pSrc++;
			b = *pSrc++;
			*pAlpha++ = *pSrc++;
			RGB2YUV_601( r, g, b, y1, u1 , v1 );
			*pDst++ = y0;
			*pDst++ = (u0+u1) >> 1;
			*pDst++ = y1;
			*pDst++ = (v0+v1) >> 1;
		}
		if ( src->width % 2 )
		{
			r = *pSrc++;
			g = *pSrc++;
			b = *pSrc++;
			*pAlpha++ = *pSrc++;
			RGB2YUV_601( r, g, b, y0, u0 , v0 );
			*pDst++ = y0;
			*pDst++ = u0;
		}
	}
}

static void convert_rgb_to_yuv422( mlt_image src, mlt_image dst )
{
	int y0, y1, u0, u1, v0, v1;
	int r, g, b;

	mlt_image_set_values( dst, NULL, mlt_image_yuv422, src->width, src->height );
	mlt_image_alloc_data( dst );

	for ( int line = 0; line < src->height; line++ )
	{
		uint8_t* pSrc = src->planes[0] + src->strides[0] * line;
		uint8_t* pDst = dst->planes[0] + dst->strides[0] * line;
		int j = src->width / 2 + 1;
		while ( --j )
		{
			r = *pSrc++;
			g = *pSrc++;
			b = *pSrc++;
			RGB2YUV_601( r, g, b, y0, u0 , v0 );
			r = *pSrc++;
			g = *pSrc++;
			b = *pSrc++;
			RGB2YUV_601( r, g, b, y1, u1 , v1 );
			*pDst++ = y0;
			*pDst++ = (u0+u1) >> 1;
			*pDst++ = y1;
			*pDst++ = (v0+v1) >> 1;
		}
		if ( src->width % 2 )
		{
			r = *pSrc++;
			g = *pSrc++;
			b = *pSrc++;
			RGB2YUV_601( r, g, b, y0, u0 , v0 );
			*pDst++ = y0;
			*pDst++ = u0;
		}
	}
}

static void convert_yuv420p_to_yuv422( mlt_image src, mlt_image dst )
{
	mlt_image_set_values( dst, NULL, mlt_image_yuv422, src->width, src->height );
	mlt_image_alloc_data( dst );

	for ( int line = 0; line < src->height; line++ )
	{
		uint8_t* pSrcY = src->planes[0] + src->strides[0] * line;
		uint8_t* pSrcU = src->planes[1] + src->strides[1] * line / 2;
		uint8_t* pSrcV = src->planes[2] + src->strides[2] * line / 2;
		uint8_t* pDst = dst->planes[0] + dst->strides[0] * line;
		int j = src->width / 2 + 1;
		while ( --j )
		{
			*pDst++ = *pSrcY++;
			*pDst++ = *pSrcU++;
			*pDst++ = *pSrcY++;
			*pDst++ = *pSrcV++;
		}
	}
}

static void convert_yuv420p_to_rgb( mlt_image src, mlt_image dst )
{
	int yy, uu, vv;
	int r,g,b;

	mlt_image_set_values( dst, NULL, mlt_image_rgb, src->width, src->height );
	mlt_image_alloc_data( dst );

	for ( int line = 0; line < src->height; line++ )
	{
		uint8_t* pSrcY = src->planes[0] + src->strides[0] * line;
		uint8_t* pSrcU = src->planes[1] + src->strides[1] * line / 2;
		uint8_t* pSrcV = src->planes[2] + src->strides[2] * line / 2;
		uint8_t* pDst = dst->planes[0] + dst->strides[0] * line;
		int total = src->width / 2 + 1;
		while ( --total )
		{
			yy = *pSrcY++;
			uu = *pSrcU++;
			vv = *pSrcV++;
			YUV2RGB_601( yy, uu, vv, r, g, b );
			pDst[0] = r;
			pDst[1] = g;
			pDst[2] = b;
			yy = *pSrcY++;
			YUV2RGB_601( yy, uu, vv, r, g, b );
			pDst[3] = r;
			pDst[4] = g;
			pDst[5] = b;
			pDst += 6;
		}
	}
}

static void convert_yuv420p_to_rgba( mlt_image src, mlt_image dst )
{
	int yy, uu, vv;
	int r,g,b;

	mlt_image_set_values( dst, NULL, mlt_image_rgba, src->width, src->height );
	mlt_image_alloc_data( dst );

	for ( int line = 0; line < src->height; line++ )
	{
		uint8_t* pSrcY = src->planes[0] + src->strides[0] * line;
		uint8_t* pSrcU = src->planes[1] + src->strides[1] * line / 2;
		uint8_t* pSrcV = src->planes[2] + src->strides[2] * line / 2;
		uint8_t* pSrcA = src->planes[3] + src->strides[3] * line;
		uint8_t* pDst = dst->planes[0] + dst->strides[0] * line;
		int total = src->width / 2 + 1;
		if ( pSrcA )
			while ( --total )
			{
				yy = *pSrcY++;
				uu = *pSrcU++;
				vv = *pSrcV++;
				YUV2RGB_601( yy, uu, vv, r, g, b );
				pDst[0] = r;
				pDst[1] = g;
				pDst[2] = b;
				pDst[3] = *pSrcA++;
				yy = *pSrcY++;
				YUV2RGB_601( yy, uu, vv, r, g, b );
				pDst[4] = r;
				pDst[5] = g;
				pDst[6] = b;
				pDst[7] = *pSrcA++;
				pDst += 8;
			}
		else
			while ( --total )
			{
				yy = *pSrcY++;
				uu = *pSrcU++;
				vv = *pSrcV++;
				YUV2RGB_601( yy, uu, vv, r, g, b );
				pDst[0] = r;
				pDst[1] = g;
				pDst[2] = b;
				pDst[3] = 0xff;
				yy = *pSrcY++;
				YUV2RGB_601( yy, uu, vv, r, g, b );
				pDst[4] = r;
				pDst[5] = g;
				pDst[6] = b;
				pDst[7] = 0xff;
				pDst += 8;
			}
	}
}

static void convert_yuv422_to_yuv420p( mlt_image src, mlt_image dst )
{
	int lines = src->height;
	int pixels = src->width;

	mlt_image_set_values( dst, NULL, mlt_image_yuv420p, src->width, src->height );
	mlt_image_alloc_data( dst );

	// Y
	for ( int line = 0; line < lines; line++ )
	{
		uint8_t* pSrc = src->planes[0] + src->strides[0] * line;
		uint8_t* pDst = dst->planes[0] + dst->strides[0] * line;
		for ( int pixel = 0; pixel < pixels; pixel++ )
		{
			*pDst++ = *pSrc;
			pSrc += 2;
		}
	}

	lines = src->height / 2;
	pixels = src->width / 2;

	// U
	for ( int line = 0; line < lines; line++ )
	{
		uint8_t* pSrc = src->planes[0] + src->strides[0] * line * 2 + 1;
		uint8_t* pDst = dst->planes[1] + dst->strides[1] * line;
		for ( int pixel = 0; pixel < pixels; pixel++ )
		{
			*pDst++ = *pSrc;
			pSrc += 4;
		}
	}

	// V
	for ( int line = 0; line < lines; line++ )
	{
		uint8_t* pSrc = src->planes[0] + src->strides[0] * line * 2 + 3;
		uint8_t* pDst = dst->planes[2] + dst->strides[2] * line;
		for ( int pixel = 0; pixel < pixels; pixel++ )
		{
			*pDst++ = *pSrc;
			pSrc += 4;
		}
	}
}

static void convert_rgb_to_rgba( mlt_image src, mlt_image dst )
{
	mlt_image_set_values( dst, NULL, mlt_image_rgba, src->width, src->height );
	mlt_image_alloc_data( dst );

	for ( int line = 0; line < src->height; line++ )
	{
		uint8_t* pSrc = src->planes[0] + src->strides[0] * line;
		uint8_t* pAlpha = src->planes[3] + src->strides[3] * line;
		uint8_t* pDst = dst->planes[0] + dst->strides[0] * line;
		int total = src->width + 1;
		if ( pAlpha )
			while ( --total )
			{
				*pDst++ = pSrc[0];
				*pDst++ = pSrc[1];
				*pDst++ = pSrc[2];
				*pDst++ = *pAlpha++;
				pSrc += 3;
			}
		else
			while ( --total )
			{
				*pDst++ = pSrc[0];
				*pDst++ = pSrc[1];
				*pDst++ = pSrc[2];
				*pDst++ = 0xff;
				pSrc += 3;
			}
	}
}

static void convert_rgba_to_rgb( mlt_image src, mlt_image dst )
{
	mlt_image_set_values( dst, NULL, mlt_image_rgb, src->width, src->height );
	mlt_image_alloc_data( dst );
	mlt_image_alloc_alpha( dst );

	for ( int line = 0; line < src->height; line++ )
	{
		uint8_t* pSrc = src->planes[0] + src->strides[0] * line;
		uint8_t* pDst = dst->planes[0] + dst->strides[0] * line;
		uint8_t* pAlpha = dst->planes[3] + dst->strides[3] * line;
		int total = src->width + 1;
		while ( --total )
		{
			*pDst++ = pSrc[0];
			*pDst++ = pSrc[1];
			*pDst++ = pSrc[2];
			*pAlpha++ = pSrc[3];
			pSrc += 4;
		}
	}
}

typedef void ( *conversion_function )( mlt_image src, mlt_image dst );

static conversion_function conversion_matrix[ mlt_image_invalid - 1 ][ mlt_image_invalid - 1 ] = {
	{ NULL, convert_rgb_to_rgba, convert_rgb_to_yuv422, NULL, NULL, NULL, NULL },
	{ convert_rgba_to_rgb, NULL, convert_rgba_to_yuv422, NULL, NULL, NULL, NULL },
	{ convert_yuv422_to_rgb, convert_yuv422_to_rgba, NULL, convert_yuv422_to_yuv420p, NULL, NULL, NULL },
	{ convert_yuv420p_to_rgb, convert_yuv420p_to_rgba, convert_yuv420p_to_yuv422, NULL, NULL, NULL, NULL },
	{ NULL, NULL, NULL, NULL, NULL, NULL, NULL },
	{ NULL, NULL, NULL, NULL, NULL, NULL, NULL },
	{ NULL, NULL, NULL, NULL, NULL, NULL, NULL },
};

static int convert_image( mlt_frame frame, uint8_t **buffer, mlt_image_format *format, mlt_image_format requested_format )
{
	int error = 0;
	mlt_properties properties = MLT_FRAME_PROPERTIES( frame );
	int width = mlt_properties_get_int( properties, "width" );
	int height = mlt_properties_get_int( properties, "height" );

	if ( *format != requested_format )
	{
		conversion_function converter = conversion_matrix[ *format - 1 ][ requested_format - 1 ];

		mlt_log_debug( NULL, "[filter imageconvert] %s -> %s @ %dx%d\n",
			mlt_image_format_name( *format ), mlt_image_format_name( requested_format ),
			width, height );
		if ( converter )
		{
			struct mlt_image_s src;
			struct mlt_image_s dst;
			mlt_image_set_values( &src, *buffer, *format, width, height );
			if ( requested_format == mlt_image_rgba && mlt_frame_get_alpha( frame ) )
			{
				// imageconvert leaves the alpha buffer alone except in the case of rgba.
				// For rgba input, an alpha buffer will be created and added to the frame.
				// For rgba output, the alpha buffer will be copied to the rgba and the buffer is removed from the frame.
				src.planes[3] = mlt_frame_get_alpha( frame );
				src.strides[3] = src.width;
			}
			converter( &src, &dst );
			mlt_frame_set_image( frame, dst.data, 0, dst.release_data );
			if ( requested_format == mlt_image_rgba )
			{
				// Clear the alpha buffer on the frame
				mlt_frame_set_alpha( frame, NULL, 0, NULL );
			}
			else if ( dst.alpha )
			{
				mlt_frame_set_alpha( frame, dst.alpha, 0, dst.release_alpha );
			}
			*buffer = dst.data;
			*format = dst.format;
		}
		else
		{
			mlt_log_error( NULL, "imageconvert: no conversion from %s to %s\n", mlt_image_format_name(*format), mlt_image_format_name(requested_format));
			error = 1;
		}
	}

	return error;
}

/** Filter processing.
*/

static mlt_frame filter_process( mlt_filter filter, mlt_frame frame )
{
	if ( !frame->convert_image )
		frame->convert_image = convert_image;
	return frame;
}

/** Constructor for the filter.
*/

mlt_filter filter_imageconvert_init( mlt_profile profile, mlt_service_type type, const char *id, char *arg )
{
	mlt_filter filter = calloc( 1, sizeof( struct mlt_filter_s ) );
	if ( mlt_filter_init( filter, filter ) == 0 )
	{
		filter->process = filter_process;
	}
	return filter;
}

