/*
 * 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 "ImageLoaderPnm.hpp"

#include <sstream>
#include <fstream>

namespace
{
	struct PnmHeader
	{
		char magic = 0; // 1 for P1, 2 for P2 etc
		int width = 0;
		int height = 0;
		int maxVal = 255; // Only used for formats P2, P3, P5, P6
	};

	void MaybeSkipComments(std::istream& stream)
	{
		std::string line;
		while (std::getline(stream, line))
		{
			line.erase(0, line.find_first_not_of(" \t"));

			if (line.empty() || line[0] != '#')
			{
				stream.seekg(-static_cast<std::streamoff>(line.length()) - 1, std::ios::cur);
				break;
			}
		}
	}

	PnmHeader ParseHeader(std::istream& stream)
	{
		PnmHeader header;

		MaybeSkipComments(stream);
		char pMagic;
		stream >> pMagic;
		if (pMagic != 'P')
		{
			throw std::runtime_error("Unsupported magic");
		}

		stream >> header.magic;
		header.magic -= '0';
		if (header.magic < 1 || header.magic > 6)
		{
			throw std::runtime_error("Unsupported magic number: P" + (header.magic + '0'));
		}

		MaybeSkipComments(stream);
		stream >> header.width >> header.height;
		if (header.magic != 1 && header.magic != 4)
		{
			MaybeSkipComments(stream);
			stream >> header.maxVal;
		}

		stream.ignore(std::numeric_limits<std::streamsize>::max(), '\n');
		return header;
	}

	OpenVulkano::Array<uint8_t> ReadBinaryData(std::istream& stream, const PnmHeader& header)
	{
		size_t dataSize = header.width * header.height * ((header.magic == 6) ? 3 : 1);
		OpenVulkano::Array<uint8_t> data(dataSize);

		if (header.magic == 4)
		{
			size_t rowBytes = (header.width + 7) / 8;
			OpenVulkano::Array<uint8_t> rowBuffer(rowBytes);
			
			for (int y = header.height - 1; y >= 0; --y)
			{
				stream.read(reinterpret_cast<char*>(rowBuffer.Data()), rowBytes);
				size_t rowStart = y * header.width;
				for (int x = 0; x < header.width; ++x)
				{
					size_t byteIndex = x / 8;
					size_t bitIndex = 7 - (x % 8);
					data[rowStart + x] = (rowBuffer[byteIndex] >> bitIndex) & 1 ? 255 : 0;
				}
			}
		}
		else
		{
			size_t bytesPerRow = header.width * ((header.magic == 6) ? 3 : 1);
			OpenVulkano::Array<uint8_t> rowBuffer(bytesPerRow);
			
			for (int y = header.height - 1; y >= 0; --y)
			{
				stream.read(reinterpret_cast<char*>(rowBuffer.Data()), bytesPerRow);
				size_t rowStart = y * bytesPerRow;
				std::memcpy(data.Data() + rowStart, rowBuffer.Data(), bytesPerRow);
			}
		}

		return data;
	}

	OpenVulkano::Array<uint8_t> ReadTextData(std::istream& stream, const PnmHeader& header)
	{
		const size_t numChannels = (header.magic == 3) ? 3 : 1;
		OpenVulkano::Array<uint8_t> data(header.width * header.height * numChannels);
		const int PRECOMPUTED_RATIO = 255 / header.maxVal;

		for (int y = header.height - 1; y >= 0; --y)
		{
			for (size_t x = 0; x < header.width * numChannels; ++x)
			{
				int value;
				stream >> value;

				if (header.magic == 1)
				{
					data[y * header.width * numChannels + x] = static_cast<uint8_t>(value * 255);
				}
				else
				{
					data[y * header.width * numChannels + x] = static_cast<uint8_t>(value * PRECOMPUTED_RATIO);
				}
			}
		}

		return data;
	}
}

namespace OpenVulkano::Image
{
	std::unique_ptr<Image> ImageLoaderPnm::loadFromFile(const std::string& filePath)
	{
		std::ifstream file(filePath, std::ios::binary);
		if (!file.is_open())
		{
			throw std::runtime_error("Failed to open file: " + filePath);
		}

		std::vector<uint8_t> buffer((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
		return loadFromMemory(buffer);
	}

	std::unique_ptr<Image> ImageLoaderPnm::loadFromMemory(const std::vector<uint8_t>& buffer)
	{
		std::istringstream stream(std::string(buffer.begin(), buffer.end()));
		PnmHeader header = ParseHeader(stream);

		std::unique_ptr<Image> result = std::make_unique<Image>();
		result->resolution.x = header.width;
		result->resolution.y = header.height;
		result->resolution.z = 1;

		if (header.magic <= 3)
		{
			result->data = ReadTextData(stream, header);
		}
		else if (header.magic <= 6)
		{
			result->data = ReadBinaryData(stream, header);
		}
		else
		{
			throw std::runtime_error("Unsupported format: P" + (header.magic + '0'));
		}

		if (header.magic == 1 || header.magic == 2 || header.magic == 4 || header.magic == 5)
		{
			result->dataFormat = DataFormat::Format::R8_UINT;
		}
		else
		{
			result->dataFormat = DataFormat::Format::R8G8B8_UINT;
		}

		return result;
	}

	Math::Vector2i ImageLoaderPnm::GetImageDimensions(const std::string& filename)
	{
		std::ifstream file(filename, std::ios::binary);
		if (!file.is_open())
		{
			throw std::runtime_error("Failed to open file: " + filename);
		}

		PnmHeader header = ParseHeader(file);
		return Math::Vector2i(header.width, header.height);
	}
};