/*
 * 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
	{
		inline void NV12FromChromaPlanes(unsigned char* __restrict src, unsigned char* __restrict dest, size_t chromaChannelPixelCount)
		{
			for(size_t i = 0; i < chromaChannelPixelCount; i++)
			{
				dest[i * 2] = src[i];
				dest[i * 2 + 1] = src[i + chromaChannelPixelCount];
			}
		}

		inline void ChromaPlanesFromNV12(unsigned char* __restrict src, unsigned char* __restrict dest, uint32_t chromaChannelWidth, uint32_t chromaChannelHeight, uint32_t chromaRowPadding)
		{
			uint8_t* dest2 = dest + chromaChannelWidth * chromaChannelHeight;
			for (uint32_t row = 0; row < chromaChannelHeight; row++)
			{
				for (uint32_t col = 0; col < chromaChannelWidth; col++, dest++, dest2++, src += 2)
				{
					*dest = src[0];
					*dest2 = src[1];
				}
				src += chromaRowPadding;
			}
		}

		inline 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, uint32_t srcLumRowPadding = 0, uint32_t srcChromaRowPadding = 0)
		{
			std::unique_ptr<uint8_t[]> buffer(new uint8_t[(sizeLumX / strideX) * (sizeLumY / strideY) + 2 * (sizeChromaComponentX / strideX) * (sizeChromaComponentY / strideY)]);
			if (strideX == 1 && strideY == 1)
			{
				if (srcLumRowPadding)
				{
					for(uint32_t y = 0, srcOffset = 0, destOffset = 0; y < sizeLumY; y++, srcOffset += sizeLumX + srcLumRowPadding, destOffset += sizeLumX)
					{
						memcpy(buffer.get() + destOffset, srcLum + srcOffset, sizeLumX);
					}
				}
				else
				{
					memcpy(buffer.get(), srcLum, sizeLumX * sizeLumY);
				}
				ChromaPlanesFromNV12(srcChroma, buffer.get() + sizeLumX * sizeLumY, sizeChromaComponentX, sizeChromaComponentY, srcChromaRowPadding);
			}
			else
			{
				uint8_t* dest = buffer.get();
				size_t rowStride = strideY * (sizeLumX + srcLumRowPadding);
				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 + srcChromaRowPadding);
				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;
		}
	}
}