#include "ehs/io/img/PNG.h"

namespace ehs
{
	Array<Byte, UInt_64> pngSeq = {0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A};

	PNG::PNG()
		: hashId(0)
	{
	}

	PNG::PNG(const Str_8& filePath)
	{
		id = File::ParseName_8(filePath);
		hashId = id.Hash_64();

		File file(filePath, Mode::READ, Disposition::OPEN);

		Serializer<UInt_64> data = file.ReadSerializer_64(Endianness::BE, file.Size());

		file.Release();

		Array<Byte, UInt_64> seq = data.ReadArray<Byte, UInt_64>(8);
		if (seq != pngSeq)
		{
			EHS_LOG_INT("Error", 0, "File at file path, \"" + filePath + "\", is not a valid PNG file.");
			return;
		}

		while (data.GetOffset() < data.Size())
		{
			UInt_32 length = data.Read<UInt_32>();

			Str_8 id = data.ReadStr<Char_8, UInt_64>(4);

			Serializer<UInt_64> chunkData(Endianness::BE, &data[data.GetOffset()], length);
			data.SetOffset(data.GetOffset() + length);

			Byte crc[4];
			UInt_64 crcSize = 4;

			data.ReadArray(crc, &crcSize);

			PNG_Chunk* chunk = GetChunk(id);
			if (chunk)
			{
				chunk->GetData()->SetOffset(chunk->GetData()->Size());
				chunk->GetData()->WriteSer(chunkData);
				chunk->GetData()->SetOffset(0);
			}
			else
				chunks.Push(PNG_Chunk(id, chunkData, crc));
		}
	}

	PNG::PNG(const Str_8& id, Serializer<UInt_64>& data)
		: id(id), hashId(id.Hash_64())
	{
		Array<Byte, UInt_64> seq = data.ReadArray<Byte, UInt_64>(8);
		if (seq != pngSeq)
			return;

		while (data.GetOffset() < data.Size())
		{
			UInt_32 length = data.Read<UInt_32>();

			Str_8 id = data.ReadStr<Char_8, UInt_64>(4);

			Serializer<UInt_64> chunkData(Endianness::BE, &data[data.GetOffset()], length);
			data.SetOffset(data.GetOffset() + length);

			Byte crc[4];
			UInt_64 crcSize = 4;

			data.ReadArray(crc, &crcSize);

			PNG_Chunk* chunk = GetChunk(id);
			if (chunk)
			{
				chunk->GetData()->SetOffset(chunk->GetData()->Size());
				chunk->GetData()->WriteSer(chunkData);
				chunk->GetData()->SetOffset(0);
			}
			else
				chunks.Push(PNG_Chunk(id, chunkData, crc));
		}
	}

	PNG::PNG(const PNG& png)
		: id(png.id), hashId(png.hashId), chunks(png.chunks)
	{
	}

	PNG& PNG::operator=(const PNG& png)
	{
		if (this == &png)
			return *this;

		id = png.id;
		hashId = png.hashId;
		chunks = png.chunks;

		return *this;
	}

	bool PNG::HasChunk(const UInt_64 hashId) const
	{
		for (UInt_64 i = 0; i < chunks.Size(); ++i)
			if (chunks[i].GetHashId() == hashId)
				return true;

		return false;
	}

	bool PNG::HasChunk(const Str_8& id) const
	{
		return HasChunk(id.Hash_64());
	}

	PNG_Chunk* PNG::GetChunk(const UInt_64 hashId)
	{
		for (UInt_64 i = 0; i < chunks.Size(); ++i)
			if (chunks[i].GetHashId() == hashId)
				return &chunks[i];

		return nullptr;
	}

	PNG_Chunk* PNG::GetChunk(const Str_8& id)
	{
		return GetChunk(id.Hash_64());
	}

	bool PNG::IsPNG(Serializer<UInt_32>& data)
	{
		UInt_32 oldOffset = data.GetOffset();
		data.SetOffset(0);

		Array<Byte, UInt_32> seq = data.ReadArray<Byte, UInt_32>(8);
		if (seq != pngSeq)
			return false;

		data.SetOffset(oldOffset);

		return true;
	}

	void PNG::FilterNone(const Byte* const in, Byte* const out, const UInt_8 bitDepth, const UInt_8 channels, const UInt_32 scanline)
	{
		UInt_8 bytes = bitDepth / 8;

		for (UInt_32 i = 0; i < scanline; i += bytes)
			for (UInt_8 b = 0; b < bytes; ++b)
				out[i + b] = in[i + bytes - 1 - b];
	}

	void PNG::FilterSub(const Byte* const in, Byte* const out, const UInt_8 bitDepth, const UInt_8 channels, const UInt_32 scanline)
	{
		UInt_8 bytes = bitDepth / 8;
		UInt_8 pxSize = bytes * channels;

		for (UInt_32 i = 0; i < scanline; i += bytes)
		{
			for (UInt_8 b = 0; b < bytes; ++b)
			{
				if (i >= pxSize)
					out[i + b] = in[i + bytes - 1 - b] + out[i - pxSize + b];
				else
					out[i + b] = in[i + bytes - 1 - b];
			}
		}
	}

	void PNG::FilterUp(const Byte* const in, Byte* const out, const UInt_8 bitDepth, const UInt_8 channels, const UInt_32 scanline)
	{
		UInt_8 bytes = bitDepth / 8;

		for (UInt_32 i = 0; i < scanline; i += bytes)
			for (UInt_8 b = 0; b < bytes; ++b)
				out[scanline + i + b] = in[i + bytes - 1 - b] + out[i + b];
	}

	void PNG::FilterAverage(const Byte* const in, Byte* const out, const UInt_8 bitDepth, const UInt_8 channels, const UInt_32 scanline)
	{
		UInt_8 bytes = bitDepth / 8;
		UInt_8 pxSize = bytes * channels;

		for (UInt_32 i = 0; i < scanline; i += bytes)
		{
			for (UInt_8 b = 0; b < bytes; ++b)
			{
				Byte prevPx = 0;
				if (i >= pxSize)
					prevPx = out[scanline + i - pxSize + b];

				out[scanline + i + b] = in[i + bytes - 1 - b] + (Byte)Math::Floor(((float)prevPx + (float)out[i + b]) / 2.0f);
			}
		}
	}

	void PNG::FilterPaeth(const Byte* const in, Byte* const out, const UInt_8 bitDepth, const UInt_8 channels, const UInt_32 scanline)
	{
		UInt_8 bytes = bitDepth / 8;
		UInt_8 pxSize = bytes * channels;

		for (UInt_32 i = 0; i < scanline; i += bytes)
		{
			for (UInt_8 b = 0; b < bytes; ++b)
			{
				if (i >= pxSize)
					out[scanline + i + b] = in[i + bytes - 1 - b] + PaethPredictor(out[scanline + i - pxSize + b], out[i + b], out[i - pxSize + b]);
				else
					out[scanline + i + b] = in[i + bytes - 1 - b] + out[i + b];
			}
		}
	}

	Byte PNG::PaethPredictor(const Byte a, const Byte b, const Byte c)
	{
		SInt_16 p = a + b - c;
		SInt_16 pa = Math::Abs(p - a);
		SInt_16 pb = Math::Abs(p - b);
		SInt_16 pc = Math::Abs(p - c);

		if (pa <= pb && pa <= pc)
			return a;
		else if (pb <= pc)
			return b;
		else
			return c;
	}
}