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

#include "Scene/DataFormat.hpp"

#include <algorithm>
#include <cmath>
#include <fstream>
#include <sstream>
#include <stdexcept>

namespace
{
	struct PfmImageHeader
	{
		int width;
		int height;
		bool isColor;
		bool isLittleEndian;
	};

	PfmImageHeader ParseHeader(std::istringstream& stream)
	{
		PfmImageHeader header;
		std::string magic;
		stream >> magic;
		header.isColor = (magic == "PF");

		float endianess;
		stream >> header.width >> header.height >> endianess;
		stream.ignore();

		header.isLittleEndian = endianess < 0;

		return header;
	}
}

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

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

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

		size_t numChannels = header.isColor ? 3 : 1;
		Array<float> data(header.width * header.height * numChannels);
		size_t rowSize = header.width * numChannels * sizeof(float);

		for (int y = header.height - 1; y >= 0; --y)
		{
			stream.read(reinterpret_cast<char*>(&data[y * header.width * numChannels]), rowSize);
			if (!stream)
			{
				throw std::runtime_error("Error reading PFM image data");
			}
		}

		if (!header.isLittleEndian)
		{
			for (float& value : data)
			{
				uint8_t* ptr = reinterpret_cast<uint8_t*>(&value);
				std::reverse(ptr, ptr + sizeof(float));
			}
		}

		auto image = std::make_unique<Image>();
		image->resolution = { static_cast<uint32_t>(header.width), static_cast<uint32_t>(header.height), 1 };
		image->dataFormat = header.isColor ? DataFormat::Format::R32G32B32_SFLOAT : DataFormat::Format::R32_SFLOAT;
		image->data = std::move(*reinterpret_cast<Array<uint8_t>*>(&data));
		return image;
	}

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

		std::vector<uint8_t> buffer((std::istreambuf_iterator<char>(file)), std::istreambuf_iterator<char>());
		std::istringstream stream(std::string(buffer.begin(), buffer.end()), std::ios::binary);
		PfmImageHeader header = ParseHeader(stream);
		return { header.width, header.height };
	}
}