/*
 * 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 "IO/Archive/ArchiveReader.hpp"
#include "IO/Files/Pfm.hpp"
#include "IO/Files/Pnm.hpp"
#include <string>
#include <fstream>

namespace OpenVulkano::AR::Playback
{
	struct DepthImage
	{
		PfmImage depth;
		PnmImage confidence;
	};

	struct ColorImg
	{
		enum ChannelCount { GRAY = 1, RGBA = 4, YUV = 2, NV12 = -2 };

		std::unique_ptr<uint8_t[]> dataPtr;
		uint8_t* data = nullptr;
		uint8_t* dataUV = nullptr;
		int cols = 0, rows = 0, channels = 0;
		
		void Allocate()
		{
			if (data) throw std::runtime_error("Data already set!");
			dataPtr.reset(new uint8_t[cols * rows * abs(channels)]);
			data = dataPtr.get();
		}

		bool UseRGB() const { return channels == RGBA; }

		void Decode(const Array<char>& data);
	};

	class ArPlaybackReader final
	{
		static constexpr std::string_view TAR_EXTENSIONS_REGEX = R"(\.(tar(\.gz|\.bz2|\.zst)?|tgz|tbz|tb2|tbz2))";

		ArchiveReader m_archiveMetadata, m_archiveColor, m_archiveDepth, m_archiveConfidence;
		size_t m_imgTotalSize = 0, m_imgReadSize = 0;
		bool m_hasConfidence = false;
		
	public:
		ArPlaybackReader(const std::filesystem::path& recDir)
		{
			const std::string extensions = R"((_\d+|\.part\d+)?)" + std::string(TAR_EXTENSIONS_REGEX);
			m_archiveMetadata.Open(recDir, ".*meta(data)?" + extensions);
			m_archiveColor.Open(recDir, ".*(color|image)" + extensions, &m_imgTotalSize);
			m_archiveDepth.Open(recDir, ".*depth" + extensions);
			m_hasConfidence = m_archiveConfidence.Open(recDir, ".*conf(idence)?" + extensions);
		}

		int GetNextFrameId() const
		{
			const std::string& name = m_archiveMetadata.GetDescription().path;
			return std::stoi(name.substr(0, name.length() - 5));
		}

		Array<char> ReadMetadata()
		{
			return std::move(m_archiveMetadata.GetNextFile()->second);
		}

		ColorImg ReadColorImage()
		{
			ColorImg img;
			img.channels = ColorImg::RGBA; // TODO test if defaulting to NV12
			ReadColorImage(img);
			return img;
		}

		void ReadColorImage(ColorImg& img)
		{
			std::optional<std::pair<FileDescription, Array<char>>> file = ReadColorImageRaw();
			img.Decode(file->second);
			m_imgReadSize += 1000 + file.value().first.size;
		}

		std::optional<std::pair<FileDescription, Array<char>>> ReadColorImageRaw() { return m_archiveColor.GetNextFile(); }

		auto ReadDepthImageRaw() { return m_archiveDepth.GetNextFile(); }

		auto ReadConfidenceRaw() { return m_archiveConfidence.GetNextFile(); }

		DepthImage ReadDepthImage()
		{
			DepthImage img;
			m_archiveDepth.GetNextFileAsStream([&img](const FileDescription& desc, std::istream& stream) { img.depth.TryRead(stream); });
			if (!m_archiveConfidence.GetNextFileAsStream([&img](const FileDescription&, std::istream& stream) { img.confidence.TryRead(stream); }))
			{ // No confidence image available
				img.confidence.header.width = img.confidence.header.height = 1;
				img.confidence.image = std::make_unique<char[]>(m_imgTotalSize);
				img.confidence.image[0] = 2; // TODO
			}
			
			return img;
		}
		
		[[nodiscard]] double GetProgress() const
		{
			return static_cast<double>(m_imgReadSize) / static_cast<double>(m_imgTotalSize);
		}

		[[nodiscard]] bool HasNext() const
		{
			return m_archiveMetadata.HasNext() && m_archiveDepth.HasNext() && m_archiveColor.HasNext()
					&& (m_archiveConfidence.HasNext() || !m_hasConfidence);
		}
	};
}