// 内置 inflate（RFC 1950 / 1951）—— 供 HDF5 分块与 MDF4 `##DZ` 共用。
//
// ⚠️ **不能用 zlib（`-lz`）**：链接 libz.a 会让 Release 产物多出约 46KB，
// 而这里只需要「解压」这一个功能，自带的紧凑实现只要几 KB。
#include "Deflate.hpp"

namespace jdi
{
	namespace
	{
		/** @brief 位流读取器（LSB 优先，与 DEFLATE 一致）。 */
		struct BitReader
		{
			const std::uint8_t* data = nullptr;
			std::size_t size = 0;
			std::size_t position = 0;
			std::uint32_t buffer = 0;
			int bits = 0;
			bool error = false;

			bool Fill(int count)
			{
				while (bits < count)
				{
					if (position >= size)
					{
						error = true;
						return false;
					}
					buffer |= static_cast<std::uint32_t>(data[position++]) << bits;
					bits += 8;
				}
				return true;
			}

			std::uint32_t Read(int count)
			{
				if (count <= 0)
				{
					return 0;
				}
				if (!Fill(count))
				{
					return 0;
				}
				const std::uint32_t value = buffer & ((1u << count) - 1u);
				buffer >>= count;
				bits -= count;
				return value;
			}

			void Align()
			{
				buffer = 0;
				bits = 0;
			}
		};

		/** @brief 规范 Huffman 表（只存计数与按长度排好序的符号 + 快查表）。 */
		constexpr int kFastBits = 9;
		constexpr int kFastSize = 1 << kFastBits;

		struct HuffmanTable
		{
			std::uint16_t count[16] = {};
			std::uint16_t symbol[320] = {};
			/**
			 * @brief 前缀码快查表（9 位）：`fastLength[peek] != 0` 时可直接得到符号。
			 *
			 * ⚠️ 逐位解码（每读 1 bit 就查一次表）是解压的绝对瓶颈：GB 级 `.mat` 实测
			 * 只有 ~11MB/s，而 9 位查表一次就能吃掉绝大多数符号（典型单字节字面量长度 8~9）。
			 */
			std::uint16_t fastSymbol[kFastSize] = {};
			std::uint8_t fastLength[kFastSize] = {};
		};

		bool BuildHuffman(HuffmanTable& table, const std::uint8_t* lengths, int total)
		{
			for (int index = 0; index < 16; ++index)
			{
				table.count[index] = 0;
			}
			for (int index = 0; index < total; ++index)
			{
				++table.count[lengths[index]];
			}
			if (table.count[0] == total)
			{
				return true;	// 全部长度为 0：该表没有任何码（动态块里可能是距离表）
			}
			int left = 1;
			for (int length = 1; length < 16; ++length)
			{
				left <<= 1;
				left -= table.count[length];
				if (left < 0)
				{
					return false;	// 超订（over-subscribed）
				}
			}
			std::uint16_t offsets[16] = {};
			offsets[1] = 0;
			for (int length = 1; length < 15; ++length)
			{
				offsets[length + 1] = static_cast<std::uint16_t>(offsets[length] + table.count[length]);
			}
			for (int index = 0; index < total; ++index)
			{
				if (lengths[index] != 0)
				{
					table.symbol[offsets[lengths[index]]++] = static_cast<std::uint16_t>(index);
				}
			}

			// 建 9 位快查表：按规范码给符号分配 code，再按 DEFLATE 的 LSB-first 位序反转。
			for (int index = 0; index < kFastSize; ++index)
			{
				table.fastLength[index] = 0;
			}
			int code = 0;
			int symbolIndex = 0;
			for (int length = 1; length < 16; ++length)
			{
				for (int i = 0; i < table.count[length]; ++i)
				{
					const int symbol = table.symbol[symbolIndex + i];
					if (length <= kFastBits)
					{
						int reversed = 0;
						for (int bit = 0; bit < length; ++bit)
						{
							reversed |= ((code >> bit) & 1) << (length - 1 - bit);
						}
						const int step = 1 << length;
						for (int fill = reversed; fill < kFastSize; fill += step)
						{
							table.fastSymbol[fill] = static_cast<std::uint16_t>(symbol);
							table.fastLength[fill] = static_cast<std::uint8_t>(length);
						}
					}
					++code;
				}
				code <<= 1;
				symbolIndex += table.count[length];
			}
			return true;
		}

		int DecodeSymbol(BitReader& reader, const HuffmanTable& table)
		{
			// 快路径：一次填到 9 位，查表命中则直接返回（绝大多数文字/长度符号）。
			if (reader.bits >= kFastBits || reader.Fill(kFastBits))
			{
				const int peek = static_cast<int>(reader.buffer & (kFastSize - 1));
				const int length = table.fastLength[peek];
				if (length > 0)
				{
					reader.buffer >>= length;
					reader.bits -= length;
					return table.fastSymbol[peek];
				}
			}
			else
			{
				// 结尾处不足 9 位：交给慢路径逐位解码（它才是"数据真的结束"的判据）。
				reader.error = false;
			}

			int code = 0;
			int first = 0;
			int index = 0;
			for (int length = 1; length < 16; ++length)
			{
				code |= static_cast<int>(reader.Read(1));
				if (reader.error)
				{
					return -1;
				}
				const int count = table.count[length];
				if (code - first < count)
				{
					return table.symbol[index + (code - first)];
				}
				index += count;
				first = (first + count) << 1;
				code <<= 1;
			}
			return -1;
		}

		const std::uint16_t kLengthBase[29] = {3,  4,  5,  6,  7,  8,  9,  10, 11,  13,
		                                       15, 17, 19, 23, 27, 31, 35, 43, 51,  59,
		                                       67, 83, 99, 115, 131, 163, 195, 227, 258};
		const std::uint8_t kLengthExtra[29] = {0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2,
		                                       2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0};
		const std::uint16_t kDistanceBase[30] = {1,    2,    3,    4,    5,    7,     9,     13,
		                                         17,   25,   33,   49,   65,   97,    129,   193,
		                                         257,  385,  513,  769,  1025, 1537,  2049,  3073,
		                                         4097, 6145, 8193, 12289, 16385, 24577};
		const std::uint8_t kDistanceExtra[30] = {0, 0, 0, 0, 1, 1, 2, 2,  3,  3,  4,  4,  5,  5,  6,
		                                         6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13};
	} // namespace

	bool LooksLikeZlib(const std::uint8_t* source, std::size_t sourceSize)
	{
		if (source == nullptr || sourceSize < 6)
		{
			return false;
		}
		const std::uint8_t cmf = source[0];
		const std::uint8_t flg = source[1];
		return (cmf & 0x0F) == 8 && ((static_cast<unsigned>(cmf) * 256u + flg) % 31u) == 0 &&
		       (flg & 0x20) == 0;
	}

	bool InflateZlib(const std::uint8_t* source, std::size_t sourceSize,
	                 std::vector<std::uint8_t>& out, std::size_t expectedSize)
	{
		if (!LooksLikeZlib(source, sourceSize))
		{
			return false;
		}
		BitReader reader;
		reader.data = source + 2;
		reader.size = sourceSize - 6;	// 末尾 4 字节是 Adler32
		out.clear();
		if (expectedSize > 0)
		{
			out.reserve(expectedSize);
		}
		HuffmanTable fixedLiterals;
		HuffmanTable fixedDistances;
		bool fixedReady = false;

		for (int block = 0; block < 100000; ++block)
		{
			const std::uint32_t final = reader.Read(1);
			const std::uint32_t type = reader.Read(2);
			if (reader.error)
			{
				return false;
			}
			if (type == 0)
			{
				// 存储块（未压缩）：对齐后 4 字节长度 + 反码 + 原始字节。
				reader.Align();
				if (reader.position + 4 > reader.size)
				{
					return false;
				}
				const std::uint16_t length = static_cast<std::uint16_t>(
				    reader.data[reader.position] | (reader.data[reader.position + 1] << 8));
				reader.position += 4;
				if (reader.position + length > reader.size)
				{
					return false;
				}
				out.insert(out.end(), reader.data + reader.position,
				           reader.data + reader.position + length);
				reader.position += length;
			}
			else if (type == 1 || type == 2)
			{
				HuffmanTable literals;
				HuffmanTable distances;
				if (type == 1)
				{
					if (!fixedReady)
					{
						std::uint8_t lengths[288];
						for (int index = 0; index < 144; ++index)
						{
							lengths[index] = 8;
						}
						for (int index = 144; index < 256; ++index)
						{
							lengths[index] = 9;
						}
						for (int index = 256; index < 280; ++index)
						{
							lengths[index] = 7;
						}
						for (int index = 280; index < 288; ++index)
						{
							lengths[index] = 8;
						}
						if (!BuildHuffman(fixedLiterals, lengths, 288))
						{
							return false;
						}
						std::uint8_t distanceLengths[30];
						for (int index = 0; index < 30; ++index)
						{
							distanceLengths[index] = 5;
						}
						if (!BuildHuffman(fixedDistances, distanceLengths, 30))
						{
							return false;
						}
						fixedReady = true;
					}
					literals = fixedLiterals;
					distances = fixedDistances;
				}
				else
				{
					const int literalCount = static_cast<int>(reader.Read(5)) + 257;
					const int distanceCount = static_cast<int>(reader.Read(5)) + 1;
					const int codeCount = static_cast<int>(reader.Read(4)) + 4;
					if (reader.error || literalCount > 286 || distanceCount > 30)
					{
						return false;
					}
					static const std::uint8_t kOrder[19] = {16, 17, 18, 0, 8,  7, 9,  6, 10, 5,
					                                        11, 4,  12, 3, 13, 2, 14, 1, 15};
					std::uint8_t codeLengths[19] = {};
					for (int index = 0; index < codeCount; ++index)
					{
						codeLengths[kOrder[index]] = static_cast<std::uint8_t>(reader.Read(3));
					}
					HuffmanTable codeTable;
					if (reader.error || !BuildHuffman(codeTable, codeLengths, 19))
					{
						return false;
					}
					std::uint8_t lengths[288 + 30] = {};
					int produced = 0;
					while (produced < literalCount + distanceCount)
					{
						const int symbol = DecodeSymbol(reader, codeTable);
						if (symbol < 0)
						{
							return false;
						}
						if (symbol < 16)
						{
							lengths[produced++] = static_cast<std::uint8_t>(symbol);
							continue;
						}
						int repeat = 0;
						std::uint8_t value = 0;
						if (symbol == 16)
						{
							if (produced == 0)
							{
								return false;
							}
							value = lengths[produced - 1];
							repeat = 3 + static_cast<int>(reader.Read(2));
						}
						else if (symbol == 17)
						{
							repeat = 3 + static_cast<int>(reader.Read(3));
						}
						else
						{
							repeat = 11 + static_cast<int>(reader.Read(7));
						}
						if (reader.error || produced + repeat > literalCount + distanceCount)
						{
							return false;
						}
						for (int index = 0; index < repeat; ++index)
						{
							lengths[produced++] = value;
						}
					}
					if (!BuildHuffman(literals, lengths, literalCount) ||
					    !BuildHuffman(distances, lengths + literalCount, distanceCount))
					{
						return false;
					}
				}

				for (int symbolIndex = 0; symbolIndex < 100000000; ++symbolIndex)
				{
					const int symbol = DecodeSymbol(reader, literals);
					if (symbol < 0)
					{
						return false;
					}
					if (symbol < 256)
					{
						out.push_back(static_cast<std::uint8_t>(symbol));
						continue;
					}
					if (symbol == 256)
					{
						break;	// 块结束
					}
					const int lengthCode = symbol - 257;
					if (lengthCode >= 29)
					{
						return false;
					}
					const std::size_t length =
					    kLengthBase[lengthCode] + reader.Read(kLengthExtra[lengthCode]);
					const int distanceCode = DecodeSymbol(reader, distances);
					if (distanceCode < 0 || distanceCode >= 30 || reader.error)
					{
						return false;
					}
					const std::size_t distance =
					    kDistanceBase[distanceCode] + reader.Read(kDistanceExtra[distanceCode]);
					if (reader.error || distance == 0 || distance > out.size())
					{
						return false;
					}
					const std::size_t start = out.size() - distance;
					for (std::size_t index = 0; index < length; ++index)
					{
						out.push_back(out[start + index]);
					}
				}
			}
			else
			{
				return false;	// 类型 3 非法
			}
			if (final != 0)
			{
				return true;
			}
		}
		return false;
	}
} // namespace jdi