/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/.
 */

#pragma once

#include <memory>
#include <cstring>

namespace OpenVulkano
{
	namespace YuvUtils
	{
		void NV12FromChromaPlanes(unsigned char* __restrict src, unsigned char* __restrict dest, int chromaChannelPixelCount)
		{
			for(int i = 0; i < chromaChannelPixelCount; i++)
			{
				dest[i * 2] = src[i];
			}
			for(int i=0; i < chromaChannelPixelCount; i++)
			{
				dest[i * 2 + 1] = src[i + chromaChannelPixelCount];
			}
		}

		void ChromaPlanesFromNV12(unsigned char* __restrict src, unsigned char* __restrict dest, int chromaChannelPixelCount)
		{
			for (int i = 0; i < chromaChannelPixelCount; i++)
			{
				dest[i] = src[i * 2];
				dest[i + chromaChannelPixelCount] = src[i * 2 + 1];
			}
		}

		std::unique_ptr<uint8_t[]> PlansFromNV12(
				unsigned char* __restrict srcLum, unsigned char* __restrict srcChroma, uint32_t sizeLumX, uint32_t sizeLumY,
				uint32_t sizeChromaComponentX, uint32_t sizeChromaComponentY, uint32_t strideX = 1, uint32_t strideY = 1)
		{
			std::unique_ptr<uint8_t[]> buffer(new uint8_t[(sizeLumX / strideX) * (sizeLumY / strideY) + 2 * (sizeChromaComponentX / strideX) * (sizeChromaComponentY / strideY)]);
			if (strideX == 1 && strideY == 1)
			{
				memcpy(buffer.get(), srcLum, sizeLumX * sizeLumY);
				ChromaPlanesFromNV12(srcChroma, buffer.get() + sizeLumX * sizeLumY, sizeChromaComponentX * sizeChromaComponentY);
			}
			else
			{
				uint8_t* dest = buffer.get();
				size_t rowStride = strideY * sizeLumX;
				for(uint32_t y = 0; y < sizeLumY; y += strideY, srcLum += rowStride)
				{
					for(uint32_t x = 0; x < sizeLumX; x += strideX, dest++)
					{
						*dest = srcLum[x];
					}
				}
				uint32_t chromaCount = (sizeChromaComponentX / strideX) * (sizeChromaComponentY / strideY);
				rowStride = strideY * sizeChromaComponentX * 2;
				for(uint32_t y = 0; y < sizeChromaComponentY; y += strideY, srcChroma += rowStride)
				{
					for (uint32_t x = 0; x < sizeChromaComponentX * 2; x += 2 * strideX, dest++)
					{
						dest[0] = srcChroma[x];
						dest[chromaCount] = srcChroma[x + 1];
					}
				}
			}
			return buffer;
		}
	}
}