// Copyright Epic Games, Inc. All Rights Reserved.

#include <zencore/compress.h>

#include <zencore/blake3.h>
#include <zencore/compositebuffer.h>
#include <zencore/crc32.h>
#include <zencore/endian.h>

#include "../3rdparty/Oodle/include/oodle2.h"
#if ZEN_PLATFORM_WINDOWS
#	pragma comment(lib, "oo2core_win64.lib")
#endif

#include <doctest/doctest.h>
#include <lz4.h>
#include <functional>
#include <limits>

namespace zen::detail {

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

static constexpr uint64_t DefaultBlockSize = 256 * 1024;

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

/** Method used to compress the data in a compressed buffer. */
enum class CompressionMethod : uint8_t
{
	/** Header is followed by one uncompressed block. */
	None = 0,
	/** Header is followed by an array of compressed block sizes then the compressed blocks. */
	Oodle = 3,
	/** Header is followed by an array of compressed block sizes then the compressed blocks. */
	LZ4 = 4,
};

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

/** Header used on every compressed buffer. Always stored in big-endian format. */
struct BufferHeader
{
	static constexpr uint32_t ExpectedMagic = 0xb7756362;  // <dot>ucb

	uint32_t		  Magic = ExpectedMagic;  // A magic number to identify a compressed buffer. Always 0xb7756362.
	uint32_t		  Crc32 = 0;			  // A CRC-32 used to check integrity of the buffer. Uses the polynomial 0x04c11db7
	CompressionMethod Method =
		CompressionMethod::None;	   // The method used to compress the buffer. Affects layout of data following the header
	uint8_t	 Compressor			 = 0;  // The method-specific compressor used to compress the buffer.
	uint8_t	 CompressionLevel	 = 0;  // The method-specific compression level used to compress the buffer.
	uint8_t	 BlockSizeExponent	 = 0;  // The power of two size of every uncompressed block except the last. Size is 1 << BlockSizeExponent
	uint32_t BlockCount			 = 0;  // The number of blocks that follow the header
	uint64_t TotalRawSize		 = 0;  // The total size of the uncompressed data
	uint64_t TotalCompressedSize = 0;  // The total size of the compressed data including the header
	BLAKE3	 RawHash;				   // The hash of the uncompressed data

	/** Checks validity of the buffer based on the magic number, method, and CRC-32. */
	static bool IsValid(const CompositeBuffer& CompressedData);
	static bool IsValid(const SharedBuffer& CompressedData) { return IsValid(CompositeBuffer(CompressedData)); }

	/** Read a header from a buffer that is at least sizeof(BufferHeader) without any validation. */
	static BufferHeader Read(const CompositeBuffer& CompressedData)
	{
		BufferHeader Header;
		if (sizeof(BufferHeader) <= CompressedData.GetSize())
		{
			CompressedData.CopyTo(MakeMutableMemoryView(&Header, &Header + 1));
			Header.ByteSwap();
		}
		return Header;
	}

	/**
	 * Write a header to a memory view that is at least sizeof(BufferHeader).
	 *
	 * @param HeaderView   View of the header to write, including any method-specific header data.
	 */
	void Write(MutableMemoryView HeaderView) const
	{
		BufferHeader Header = *this;
		Header.ByteSwap();
		HeaderView.CopyFrom(MakeMemoryView(&Header, &Header + 1));
		Header.ByteSwap();
		Header.Crc32 = CalculateCrc32(HeaderView);
		Header.ByteSwap();
		HeaderView.CopyFrom(MakeMemoryView(&Header, &Header + 1));
	}

	void ByteSwap()
	{
		Magic				= zen::ByteSwap(Magic);
		Crc32				= zen::ByteSwap(Crc32);
		BlockCount			= zen::ByteSwap(BlockCount);
		TotalRawSize		= zen::ByteSwap(TotalRawSize);
		TotalCompressedSize = zen::ByteSwap(TotalCompressedSize);
	}

	/** Calculate the CRC-32 from a view of a header including any method-specific header data. */
	static uint32_t CalculateCrc32(MemoryView HeaderView)
	{
		uint32_t		   Crc32		= 0;
		constexpr uint64_t MethodOffset = offsetof(BufferHeader, Method);
		for (MemoryView View = HeaderView + MethodOffset; const uint64_t ViewSize = View.GetSize();)
		{
			const int32_t Size = static_cast<int32_t>(zen::Min<uint64_t>(ViewSize, /* INT_MAX */ 2147483647u));
			Crc32			   = zen::MemCrc32(View.GetData(), Size, Crc32);
			View += Size;
		}
		return Crc32;
	}
};

static_assert(sizeof(BufferHeader) == 64, "BufferHeader is the wrong size.");

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

class BaseEncoder
{
public:
	virtual CompositeBuffer Compress(const CompositeBuffer& RawData, uint64_t BlockSize = DefaultBlockSize) const = 0;
};

class BaseDecoder
{
public:
	virtual CompositeBuffer Decompress(const BufferHeader& Header, const CompositeBuffer& CompressedData) const						 = 0;
	virtual bool TryDecompressTo(const BufferHeader& Header, const CompositeBuffer& CompressedData, MutableMemoryView RawView) const = 0;
	virtual uint64_t GetHeaderSize(const BufferHeader& Header) const																 = 0;
};

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

class NoneEncoder final : public BaseEncoder
{
public:
	[[nodiscard]] CompositeBuffer Compress(const CompositeBuffer& RawData, uint64_t /* BlockSize */) const final
	{
		BufferHeader Header;
		Header.Method			   = CompressionMethod::None;
		Header.BlockCount		   = 1;
		Header.TotalRawSize		   = RawData.GetSize();
		Header.TotalCompressedSize = Header.TotalRawSize + sizeof(BufferHeader);
		Header.RawHash			   = BLAKE3::HashBuffer(RawData);

		UniqueBuffer HeaderData = UniqueBuffer::Alloc(sizeof(BufferHeader));
		Header.Write(HeaderData);
		return CompositeBuffer(HeaderData.MoveToShared(), RawData.MakeOwned());
	}
};

class NoneDecoder final : public BaseDecoder
{
public:
	[[nodiscard]] CompositeBuffer Decompress(const BufferHeader& Header, const CompositeBuffer& CompressedData) const final
	{
		if (Header.Method == CompressionMethod::None && Header.TotalCompressedSize == CompressedData.GetSize() &&
			Header.TotalCompressedSize == Header.TotalRawSize + sizeof(BufferHeader))
		{
			return CompressedData.Mid(sizeof(BufferHeader), Header.TotalRawSize).MakeOwned();
		}
		return CompositeBuffer();
	}

	[[nodiscard]] bool TryDecompressTo(const BufferHeader&	  Header,
									   const CompositeBuffer& CompressedData,
									   MutableMemoryView	  RawView) const final
	{
		if (Header.Method == CompressionMethod::None && Header.TotalRawSize == RawView.GetSize() &&
			Header.TotalCompressedSize == CompressedData.GetSize() &&
			Header.TotalCompressedSize == Header.TotalRawSize + sizeof(BufferHeader))
		{
			CompressedData.CopyTo(RawView, sizeof(BufferHeader));
			return true;
		}
		return false;
	}

	[[nodiscard]] uint64_t GetHeaderSize(const BufferHeader&) const final { return sizeof(BufferHeader); }
};

//////////////////////////////////////////////////////////////////////////

class BlockEncoder : public BaseEncoder
{
public:
	CompositeBuffer Compress(const CompositeBuffer& RawData, uint64_t BlockSize = DefaultBlockSize) const final;

protected:
	virtual CompressionMethod GetMethod() const															 = 0;
	virtual uint8_t			  GetCompressor() const														 = 0;
	virtual uint8_t			  GetCompressionLevel() const												 = 0;
	virtual uint64_t		  CompressBlockBound(uint64_t RawSize) const								 = 0;
	virtual bool			  CompressBlock(MutableMemoryView& CompressedData, MemoryView RawData) const = 0;

private:
	uint64_t GetCompressedBlocksBound(uint64_t BlockCount, uint64_t BlockSize, uint64_t RawSize) const
	{
		switch (BlockCount)
		{
			case 0:
				return 0;
			case 1:
				return CompressBlockBound(RawSize);
			default:
				return CompressBlockBound(BlockSize) - BlockSize + RawSize;
		}
	}
};

CompositeBuffer
BlockEncoder::Compress(const CompositeBuffer& RawData, const uint64_t BlockSize) const
{
	ZEN_ASSERT(IsPow2(BlockSize) && (BlockSize <= (1u << 31)));

	const uint64_t RawSize = RawData.GetSize();
	BLAKE3Stream   RawHash;

	const uint64_t BlockCount = RoundUp(RawSize, BlockSize) / BlockSize;
	ZEN_ASSERT(BlockCount <= ~uint32_t(0));

	// Allocate the buffer for the header, metadata, and compressed blocks.
	const uint64_t MetaSize			  = BlockCount * sizeof(uint32_t);
	const uint64_t CompressedDataSize = sizeof(BufferHeader) + MetaSize + GetCompressedBlocksBound(BlockCount, BlockSize, RawSize);
	UniqueBuffer   CompressedData	  = UniqueBuffer::Alloc(CompressedDataSize);

	// Compress the raw data in blocks and store the raw data for incompressible blocks.
	std::vector<uint32_t> CompressedBlockSizes;
	CompressedBlockSizes.reserve(BlockCount);
	uint64_t CompressedSize = 0;
	{
		UniqueBuffer	  RawBlockCopy;
		MutableMemoryView CompressedBlocksView = CompressedData.GetMutableView() + sizeof(BufferHeader) + MetaSize;
		for (uint64_t RawOffset = 0; RawOffset < RawSize;)
		{
			const uint64_t	 RawBlockSize = zen::Min(RawSize - RawOffset, BlockSize);
			const MemoryView RawBlock	  = RawData.ViewOrCopyRange(RawOffset, RawBlockSize, RawBlockCopy);
			RawHash.Append(RawBlock);

			MutableMemoryView CompressedBlock = CompressedBlocksView;
			if (!CompressBlock(CompressedBlock, RawBlock))
			{
				return CompositeBuffer();
			}

			uint64_t CompressedBlockSize = CompressedBlock.GetSize();
			if (RawBlockSize <= CompressedBlockSize)
			{
				CompressedBlockSize	 = RawBlockSize;
				CompressedBlocksView = CompressedBlocksView.CopyFrom(RawBlock);
			}
			else
			{
				CompressedBlocksView += CompressedBlockSize;
			}

			CompressedBlockSizes.push_back(static_cast<uint32_t>(CompressedBlockSize));
			CompressedSize += CompressedBlockSize;
			RawOffset += RawBlockSize;
		}
	}

	// Return an uncompressed buffer if the compressed data is larger than the raw data.
	if (RawSize <= MetaSize + CompressedSize)
	{
		CompressedData.Reset();
		return NoneEncoder().Compress(RawData, BlockSize);
	}

	// Write the header and calculate the CRC-32.
	for (uint32_t& Size : CompressedBlockSizes)
	{
		Size = ByteSwap(Size);
	}
	CompressedData.GetMutableView().Mid(sizeof(BufferHeader), MetaSize).CopyFrom(MakeMemoryView(CompressedBlockSizes));

	BufferHeader Header;
	Header.Method			   = GetMethod();
	Header.Compressor		   = GetCompressor();
	Header.CompressionLevel	   = GetCompressionLevel();
	Header.BlockSizeExponent   = static_cast<uint8_t>(zen::FloorLog2_64(BlockSize));
	Header.BlockCount		   = static_cast<uint32_t>(BlockCount);
	Header.TotalRawSize		   = RawSize;
	Header.TotalCompressedSize = sizeof(BufferHeader) + MetaSize + CompressedSize;
	Header.RawHash			   = RawHash.GetHash();
	Header.Write(CompressedData.GetMutableView().Left(sizeof(BufferHeader) + MetaSize));

	const MemoryView CompositeView = CompressedData.GetView().Left(Header.TotalCompressedSize);
	return CompositeBuffer(SharedBuffer::MakeView(CompositeView, CompressedData.MoveToShared()));
}

class BlockDecoder : public BaseDecoder
{
public:
	CompositeBuffer		   Decompress(const BufferHeader& Header, const CompositeBuffer& CompressedData) const final;
	[[nodiscard]] bool	   TryDecompressTo(const BufferHeader&	  Header,
										   const CompositeBuffer& CompressedData,
										   MutableMemoryView	  RawView) const final;
	[[nodiscard]] uint64_t GetHeaderSize(const BufferHeader& Header) const final
	{
		return sizeof(BufferHeader) + sizeof(uint32_t) * uint64_t(Header.BlockCount);
	}

protected:
	virtual bool DecompressBlock(MutableMemoryView RawData, MemoryView CompressedData) const = 0;
};

CompositeBuffer
BlockDecoder::Decompress(const BufferHeader& Header, const CompositeBuffer& CompressedData) const
{
	if (Header.BlockCount == 0 || Header.TotalCompressedSize != CompressedData.GetSize())
	{
		return CompositeBuffer();
	}

	// The raw data cannot reference the compressed data unless it is owned.
	// An empty raw buffer requires an empty segment, which this path creates.
	if (!CompressedData.IsOwned() || Header.TotalRawSize == 0)
	{
		UniqueBuffer Buffer = UniqueBuffer::Alloc(Header.TotalRawSize);
		return TryDecompressTo(Header, CompressedData, Buffer) ? CompositeBuffer(Buffer.MoveToShared()) : CompositeBuffer();
	}

	std::vector<uint32_t> CompressedBlockSizes;
	CompressedBlockSizes.resize(Header.BlockCount);
	CompressedData.CopyTo(MakeMutableMemoryView(CompressedBlockSizes), sizeof(BufferHeader));

	for (uint32_t& Size : CompressedBlockSizes)
	{
		Size = ByteSwap(Size);
	}

	// Allocate the buffer for the raw blocks that were compressed.
	SharedBuffer	  RawData;
	MutableMemoryView RawDataView;
	const uint64_t	  BlockSize = uint64_t(1) << Header.BlockSizeExponent;
	{
		uint64_t RawDataSize	  = 0;
		uint64_t RemainingRawSize = Header.TotalRawSize;
		for (const uint32_t CompressedBlockSize : CompressedBlockSizes)
		{
			const uint64_t RawBlockSize = zen::Min(RemainingRawSize, BlockSize);
			if (CompressedBlockSize < BlockSize)
			{
				RawDataSize += RawBlockSize;
			}
			RemainingRawSize -= RawBlockSize;
		}
		UniqueBuffer RawDataBuffer = UniqueBuffer::Alloc(RawDataSize);
		RawDataView				   = RawDataBuffer;
		RawData					   = RawDataBuffer.MoveToShared();
	}

	// Decompress the compressed data in blocks and reference the uncompressed blocks.
	uint64_t				  PendingCompressedSegmentOffset = sizeof(BufferHeader) + uint64_t(Header.BlockCount) * sizeof(uint32_t);
	uint64_t				  PendingCompressedSegmentSize	 = 0;
	uint64_t				  PendingRawSegmentOffset		 = 0;
	uint64_t				  PendingRawSegmentSize			 = 0;
	std::vector<SharedBuffer> Segments;

	const auto CommitPendingCompressedSegment =
		[&PendingCompressedSegmentOffset, &PendingCompressedSegmentSize, &CompressedData, &Segments] {
			if (PendingCompressedSegmentSize)
			{
				CompressedData.IterateRange(PendingCompressedSegmentOffset,
											PendingCompressedSegmentSize,
											[&Segments](MemoryView View, const SharedBuffer& ViewOuter) {
												Segments.push_back(SharedBuffer::MakeView(View, ViewOuter));
											});
				PendingCompressedSegmentOffset += PendingCompressedSegmentSize;
				PendingCompressedSegmentSize = 0;
			}
		};

	const auto CommitPendingRawSegment = [&PendingRawSegmentOffset, &PendingRawSegmentSize, &RawData, &Segments] {
		if (PendingRawSegmentSize)
		{
			const MemoryView PendingSegment = RawData.GetView().Mid(PendingRawSegmentOffset, PendingRawSegmentSize);
			Segments.push_back(SharedBuffer::MakeView(PendingSegment, RawData));
			PendingRawSegmentOffset += PendingRawSegmentSize;
			PendingRawSegmentSize = 0;
		}
	};

	UniqueBuffer CompressedBlockCopy;
	uint64_t	 RemainingRawSize		 = Header.TotalRawSize;
	uint64_t	 RemainingCompressedSize = CompressedData.GetSize();
	for (const uint32_t CompressedBlockSize : CompressedBlockSizes)
	{
		if (RemainingCompressedSize < CompressedBlockSize)
		{
			return CompositeBuffer();
		}

		const uint64_t RawBlockSize = zen::Min(RemainingRawSize, BlockSize);
		if (RawBlockSize == CompressedBlockSize)
		{
			CommitPendingRawSegment();
			PendingCompressedSegmentSize += RawBlockSize;
		}
		else
		{
			CommitPendingCompressedSegment();
			const MemoryView CompressedBlock =
				CompressedData.ViewOrCopyRange(PendingCompressedSegmentOffset, CompressedBlockSize, CompressedBlockCopy);
			if (!DecompressBlock(RawDataView.Left(RawBlockSize), CompressedBlock))
			{
				return CompositeBuffer();
			}
			PendingCompressedSegmentOffset += CompressedBlockSize;
			PendingRawSegmentSize += RawBlockSize;
			RawDataView += RawBlockSize;
		}

		RemainingCompressedSize -= CompressedBlockSize;
		RemainingRawSize -= RawBlockSize;
	}

	CommitPendingCompressedSegment();
	CommitPendingRawSegment();

	return CompositeBuffer(std::move(Segments));
}

bool
BlockDecoder::TryDecompressTo(const BufferHeader& Header, const CompositeBuffer& CompressedData, MutableMemoryView RawView) const
{
	if (Header.TotalRawSize != RawView.GetSize() || Header.TotalCompressedSize != CompressedData.GetSize())
	{
		return false;
	}

	std::vector<uint32_t> CompressedBlockSizes;
	CompressedBlockSizes.resize(Header.BlockCount);
	CompressedData.CopyTo(MakeMutableMemoryView(CompressedBlockSizes), sizeof(BufferHeader));

	for (uint32_t& Size : CompressedBlockSizes)
	{
		Size = ByteSwap(Size);
	}

	UniqueBuffer   CompressedBlockCopy;
	const uint64_t BlockSize			   = uint64_t(1) << Header.BlockSizeExponent;
	uint64_t	   CompressedOffset		   = sizeof(BufferHeader) + uint64_t(Header.BlockCount) * sizeof(uint32_t);
	uint64_t	   RemainingRawSize		   = Header.TotalRawSize;
	uint64_t	   RemainingCompressedSize = CompressedData.GetSize();
	for (uint32_t CompressedBlockSize : CompressedBlockSizes)
	{
		if (RemainingCompressedSize < CompressedBlockSize)
		{
			return false;
		}

		const uint64_t RawBlockSize = zen::Min(RemainingRawSize, BlockSize);
		if (RawBlockSize == CompressedBlockSize)
		{
			CompressedData.CopyTo(RawView.Left(RawBlockSize), CompressedOffset);
		}
		else
		{
			const MemoryView CompressedBlock = CompressedData.ViewOrCopyRange(CompressedOffset, CompressedBlockSize, CompressedBlockCopy);
			if (!DecompressBlock(RawView.Left(RawBlockSize), CompressedBlock))
			{
				return false;
			}
		}

		RemainingCompressedSize -= CompressedBlockSize;
		RemainingRawSize -= RawBlockSize;
		CompressedOffset += CompressedBlockSize;
		RawView += RawBlockSize;
	}

	return RemainingRawSize == 0;
}

//////////////////////////////////////////////////////////////////////////

struct OodleInit
{
	OodleInit()
	{
		OodleConfigValues Config;
		Oodle_GetConfigValues(&Config);
		// Always read/write Oodle v9 binary data.
		Config.m_OodleLZ_BackwardsCompatible_MajorVersion = 9;
		Oodle_SetConfigValues(&Config);
	}
};

OodleInit InitOodle;

class OodleEncoder final : public BlockEncoder
{
public:
	OodleEncoder(OodleCompressor InCompressor, OodleCompressionLevel InCompressionLevel)
	: Compressor(InCompressor)
	, CompressionLevel(InCompressionLevel)
	{
	}

protected:
	CompressionMethod GetMethod() const final { return CompressionMethod::Oodle; }
	uint8_t			  GetCompressor() const final { return static_cast<uint8_t>(Compressor); }
	uint8_t			  GetCompressionLevel() const final { return static_cast<uint8_t>(CompressionLevel); }

	uint64_t CompressBlockBound(uint64_t RawSize) const final
	{
		return static_cast<uint64_t>(OodleLZ_GetCompressedBufferSizeNeeded(OodleLZ_Compressor_Kraken, static_cast<OO_SINTa>(RawSize)));
	}

	bool CompressBlock(MutableMemoryView& CompressedData, MemoryView RawData) const final
	{
		const OodleLZ_Compressor	   LZCompressor		  = GetOodleLZCompressor(Compressor);
		const OodleLZ_CompressionLevel LZCompressionLevel = GetOodleLZCompressionLevel(CompressionLevel);
		if (LZCompressor == OodleLZ_Compressor_Invalid || LZCompressionLevel == OodleLZ_CompressionLevel_Invalid ||
			LZCompressionLevel == OodleLZ_CompressionLevel_None)
		{
			return false;
		}

		const OO_SINTa RawSize = static_cast<OO_SINTa>(RawData.GetSize());
		if (static_cast<OO_SINTa>(CompressedData.GetSize()) < OodleLZ_GetCompressedBufferSizeNeeded(LZCompressor, RawSize))
		{
			return false;
		}

		const OO_SINTa Size = OodleLZ_Compress(LZCompressor, RawData.GetData(), RawSize, CompressedData.GetData(), LZCompressionLevel);
		CompressedData.LeftInline(static_cast<uint64_t>(Size));
		return Size > 0;
	}

	static OodleLZ_Compressor GetOodleLZCompressor(OodleCompressor Compressor)
	{
		switch (Compressor)
		{
			case OodleCompressor::Selkie:
				return OodleLZ_Compressor_Selkie;
			case OodleCompressor::Mermaid:
				return OodleLZ_Compressor_Mermaid;
			case OodleCompressor::Kraken:
				return OodleLZ_Compressor_Kraken;
			case OodleCompressor::Leviathan:
				return OodleLZ_Compressor_Leviathan;
			case OodleCompressor::NotSet:
			default:
				return OodleLZ_Compressor_Invalid;
		}
	}

	static OodleLZ_CompressionLevel GetOodleLZCompressionLevel(OodleCompressionLevel Level)
	{
		const int IntLevel = (int)Level;
		if (IntLevel < (int)OodleLZ_CompressionLevel_Min || IntLevel > (int)OodleLZ_CompressionLevel_Max)
		{
			return OodleLZ_CompressionLevel_Invalid;
		}
		return OodleLZ_CompressionLevel(IntLevel);
	}

private:
	const OodleCompressor		Compressor;
	const OodleCompressionLevel CompressionLevel;
};

class OodleDecoder final : public BlockDecoder
{
protected:
	bool DecompressBlock(MutableMemoryView RawData, MemoryView CompressedData) const final
	{
		const OO_SINTa RawSize = static_cast<OO_SINTa>(RawData.GetSize());
		const OO_SINTa Size	   = OodleLZ_Decompress(CompressedData.GetData(),
													static_cast<OO_SINTa>(CompressedData.GetSize()),
													RawData.GetData(),
													RawSize,
													OodleLZ_FuzzSafe_Yes,
													OodleLZ_CheckCRC_Yes,
													OodleLZ_Verbosity_None);
		return Size == RawSize;
	}
};

///////////////////////////////////////////////////////////////////////////////////////////////////////////////////////

class LZ4Decoder final : public BlockDecoder
{
protected:
	bool DecompressBlock(MutableMemoryView RawData, MemoryView CompressedData) const final
	{
		if (CompressedData.GetSize() <= std::numeric_limits<int>::max())
		{
			const int Size = LZ4_decompress_safe(static_cast<const char*>(CompressedData.GetData()),
												 static_cast<char*>(RawData.GetData()),
												 static_cast<int>(CompressedData.GetSize()),
												 static_cast<int>(zen::Min<uint64_t>(RawData.GetSize(), uint64_t(LZ4_MAX_INPUT_SIZE))));
			return static_cast<uint64_t>(Size) == RawData.GetSize();
		}
		return false;
	}
};

//////////////////////////////////////////////////////////////////////////

static const BaseDecoder*
GetDecoder(CompressionMethod Method)
{
	static NoneDecoder	None;
	static OodleDecoder Oodle;
	static LZ4Decoder	LZ4;

	switch (Method)
	{
		default:
			return nullptr;
		case CompressionMethod::None:
			return &None;
		case CompressionMethod::Oodle:
			return &Oodle;
		case CompressionMethod::LZ4:
			return &LZ4;
	}
}

//////////////////////////////////////////////////////////////////////////

bool
BufferHeader::IsValid(const CompositeBuffer& CompressedData)
{
	if (sizeof(BufferHeader) <= CompressedData.GetSize())
	{
		const BufferHeader Header = Read(CompressedData);
		if (Header.Magic == BufferHeader::ExpectedMagic)
		{
			if (const BaseDecoder* const Decoder = GetDecoder(Header.Method))
			{
				UniqueBuffer	 HeaderCopy;
				const MemoryView HeaderView = CompressedData.ViewOrCopyRange(0, Decoder->GetHeaderSize(Header), HeaderCopy);
				if (Header.Crc32 == BufferHeader::CalculateCrc32(HeaderView))
				{
					return true;
				}
			}
		}
	}
	return false;
}

//////////////////////////////////////////////////////////////////////////

template<typename BufferType>
inline CompositeBuffer
ValidBufferOrEmpty(BufferType&& CompressedData)
{
	return BufferHeader::IsValid(CompressedData) ? CompositeBuffer(std::forward<BufferType>(CompressedData)) : CompositeBuffer();
}

}  // namespace zen::detail

namespace zen {

const CompressedBuffer CompressedBuffer::Null;

CompressedBuffer
CompressedBuffer::Compress(const CompositeBuffer& RawData,
						   OodleCompressor		  Compressor,
						   OodleCompressionLevel  CompressionLevel,
						   uint64_t				  BlockSize)
{
	using namespace detail;

	if (BlockSize == 0)
	{
		BlockSize = DefaultBlockSize;
	}

	CompressedBuffer Local;
	if (CompressionLevel == OodleCompressionLevel::None)
	{
		Local.CompressedData = NoneEncoder().Compress(RawData, BlockSize);
	}
	else
	{
		Local.CompressedData = OodleEncoder(Compressor, CompressionLevel).Compress(RawData, BlockSize);
	}
	return Local;
}

CompressedBuffer
CompressedBuffer::Compress(const SharedBuffer&	 RawData,
						   OodleCompressor		 Compressor,
						   OodleCompressionLevel CompressionLevel,
						   uint64_t				 BlockSize)
{
	return Compress(CompositeBuffer(RawData), Compressor, CompressionLevel, BlockSize);
}

CompressedBuffer
CompressedBuffer::FromCompressed(const CompositeBuffer& InCompressedData)
{
	CompressedBuffer Local;
	Local.CompressedData = detail::ValidBufferOrEmpty(InCompressedData);
	return Local;
}

CompressedBuffer
CompressedBuffer::FromCompressed(CompositeBuffer&& InCompressedData)
{
	CompressedBuffer Local;
	Local.CompressedData = detail::ValidBufferOrEmpty(std::move(InCompressedData));
	return Local;
}

CompressedBuffer
CompressedBuffer::FromCompressed(const SharedBuffer& InCompressedData)
{
	CompressedBuffer Local;
	Local.CompressedData = detail::ValidBufferOrEmpty(InCompressedData);
	return Local;
}

CompressedBuffer
CompressedBuffer::FromCompressed(SharedBuffer&& InCompressedData)
{
	CompressedBuffer Local;
	Local.CompressedData = detail::ValidBufferOrEmpty(std::move(InCompressedData));
	return Local;
}

uint64_t
CompressedBuffer::GetRawSize() const
{
	return CompressedData ? detail::BufferHeader::Read(CompressedData).TotalRawSize : 0;
}

BLAKE3
CompressedBuffer::GetRawHash() const
{
	return CompressedData ? detail::BufferHeader::Read(CompressedData).RawHash : BLAKE3();
}

bool
CompressedBuffer::TryDecompressTo(MutableMemoryView RawView) const
{
	using namespace detail;
	if (CompressedData)
	{
		const BufferHeader Header = BufferHeader::Read(CompressedData);
		if (const BaseDecoder* const Decoder = GetDecoder(Header.Method))
		{
			return Decoder->TryDecompressTo(Header, CompressedData, RawView);
		}
	}
	return false;
}

SharedBuffer
CompressedBuffer::Decompress() const
{
	using namespace detail;
	if (CompressedData)
	{
		const BufferHeader Header = BufferHeader::Read(CompressedData);
		if (const BaseDecoder* const Decoder = GetDecoder(Header.Method))
		{
			if (Header.Method == CompressionMethod::None)
			{
				return Decoder->Decompress(Header, CompressedData).Flatten();
			}
			UniqueBuffer RawData = UniqueBuffer::Alloc(Header.TotalRawSize);
			if (Decoder->TryDecompressTo(Header, CompressedData, RawData))
			{
				return RawData.MoveToShared();
			}
		}
	}
	return SharedBuffer();
}

CompositeBuffer
CompressedBuffer::DecompressToComposite() const
{
	using namespace detail;
	if (CompressedData)
	{
		const BufferHeader Header = BufferHeader::Read(CompressedData);
		if (const BaseDecoder* const Decoder = GetDecoder(Header.Method))
		{
			return Decoder->Decompress(Header, CompressedData);
		}
	}
	return CompositeBuffer();
}

bool
CompressedBuffer::TryGetCompressParameters(OodleCompressor& OutCompressor, OodleCompressionLevel& OutCompressionLevel) const
{
	using namespace detail;
	if (CompressedData)
	{
		switch (const BufferHeader Header = BufferHeader::Read(CompressedData); Header.Method)
		{
			case CompressionMethod::None:
				OutCompressor		= OodleCompressor::NotSet;
				OutCompressionLevel = OodleCompressionLevel::None;
				return true;
			case CompressionMethod::Oodle:
				OutCompressor		= OodleCompressor(Header.Compressor);
				OutCompressionLevel = OodleCompressionLevel(Header.CompressionLevel);
				return true;
			default:
				break;
		}
	}
	return false;
}

/**
	______________________ _____________________________
	\__    ___/\_   _____//   _____/\__    ___/   _____/
	  |    |    |    __)_ \_____  \   |    |  \_____  \
	  |    |    |        \/        \  |    |  /        \
	  |____|   /_______  /_______  /  |____| /_______  /
					   \/        \/                  \/
 */

TEST_CASE("CompressedBuffer")
{
	uint8_t Zeroes[1024]{};
	uint8_t Ones[1024];
	memset(Ones, 1, sizeof Ones);

	{
		CompressedBuffer Buffer = CompressedBuffer::Compress(CompositeBuffer(SharedBuffer::MakeView(MakeMemoryView(Zeroes))),
															 OodleCompressor::NotSet,
															 OodleCompressionLevel::None);

		CHECK(Buffer.GetRawSize() == sizeof(Zeroes));
		CHECK(Buffer.GetCompressedSize() == (sizeof(Zeroes) + sizeof(detail::BufferHeader)));

		CompositeBuffer	 Compressed = Buffer.GetCompressed();
		CompressedBuffer BufferD	= CompressedBuffer::FromCompressed(Compressed);

		CHECK(BufferD.IsNull() == false);

		CompositeBuffer Decomp = BufferD.DecompressToComposite();

		CHECK(Decomp.GetSize() == Buffer.GetRawSize());
		CHECK(BLAKE3::HashBuffer(Decomp) == BufferD.GetRawHash());
	}

	{
		CompressedBuffer Buffer = CompressedBuffer::Compress(
			CompositeBuffer(SharedBuffer::MakeView(MakeMemoryView(Zeroes)), SharedBuffer::MakeView(MakeMemoryView(Ones))),
			OodleCompressor::NotSet,
			OodleCompressionLevel::None);

		CHECK(Buffer.GetRawSize() == (sizeof(Zeroes) + sizeof(Ones)));
		CHECK(Buffer.GetCompressedSize() == (sizeof(Zeroes) + sizeof(Ones) + sizeof(detail::BufferHeader)));

		CompositeBuffer	 Compressed = Buffer.GetCompressed();
		CompressedBuffer BufferD	= CompressedBuffer::FromCompressed(Compressed);

		CHECK(BufferD.IsNull() == false);

		CompositeBuffer Decomp = BufferD.DecompressToComposite();

		CHECK(Decomp.GetSize() == Buffer.GetRawSize());
		CHECK(BLAKE3::HashBuffer(Decomp) == BufferD.GetRawHash());
	}

	{
		CompressedBuffer Buffer = CompressedBuffer::Compress(CompositeBuffer(SharedBuffer::MakeView(MakeMemoryView(Zeroes))));

		CHECK(Buffer.GetRawSize() == sizeof(Zeroes));
		CHECK(Buffer.GetCompressedSize() < sizeof(Zeroes));

		CompositeBuffer	 Compressed = Buffer.GetCompressed();
		CompressedBuffer BufferD	= CompressedBuffer::FromCompressed(Compressed);

		CHECK(BufferD.IsNull() == false);

		CompositeBuffer Decomp = BufferD.DecompressToComposite();

		CHECK(Decomp.GetSize() == Buffer.GetRawSize());
		CHECK(BLAKE3::HashBuffer(Decomp) == BufferD.GetRawHash());
	}

	{
		CompressedBuffer Buffer = CompressedBuffer::Compress(
			CompositeBuffer(SharedBuffer::MakeView(MakeMemoryView(Zeroes)), SharedBuffer::MakeView(MakeMemoryView(Ones))));

		CHECK(Buffer.GetRawSize() == (sizeof(Zeroes) + sizeof(Ones)));
		CHECK(Buffer.GetCompressedSize() < (sizeof(Zeroes) + sizeof(Ones)));

		CompositeBuffer	 Compressed = Buffer.GetCompressed();
		CompressedBuffer BufferD	= CompressedBuffer::FromCompressed(Compressed);

		CHECK(BufferD.IsNull() == false);

		CompositeBuffer Decomp = BufferD.DecompressToComposite();

		CHECK(Decomp.GetSize() == Buffer.GetRawSize());
		CHECK(BLAKE3::HashBuffer(Decomp) == BufferD.GetRawHash());
	}
}

void
compress_forcelink()
{
}

}  // namespace zen