#include "SpriteAssetLoader.h"
#include <cstdint>
#include <fstream>
#include <limits>
#include <vector>

namespace
{
	static const char SpriteMagic[4] = { 'S', 'P', 'R', 'T' };
	static const size_t SpriteHeaderSize = 20;
	static const uint32_t SpriteVersion = 1;
	static const uint32_t SpriteFormatRgba8888 = 1;
	static const uint64_t MaxSpritePixels = 8192ull * 8192ull;

	static bool ReadExact(std::ifstream& file, char* data, size_t size)
	{
		file.read(data, static_cast<std::streamsize>(size));
		return static_cast<size_t>(file.gcount()) == size;
	}

	static bool ReadU32LE(const std::vector<uint8_t>& bytes, size_t offset, uint32_t& value)
	{
		if (offset + 4 > bytes.size())
		{
			return false;
		}

		value =
			static_cast<uint32_t>(bytes[offset]) |
			(static_cast<uint32_t>(bytes[offset + 1]) << 8) |
			(static_cast<uint32_t>(bytes[offset + 2]) << 16) |
			(static_cast<uint32_t>(bytes[offset + 3]) << 24);
		return true;
	}

	static void WriteU32LE(std::ofstream& file, uint32_t value)
	{
		const char bytes[4] = {
			static_cast<char>(value & 0xFF),
			static_cast<char>((value >> 8) & 0xFF),
			static_cast<char>((value >> 16) & 0xFF),
			static_cast<char>((value >> 24) & 0xFF)
		};
		file.write(bytes, sizeof(bytes));
	}

	static bool CheckPixelCount(uint32_t width, uint32_t height, size_t& pixelCount)
	{
		if (width == 0 || height == 0)
		{
			return false;
		}

		const uint64_t total = static_cast<uint64_t>(width) * static_cast<uint64_t>(height);
		if (total > MaxSpritePixels ||
			total > static_cast<uint64_t>(std::numeric_limits<size_t>::max() / sizeof(uint32_t)))
		{
			return false;
		}

		pixelCount = static_cast<size_t>(total);
		return true;
	}

	static bool ReadPixels(std::ifstream& file, std::vector<uint32_t>& pixels)
	{
		std::vector<uint8_t> bytes(pixels.size() * sizeof(uint32_t), 0);
		if (!ReadExact(file, reinterpret_cast<char*>(bytes.data()), bytes.size()))
		{
			return false;
		}

		for (size_t i = 0; i < pixels.size(); ++i)
		{
			uint32_t value = 0;
			if (!ReadU32LE(bytes, i * sizeof(uint32_t), value))
			{
				return false;
			}
			pixels[i] = value;
		}

		return true;
	}

	static void WritePixels(std::ofstream& file, const RenderData::Image& image)
	{
		for (size_t i = 0; i < image.total_pixels(); ++i)
		{
			WriteU32LE(file, image.data()[i]);
		}
	}
}

namespace Asset
{
	bool SpriteAssetLoader::Load(const std::string& path, RenderData::Image& image)
	{
		std::ifstream file(path.c_str(), std::ios::binary);
		if (!file.good())
		{
			return false;
		}

		std::vector<uint8_t> header(SpriteHeaderSize, 0);
		if (!ReadExact(file, reinterpret_cast<char*>(header.data()), header.size()))
		{
			return false;
		}

		if (header[0] != SpriteMagic[0] ||
			header[1] != SpriteMagic[1] ||
			header[2] != SpriteMagic[2] ||
			header[3] != SpriteMagic[3])
		{
			return false;
		}

		uint32_t version = 0;
		uint32_t width = 0;
		uint32_t height = 0;
		uint32_t format = 0;
		if (!ReadU32LE(header, 4, version) ||
			!ReadU32LE(header, 8, width) ||
			!ReadU32LE(header, 12, height) ||
			!ReadU32LE(header, 16, format))
		{
			return false;
		}

		size_t pixelCount = 0;
		if (version != SpriteVersion ||
			format != SpriteFormatRgba8888 ||
			!CheckPixelCount(width, height, pixelCount))
		{
			return false;
		}

		std::vector<uint32_t> pixels(pixelCount, 0);
		if (!ReadPixels(file, pixels))
		{
			return false;
		}

		char trailingByte = 0;
		file.read(&trailingByte, 1);
		if (file.gcount() != 0)
		{
			return false;
		}

		image = RenderData::Image(
			static_cast<int32_t>(width),
			static_cast<int32_t>(height),
			pixels);
		return image.is_valid();
	}

	bool SpriteAssetLoader::Save(const std::string& path, const RenderData::Image& image)
	{
		if (!image.is_valid())
		{
			return false;
		}

		std::ofstream file(path.c_str(), std::ios::binary);
		if (!file.good())
		{
			return false;
		}

		file.write(SpriteMagic, sizeof(SpriteMagic));
		WriteU32LE(file, SpriteVersion);
		WriteU32LE(file, static_cast<uint32_t>(image.get_width()));
		WriteU32LE(file, static_cast<uint32_t>(image.get_height()));
		WriteU32LE(file, SpriteFormatRgba8888);
		WritePixels(file, image);
		return file.good();
	}
}