/*
 * 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/.
 */

#include "ArPlaybackReader.hpp"

#include "Base/BlockProfiler.hpp"
#include "Base/Wrapper.hpp"
#include "Image/YuvUtils.hpp"

#if __has_include("turbojpeg.h")
#include <turbojpeg.h>

namespace OpenVulkano::AR::Playback
{
	void ColorImg::Decode(const Array<char>& dataIn)
	{
		//BlockProfiler profiler("Load jpeg");

		long unsigned int jpegSize = dataIn.Size();
		unsigned char* compressedImage = const_cast<uint8_t*>(reinterpret_cast<const uint8_t*>(dataIn.Data()));

		int jpegSubsamp;
		tjhandle jpegDecompressor = tjInitDecompress();
		tjDecompressHeader2(jpegDecompressor, compressedImage, jpegSize, &cols, &rows, &jpegSubsamp);

		if (UseRGB())
		{
			channels = 4;
			if (!data) Allocate();

			tjDecompress2(jpegDecompressor, compressedImage, jpegSize, data, cols, 0/*pitch*/, rows, TJPF_RGBA, TJFLAG_FASTDCT);
		}
		else if (channels == 1)
		{
			if (!data) Allocate();

			tjDecompress2(jpegDecompressor, compressedImage, jpegSize, data, cols, 0/*pitch*/, rows, TJPF_GRAY, TJFLAG_FASTDCT);
		}
		else if (channels != 0)
		{
			if (!data) Allocate();

			std::array<uint8_t*, 3> planes;
			size_t chromaSize = cols * rows / 4;
			planes[0] = data;
			if (!dataUV) dataUV = data + (cols * rows);
			Unique<uint8_t[]> chromaTmp;
			if (channels == ColorImg::NV12 && !dataPtr)
			{
				chromaTmp = std::make_unique<uint8_t[]>(cols * rows / 2);
				planes[1] = chromaTmp.get();
			}
			else
			{
				planes[1] = dataUV;
				if (channels == ColorImg::NV12) planes[1] += chromaSize;
			}
			planes[2] = planes[1] + chromaSize;
			std::array<int, 3> rowSize = { cols, cols / 2, cols / 2 };

			tjDecompressToYUVPlanes(jpegDecompressor, compressedImage, jpegSize, planes.data(),
				cols, rowSize.data(), rows, TJFLAG_FASTDCT);

			if (channels == ColorImg::NV12)
			{
				YuvUtils::NV12FromChromaPlanes(planes[1], data + (cols * rows), chromaSize);
			}
		}
		tjDestroy(jpegDecompressor);
	}
}

#else
#define STB_IMAGE_IMPLEMENTATION
#include <stb_image.h>

namespace OpenVulkano::AR::Playback
{
	ColorImg ArPlaybackReader::ReadColorImage()
	{
		ColorImg img;
		auto file = m_archiveColor.GetNextFile();
		img.dataPtr = std::shared_ptr<uint8_t>(
				stbi_load_from_memory(reinterpret_cast<stbi_uc*>(data.Data()), data.Size(), &img.cols, &img.rows, &img.channels, 3),
				&stbi_image_free);
		img.data = img.dataPtr.get();
		m_imgReadSize += 1000 + file.value().first.size;
		return img;
	}
}
#endif