/*
 * 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 "Base/Utils.hpp"
#include <string>
#include <sstream>
#include <iomanip>
#include <memory>

namespace OpenVulkano
{
	//TODO handle comments
	class PfmHeader
	{
	public:
		static constexpr std::string_view FILE_TYPE = "PFM";
		static constexpr std::string_view FILE_EXT = ".pfm";

		uint32_t width, height;
		float maxValue;
		bool color, littleEndian;

		PfmHeader() : width(0), height(0), maxValue(0), color(false), littleEndian(std::endian::native == std::endian::little) {}

		PfmHeader(uint32_t width, uint32_t height, float maxValue, bool color) :
				width(width), height(height), maxValue(maxValue), color(color), littleEndian(std::endian::native == std::endian::little)
		{}

		constexpr PfmHeader(uint32_t width, uint32_t height, float maxValue, bool color, bool littleEndian) :
				width(width), height(height), maxValue(maxValue), color(color), littleEndian(littleEndian)
		{}

		[[nodiscard]] std::string ToString() const
		{
			std::stringstream headerStream;
			headerStream << *this;
			return headerStream.str();
		}

		friend std::ostream& operator<< (std::ostream& outStream, const PfmHeader& pfmHeader)
		{
			outStream << (pfmHeader.color ? "PF" : "Pf") << '\n';
			outStream << pfmHeader.width << ' ' << pfmHeader.height << '\n';
			if constexpr (std::endian::native == std::endian::little) outStream << '-';
			outStream << std::fixed << std::setprecision(1) << pfmHeader.maxValue << '\n';
			return outStream;
		}

		friend std::istream& operator>> (std::istream& inStream, PfmHeader& pfmHeader)
		{
			std::string marker;
			inStream >> marker;
			if (marker == "PF") pfmHeader.color = true;
			else if (marker == "Pf") pfmHeader.color = false;
			else throw std::runtime_error("Malformed PFM header!");
			inStream >> pfmHeader.width >> pfmHeader.height;
			inStream >> pfmHeader.maxValue;
			inStream.get();
			if (pfmHeader.maxValue < 0)
			{
				pfmHeader.maxValue *= -1;
				pfmHeader.littleEndian = true;
			}
			else
			{
				pfmHeader.littleEndian = false;
			}
			return inStream;
		}

		[[nodiscard]] constexpr size_t GetElementCount() const
		{
			size_t size = height * width;
			if (color) size *= 3;
			return size;
		}

		[[nodiscard]] constexpr size_t GetImageSize() const
		{
			return GetElementCount() * sizeof(float);
		}
	};

	struct PfmImage
	{
		PfmHeader header;
		std::unique_ptr<float[]> image;

		void Read(std::istream& inStream)
		{
			inStream >> header;

			if (header.width == UINT32_MAX || header.height == UINT32_MAX)
			{
				header.width = header.height = 0;
				image = nullptr;
				return;
			}

			size_t size = header.GetElementCount();

			image = std::make_unique<float[]>(size);

			inStream.read(reinterpret_cast<char*>(image.get()), size * sizeof(float));

			if ((std::endian::native == std::endian::little) != header.littleEndian)
			{
				char* data = reinterpret_cast<char*>(image.get());
				for(size_t i = 0; i < size; i++)
				{
					size_t idx = i * sizeof(float);
					std::reverse(&data[idx], &data[idx + sizeof(float)]);
				}
			}
		}

		void TryRead(std::istream& inStream)
		{
			try { Read(inStream); }
			catch (const std::exception& e)
			{
				header.height = header.width = 0;
				image = nullptr;
			}
		}

		operator bool() const { return static_cast<bool>(image); }

		static PfmImage ReadImage(std::istream& inStream)
		{
			PfmImage image;
			image.Read(inStream);
			return image;
		}

		static PfmImage TryReadImage(std::istream& inStream)
		{
			PfmImage image;
			image.TryRead(inStream);
			return image;
		}
	};
}