#include "ehs/EHS.h"
#include "ehs/Log.h"
#include "ehs/Version.h"
#include "ehs/GarbageCollector.h"
#include "ehs/io/audio/Audio.h"
#include "ehs/io/img/Img.h"
#include "ehs/io/img/PNG.h"
#include "ehs/io/RIFF.h"

#include <zlib.h>

#if defined(EHS_OS_LINUX)
    #include <sys/mman.h>
#endif

namespace ehs
{
    constexpr Char_32 name_32[] = U"Event Horizon Suite";
    constexpr Char_16 name_16[] = L"Event Horizon Suite";
    constexpr Char_8 name_8[] = "Event Horizon Suite";
    constexpr Char_32 acronym_32[] = U"EHS";
    constexpr Char_16 acronym_16[] = L"EHS";
    constexpr Char_8 acronym_8[] = "EHS";
    constexpr Char_32 versionId_32[] = U"Release";
    constexpr Char_16 versionId_16[] = L"Release";
    constexpr Char_8 versionId_8[] = "Release";

	Str_8 appName;
	Str_8 appVerId;
	Version appVer;

    const Char_32* GetName_32()
    {
        return name_32;
    }

    const Char_16* GetName_16()
    {
        return name_16;
    }

    const Char_8* GetName_8()
    {
	    return name_8;
    }

	Str_8 GetAppName_8()
	{
		return appName;
	}

    const Char_32* GetAcronym_32()
    {
        return acronym_32;
    }

    const Char_16* GetAcronym_16()
    {
        return acronym_16;
    }

    const Char_8* GetAcronym_8()
    {
        return acronym_8;
    }

    const Char_32* GetVersionId_32()
    {
        return versionId_32;
    }

    const Char_16* GetVersionId_16()
    {
        return versionId_16;
    }

    const Char_8* GetVersionId_8()
    {
        return versionId_8;
    }

	Str_8 GetAppVersionId_8()
	{
		return appVerId;
	}

    Version GetVersion()
    {
        return {1, 2, 0};
    }

	Version GetAppVersion()
	{
		return appVer;
	}

	bool DecodeWAV(const ehs::AudioCodec* const codec, ehs::Serializer<ehs::UInt_64>& in, ehs::Audio* out)
	{
		RIFF riff(in);

		if (riff.GetType() != "WAVE")
		{
			EHS_LOG_INT("Error", 0, "Data is not in WAVE format.");
			return false;
		}

		RIFF_Chunk fmt = riff.GetChunk("fmt ");
		if (!fmt.IsValid())
		{
			EHS_LOG_INT("Error", 1, "Wave does not have a format chunk.");
			return false;
		}

		Serializer<> fmtSer = fmt.GetData();

		RIFF_Chunk dChunk = riff.GetChunk("data");
		if (!dChunk.IsValid())
		{
			EHS_LOG_INT("Error", 2, "Wave does not have a data chunk.");
			return false;
		}

		UInt_16 compression = fmtSer.Read<UInt_16>();
		if (compression == 0x2)
		{
			EHS_LOG_INT("Error", 3, "Microsoft ADPCM compression unsupported.");
			return false;
		}
		else if (compression == 0x6)
		{
			EHS_LOG_INT("Error", 4, "ITU G.711 a-law compression unsupported.");
			return false;
		}
		else if (compression == 0x7)
		{
			EHS_LOG_INT("Error", 5, "ITU G.711 µ-law compression unsupported.");
			return false;
		}
		else if (compression == 0x11)
		{
			EHS_LOG_INT("Error", 6, "IMA ADPCM compression unsupported.");
			return false;
		}
		else if (compression == 0x16)
		{
			EHS_LOG_INT("Error", 7, "TU G.723 ADPCM (Yamaha) compression unsupported.");
			return false;
		}
		else if (compression == 0x31)
		{
			EHS_LOG_INT("Error", 8, "GSM 6.10 compression unsupported.");
			return false;
		}
		else if (compression == 0x40)
		{
			EHS_LOG_INT("Error", 9, "ITU G.721 ADPCM compression unsupported.");
			return false;
		}
		else if (compression == 0x50)
		{
			EHS_LOG_INT("Error", 10, "MPEG compression unsupported.");
			return false;
		}
		else if (compression == 0xFFFF)
		{
			EHS_LOG_INT("Error", 11, "Experimental compression unsupported.");
			return false;
		}
		else if (compression != 0x1 && compression != 0x3)
		{
			EHS_LOG_INT("Error", 12, "Wave has unknown compression of " + Str_8::FromNum(compression) + ".");
			return false;
		}

		UInt_16 channels = fmtSer.Read<UInt_16>();
		UInt_32 sampleRate = fmtSer.Read<UInt_32>();
		fmtSer.SetOffset(fmtSer.GetOffset() + 6);
		UInt_8 byteDepth = (UInt_8)(fmtSer.Read<UInt_16>() / 8);

		DataType dataType;
		if (byteDepth == 1)
			dataType = DataType::SINT_8;
		else if (byteDepth == 2)
			dataType = DataType::SINT_16;
		else if (byteDepth == 3)
			dataType = DataType::SINT_24;
		else if (byteDepth == 4 && compression == 0x3)
			dataType = DataType::FLOAT;
		else if (byteDepth == 4)
			dataType = DataType::SINT_32;
		else if (byteDepth == 8)
			dataType = DataType::SINT_64;
		else
			return false;

		UInt_64 size = dChunk.GetData().Size();
		UInt_64 frames = size / byteDepth / channels;

		*out = std::move(Audio(out->GetId(), sampleRate, dataType, channels, frames));

		Serializer<> dataSer = dChunk.GetData();

		for (UInt_32 i = 0; i < dataSer.Size(); i += byteDepth)
		{
			if (byteDepth == 1)
			{
				*(SInt_8*)&(*out)[i] = dataSer.Read<SInt_8>();
				if ((*out)[i] > *(SInt_8*)out->GetPeak())
					out->SetPeak(sizeof(SInt_8), &(*out)[i]);
			}
			else if (byteDepth == 2)
			{
				*(SInt_16*)&(*out)[i] = dataSer.Read<SInt_16>();
				if (*(SInt_16*)&(*out)[i] > *(SInt_16*)out->GetPeak())
					out->SetPeak(sizeof(SInt_16), &(*out)[i]);
			}
			else if (byteDepth == 3)
			{
				*(SInt_16*)&(*out)[i + 1] = dataSer.Read<SInt_16>();
				(*out)[i] = dataSer.Read<Byte>();

				SInt_32 signal = 0;
				signal |= (*out)[i];
				signal |= (*out)[i + 1] << 8;
				signal |= (*out)[i + 2] << 16;

				SInt_32 peak = 0;
				peak |= out->GetPeak()[0];
				peak |= out->GetPeak()[1] << 8;
				peak |= out->GetPeak()[2] << 16;

				if (signal > peak)
					out->SetPeak(3, &(*out)[i]);
			}
			else if (byteDepth == 4 && compression == 0x3)
			{
				*(float*)&(*out)[i] = dataSer.Read<float>();
				if (*(float*)&(*out)[i] > *(float*)out->GetPeak())
					out->SetPeak(sizeof(float), &(*out)[i]);
			}
			else if (byteDepth == 4)
			{
				*(SInt_32*)&(*out)[i] = dataSer.Read<SInt_32>();
				if (*(SInt_32*)&(*out)[i] > *(SInt_32*)out->GetPeak())
					out->SetPeak(sizeof(SInt_32), &(*out)[i]);
			}
			else if (byteDepth == 8)
			{
				*(SInt_64*)&(*out)[i] = dataSer.Read<SInt_64>();
				if (*(SInt_64*)&(*out)[i] > *(SInt_64*)out->GetPeak())
					out->SetPeak(sizeof(SInt_64), &(*out)[i]);
			}
		}

		return true;
	}

	bool EncodeEHA(const ehs::AudioCodec* const codec, ehs::Serializer<ehs::UInt_64>& out, const ehs::Audio* in)
	{
		Serializer<UInt_64> result(codec->GetEndianness());
		result.WriteVersion({1, 0, 0});
		result.Write(in->GetSampleRate());
		result.Write(in->GetDataType());
		result.Write(in->GetByteDepth());
		result.Write(in->GetChannels());
		result.Write(in->GetFrameCount());

		UInt_64 size = in->GetSize();
		UInt_8 byteDepth = in->GetByteDepth();

		result.Resize(result.Size() + size + byteDepth);
		Util::Copy(&result[result.GetOffset()], &in[0], size);
		result.SetOffset(result.GetOffset() + size);

		Util::Copy(&result[result.GetOffset()], in->GetPeak(), byteDepth);

		return true;
	}

	bool DecodeEHA(const ehs::AudioCodec* const codec, ehs::Serializer<ehs::UInt_64>& in, ehs::Audio* out)
	{
		Version version = in.ReadVersion();
		if (version != Version(1, 0, 0))
		{
			EHS_LOG_INT("Error", 0, "Incompatible audio file version.");
			return false;
		}

		UInt_64 sampleRate = in.Read<UInt_64>();
		DataType dataType = in.Read<DataType>();
		UInt_8 byteDepth = in.Read<UInt_8>();
		UInt_8 channels = in.Read<UInt_8>();
		UInt_64 frames = in.Read<UInt_64>();

		*out = Audio(out->GetId(), sampleRate, dataType, channels, frames);

		UInt_64 size = out->GetSize();
		Util::Copy(&(*out)[0], &in[in.GetOffset()], size);
		in.SetOffset(in.GetOffset() + size);

		out->SetPeak(byteDepth, &in[in.GetOffset()]);

		return true;
	}

	bool DecodePNG(const ehs::ImgCodec* const codec, ehs::Serializer<ehs::UInt_64>& in, ehs::Img* out)
	{
		PNG png(out->GetId(), in);

		PNG_Chunk* ihdr = png.GetChunk("IHDR");
		Serializer<UInt_64>* ihdrData = ihdr->GetData();

		UInt_32 width = ihdrData->Read<UInt_32>();
		UInt_32 height = ihdrData->Read<UInt_32>();
		UInt_8 bitDepth = ihdrData->Read<UInt_8>();

		UInt_8 colorType = ihdrData->Read<UInt_8>();
		if (colorType == 3)
		{
			EHS_LOG_INT("Error", 1, "Color type of " + Str_8::FromNum(colorType) + " is unsupported.");
			return false;
		}

		UInt_8 channels = 1;
		if (colorType == 2)
			channels = 3;
		else if (colorType == 4)
			channels = 2;
		else if (colorType == 6)
			channels = 4;

		*out = Img(out->GetId(), bitDepth, channels, {width, height});

		UInt_8 compression = ihdrData->Read<UInt_8>();
		if (compression)
		{
			EHS_LOG_INT("Error", 2, "Compression method of " + Str_8::FromNum(compression) + " is unsupported.");
			return false;
		}

		UInt_8 filter = ihdrData->Read<UInt_8>();
		if (filter)
		{
			EHS_LOG_INT("Error", 3, "Filter method of " + Str_8::FromNum(filter) + " is unsupported.");
			return false;
		}

		UInt_8 interlaced = ihdrData->Read<UInt_8>();
		if (interlaced)
		{
			EHS_LOG_INT("Error", 4, "Interlacing method of " + Str_8::FromNum(interlaced) + " is unsupported.");
			return false;
		}

		UInt_32 scanline = width * (bitDepth / 8) * channels;
		UInt_32 scanLineF = scanline + 1;
		UInt_32 bufferSize = scanline * height + height;

		Byte* buffer = new Byte[bufferSize];

		PNG_Chunk* idat = png.GetChunk("IDAT");
		Serializer<UInt_64>* idatData = idat->GetData();

		z_stream strm = {};
		strm.zalloc = Z_NULL;
		strm.zfree = Z_NULL;
		strm.opaque = Z_NULL;
		strm.avail_in = idatData->Size();
		strm.next_in = *idatData;
		strm.avail_out = bufferSize;
		strm.next_out = buffer;

		int code = inflateInit(&strm);
		if (code != Z_OK)
		{
			EHS_LOG_INT("Error", 5, "Failed to initialize zlib inflate with error #" + Str_8::FromNum(code) + ".");
			delete[] buffer;
			return false;
		}

		do
		{
			code = inflate(&strm, Z_NO_FLUSH);
			if (code != Z_STREAM_END && code != Z_OK)
			{
				EHS_LOG_INT("Error", 6, "Failed to zlib inflate with error #" + Str_8::FromNum(code) + ".");
				delete[] buffer;
				return false;
			}
		} while (strm.avail_out);

		code = inflateEnd(&strm);
		if (code != Z_OK)
		{
			EHS_LOG_INT("Error", 7, "Failed to uninitialize zlib inflate with error #" + Str_8::FromNum(code) + ".");
			delete[] buffer;
			return false;
		}

		for (UInt_32 i = 0, o = 0; i < bufferSize; i += scanLineF, o += scanline)
		{
			UInt_8 fCode = buffer[i];
			if (fCode == 0)
				PNG::FilterNone(&buffer[i + 1], &(*out)[o], bitDepth, channels, scanline);
			else if (fCode == 1)
				PNG::FilterSub(&buffer[i + 1], &(*out)[o], bitDepth, channels, scanline);
			else if (fCode == 2)
				PNG::FilterUp(&buffer[i + 1], &(*out)[o - scanline], bitDepth, channels, scanline);
			else if (fCode == 3)
				PNG::FilterAverage(&buffer[i + 1], &(*out)[o - scanline], bitDepth, channels, scanline);
			else if (fCode == 4)
				PNG::FilterPaeth(&buffer[i + 1], &(*out)[o - scanline], bitDepth, channels, scanline);
		}

		delete[] buffer;

		return true;
	}

	bool EncodeQOI(const ehs::ImgCodec* const codec, ehs::Serializer<ehs::UInt_64>& out, const ehs::Img* in)
	{
		UInt_8 channels = in->GetChannels();
		Vec2_u64 resolution = in->GetResolution();

		UInt_32 px_len = resolution.x * resolution.y * channels;
		UInt_32 px_end = px_len - channels;

		Byte index[256];
		for (UInt_64 i = 0; i < 64; ++i)
			*(UInt_32*)&index[i * 4] = 0;

		Byte prevPixel[4] = {0, 0, 0, 255};
		Byte pixel[4] = {0, 0, 0, 255};

		Serializer<UInt_32> result(Endianness::BE, resolution.x * resolution.y * (channels + 1) + 22);

		result.Write('q');
		result.Write('o');
		result.Write('i');
		result.Write('f');
		result.Write<UInt_32>(resolution.x);
		result.Write<UInt_32>(resolution.y);
		result.Write(in->GetChannels());
		result.Write(1);

		for (UInt_32 px_pos = 0, run = 0; px_pos < px_len; px_pos += channels)
		{
			if (channels == 4)
			{
				*(UInt_32*)pixel = *(UInt_32*)&(*in)[px_pos];
			}
			else
			{
				pixel[0] = (*in)[px_pos];
				pixel[1] = (*in)[px_pos + 1];
				pixel[2] = (*in)[px_pos + 2];
			}

			if (*(UInt_32*)pixel == *(UInt_32*)prevPixel)
			{
				run++;
				if (run == 62 || px_pos == px_end)
				{
					result.Write<UInt_8>(0xc0 | (run - 1));
					run = 0;
				}
			}
			else
			{
				if (run > 0)
				{
					result.Write<UInt_8>(0xc0 | (run - 1));
					run = 0;
				}

				UInt_32 index_pos = (prevPixel[0] * 3 + prevPixel[1] * 5 + prevPixel[2] * 7 + prevPixel[3] * 11) % 64 * channels;

				if (*(UInt_32*)&index[index_pos] == *(UInt_32*)pixel)
				{
					result.Write<UInt_8>(0x00 | (index_pos / channels));
				}
				else
				{
					*(UInt_32*)&index[index_pos] = *(UInt_32*)pixel;

					if (pixel[3] == prevPixel[3])
					{
						SInt_8 vr = pixel[0] - prevPixel[0];
						SInt_8 vg = pixel[1] - prevPixel[1];
						SInt_8 vb = pixel[2] - prevPixel[2];

						SInt_8 vg_r = vr - vg;
						SInt_8 vg_b = vb - vg;

						if (
							vr > -3 && vr < 2 &&
							vg > -3 && vg < 2 &&
							vb > -3 && vb < 2
							)
						{
							result.Write<UInt_8>(0x40 | (vr + 2) << 4 | (vg + 2) << 2 | (vb + 2));
						}
						else if (
							vg_r >  -9 && vg_r <  8 &&
							vg   > -33 && vg   < 32 &&
							vg_b >  -9 && vg_b <  8
							)
						{
							result.Write<UInt_8>(0x80 | (vg   + 32));
							result.Write<UInt_8>((vg_r + 8) << 4 | (vg_b +  8));
						}
						else
						{
							result.Write<UInt_8>(0xfe);
							result.Write(pixel[0]);
							result.Write(pixel[1]);
							result.Write(pixel[2]);
						}
					}
					else
					{
						result.Write<UInt_8>(0xff);
						result.SetEndianness(CPU::GetEndianness());
						result.Write(*(UInt_32*)pixel);
						result.SetEndianness(Endianness::BE);
					}
				}
			}

			*(UInt_32*)prevPixel = *(UInt_32*)pixel;
		}

		result.Write(0x100000000000000);

		return true;
	}

	bool DecodeQOI(const ehs::ImgCodec* const codec, ehs::Serializer<ehs::UInt_64>& in, ehs::Img* out)
	{
		Str_8 imgType = in.ReadStr<Char_8, UInt_64>(4);
		if (imgType != "qoif")
		{
			EHS_LOG_INT("Error", 0, "Given data is not in the qoif format.");
			return false;
		}

		UInt_64 width = in.Read<UInt_32>();
		UInt_64 height = in.Read<UInt_32>();

		UInt_8 channels = in.Read<UInt_8>();
		channels = 4;
		UInt_8 space = in.Read<UInt_8>();
		UInt_8 bitDepth = 8;

		UInt_64 size = width * channels * height;

		*out = Img(out->GetId(), bitDepth, channels, {width, height});

		Byte prevPixel[4] = {0, 0, 0, 255};

		Byte index[256];
		for (UInt_64 i = 0; i < 64; ++i)
			*(UInt_32*)&index[i * 4] = 0;

		UInt_32 chunksLen = in.Size() - 8;
		for (UInt_32 pos = 0, run = 0; pos < size; pos += channels)
		{
			if (run > 0)
				--run;
			else if (in.GetOffset() < chunksLen)
			{
				UInt_32 chunkType = (UInt_32)in.Read<UInt_8>();
				if (chunkType == 0xfe) // RGB
				{
					prevPixel[0] = in.Read<UInt_8>(); // R-value
					prevPixel[1] = in.Read<UInt_8>(); // G-value
					prevPixel[2] = in.Read<UInt_8>(); // B-value
				}
				else if (chunkType == 0xff) // RGBA
				{
					*(UInt_32*)prevPixel = in.Read<UInt_32>();
				}
				else if ((chunkType & 0xc0) == 0x00) // Index
				{
					*(UInt_32*)prevPixel = *(UInt_32*)&index[chunkType * channels];
				}
				else if ((chunkType & 0xc0) == 0x40) // Diff
				{
					prevPixel[0] += ((chunkType >> 4) & 0x03) - 2; // R-value
					prevPixel[1] += ((chunkType >> 2) & 0x03) - 2; // G-value
					prevPixel[2] += (chunkType & 0x03) - 2; // B-value
				}
				else if ((chunkType & 0xc0) == 0x80) // Luma
				{
					UInt_32 mod = (UInt_32)in.Read<UInt_8>();
					UInt_32 vg = (chunkType & 0x3f) - 32;
					prevPixel[0] += vg - 8 + ((mod >> 4) & 0x0f); // R-value
					prevPixel[1] += vg; // G-value
					prevPixel[2] += vg - 8 + (mod & 0x0f); // B-value
				}
				else if ((chunkType & 0xc0) == 0xc0) // Run
					run = (chunkType & 0x3f);

				*(UInt_32*)&index[(prevPixel[0] * 3 + prevPixel[1] * 5 + prevPixel[2] * 7 + prevPixel[3] * 11) % 64 * channels] = *(UInt_32*)prevPixel;
			}

			if (channels == 4)
			{
				*((UInt_32*)&(*out)[pos]) = *(UInt_32*)prevPixel;
			}
			else
			{
				(*out)[pos] = prevPixel[0];
				(*out)[pos + 1] = prevPixel[1];
				(*out)[pos + 2] = prevPixel[2];
			}
		}

		return true;
	}
}

int main()
{
	ehs::Audio::AddCodec({
		"Waveform Audio",
		"wav",
		ehs::Endianness::LE,
		nullptr,
		ehs::DecodeWAV
	});

	ehs::Audio::AddCodec({
		"Event Horizon Audio",
		"eha",
		ehs::Endianness::LE,
		ehs::EncodeEHA,
		ehs::DecodeEHA
	});

	ehs::Img::AddCodec({
		"Portable Network Graphic",
		"png",
		ehs::Endianness::BE,
		nullptr,
		ehs::DecodePNG
	});

	ehs::Img::AddCodec({
		"Quite OK Image",
		"qoi",
		ehs::Endianness::BE,
		ehs::EncodeQOI,
		ehs::DecodeQOI
	});

	ehs::GarbageCollector::Start();

	const ehs::SInt_32 code = Main(&ehs::appName, &ehs::appVerId, &ehs::appVer);
	if (code)
		EHS_LOG("Warning", 0, "Executable exited with code #" + ehs::Str_8::FromNum(code) + ".");

	ehs::GarbageCollector::Stop();

	return code;
}