/*
 * 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 <string>
#include <sstream>
#include <memory>
#include <cmath>

namespace OpenVulkano
{
	//TODO handle comments
	/**
	 * Implements Portable Anymap file headers
	 * See https://de.wikipedia.org/wiki/Portable_Anymap
	 */
	class PnmHeader
	{
		static constexpr std::string_view FILE_TYPES[] = { "PBM", "PGM", "PPM" };
		static constexpr std::string_view FILE_EXT[] = { ".pbm", ".pgm", ".ppm" };
	public:
		enum class ColorMode : uint8_t { BW = 1, GRAY, COLOR };

		uint32_t width, height, maxValue;
		ColorMode color;
		bool ascii;

		constexpr PnmHeader() : width(0), height(0), maxValue(0), color(ColorMode::GRAY), ascii(false) {}

		constexpr PnmHeader(uint32_t width, uint32_t height, bool color, uint32_t maxValue = 255, bool ascii = false) :
				width(width), height(height), maxValue(maxValue), color(color ? ColorMode::COLOR : ColorMode::GRAY), ascii(ascii)
		{}

		constexpr PnmHeader(uint32_t width, uint32_t height, ColorMode color = ColorMode::COLOR, uint32_t maxValue = 255, bool ascii = false) :
				width(width), height(height), maxValue(maxValue), color(color), ascii(ascii)
		{}

		[[nodiscard]] constexpr std::string_view GetFileType() const
		{
			return FILE_TYPES[static_cast<uint8_t>(color) - 1];
		}

		[[nodiscard]] constexpr std::string_view GetFileExtension() const
		{
			return FILE_EXT[static_cast<uint8_t>(color) - 1];
		}

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

		[[nodiscard]] constexpr uint8_t GetMagicNumberValue() const
		{
			uint8_t value = static_cast<uint8_t>(color);
			if (!ascii) value += 3;
			return value;
		}

		[[nodiscard]] constexpr char GetMagicNumberChar() const
		{
			return '0' + static_cast<char>(GetMagicNumberValue());
		}

		friend std::ostream& operator<< (std::ostream& outStream, const PnmHeader& pnmHeader)
		{
			outStream << 'P';
			outStream << pnmHeader.GetMagicNumberChar();
			outStream << '\n';
			outStream << pnmHeader.width << ' ' << pnmHeader.height << '\n';
			outStream << pnmHeader.maxValue << '\n';
			return outStream;
		}

		friend std::istream& operator>> (std::istream& inStream, PnmHeader& pnmHeader)
		{
			std::string marker;
			inStream >> marker;
			if (marker.length() != 2 || marker[0] != 'P') throw std::runtime_error("Malformed PNM header!");
			uint8_t val = marker[1] - '0';
			if (val > 6) throw std::runtime_error("Malformed PNM header!");
			pnmHeader.ascii = val < 4;
			pnmHeader.color = static_cast<ColorMode>(val - 3);

			if (!(inStream >> pnmHeader.width >> pnmHeader.height)) [[unlikely]]
			{
				throw std::runtime_error("Malformed PNM header: Unable to read width and height!");
			}
			if (pnmHeader.width == 0 || pnmHeader.height == 0) [[unlikely]]
			{
				throw std::runtime_error("Malformed PNM header: Width and height must be positive!");
			}

			if (!(inStream >> pnmHeader.maxValue)) [[unlikely]]
			{
				throw std::runtime_error("Malformed PNM header: Unable to read maxValue!");
			}
			if (pnmHeader.maxValue == 0 || pnmHeader.maxValue > 65535) [[unlikely]]
			{
				throw std::runtime_error("Malformed PNM header: Invalid maxValue (must be between 1 and 65535)!");
			}

			inStream.get();
			return inStream;
		}

		[[nodiscard]] constexpr size_t GetImageSize() const
		{
			size_t size = height * width;
			if (color == ColorMode::COLOR)
			{
				size *= 3;
				if (maxValue > 255) size *= 2;
				return size;
			}
			if (color == ColorMode::BW)
			{
				return std::ceil(size / 8.0);
			}
			return size;
		}
	};

	struct PnmImage
	{
		PnmHeader header;
		std::unique_ptr<char[]> image;

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

			// TODO handle ascii mode
			if (header.ascii) throw std::runtime_error("ascii mode not supported!");

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

			size_t size = header.GetImageSize();
			image = std::make_unique<char[]>(size);
			inStream.read(image.get(), size);
		}

		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 PnmImage ReadImage(std::istream& inStream)
		{
			PnmImage image;
			image.Read(inStream);
			return image;
		}
		
		static PnmImage TryReadImage(std::istream& inStream)
		{
			PnmImage image;
			image.TryRead(inStream);
			return image;
		}
	};
}