//===----------------------------------------------------------------------===//
//                         DuckDB
//
// miniz_wrapper.hpp
//
//
//===----------------------------------------------------------------------===//

#pragma once

#include "miniz.hpp"
#include <string>
#include <stdexcept>

namespace duckdb {

enum class MiniZStreamType { MINIZ_TYPE_NONE, MINIZ_TYPE_INFLATE, MINIZ_TYPE_DEFLATE };

struct MiniZStream {
	static constexpr uint8_t GZIP_HEADER_MINSIZE = 10;
	static constexpr uint8_t GZIP_FOOTER_SIZE = 8;
	static constexpr uint8_t GZIP_COMPRESSION_DEFLATE = 0x08;
	static constexpr unsigned char GZIP_FLAG_UNSUPPORTED = 0x1 | 0x2 | 0x4 | 0x10 | 0x20;

public:
	MiniZStream() : type(MiniZStreamType::MINIZ_TYPE_NONE) {
		ResetStreamInternal();
	}

	~MiniZStream() {
		switch (type) {
		case MiniZStreamType::MINIZ_TYPE_INFLATE:
			duckdb_miniz::mz_inflateEnd(&stream);
			break;
		case MiniZStreamType::MINIZ_TYPE_DEFLATE:
			duckdb_miniz::mz_deflateEnd(&stream);
			break;
		default:
			break;
		}
	}

	void FormatException(const std::string &error_msg) {
		throw std::runtime_error(error_msg);
	}

	void FormatException(const char *error_msg, int mz_ret) {
		auto err = duckdb_miniz::mz_error(mz_ret);
		FormatException(error_msg + std::string(": ") + (err ? err : "Unknown error code"));
	}

	void Decompress(const char *compressed_data, size_t compressed_size, char *out_data, size_t out_size) {
		type = MiniZStreamType::MINIZ_TYPE_INFLATE;

		// Loop over blocks
		while (compressed_size > 0) {
			// Read block header
			if (compressed_size < GZIP_HEADER_MINSIZE) {
				FormatException("Failed to decompress GZIP block: compressed size is less than gzip header size");
			}
			auto gzip_hdr = reinterpret_cast<const unsigned char *>(compressed_data);
			if (gzip_hdr[0] != 0x1F || gzip_hdr[1] != 0x8B || gzip_hdr[2] != GZIP_COMPRESSION_DEFLATE ||
			    gzip_hdr[3] & GZIP_FLAG_UNSUPPORTED) {
				FormatException("Input is invalid/unsupported GZIP stream");
			}
			compressed_data += GZIP_HEADER_MINSIZE;
			compressed_size -= GZIP_HEADER_MINSIZE;

			// Initialize stream
			auto mz_ret = mz_inflateInit2(&stream, -MZ_DEFAULT_WINDOW_BITS);
			if (mz_ret != duckdb_miniz::MZ_OK) {
				FormatException("Failed to initialize miniz", mz_ret);
			}

			// Set up in/out
			stream.next_in = reinterpret_cast<const unsigned char *>(compressed_data);
			stream.avail_in = static_cast<unsigned int>(compressed_size);
			stream.next_out = reinterpret_cast<unsigned char *>(out_data);
			stream.avail_out = static_cast<unsigned int>(out_size);

			// Decompress and uninitialize stream
			mz_ret = mz_inflate(&stream, duckdb_miniz::MZ_FINISH);
			if (mz_ret != duckdb_miniz::MZ_OK && mz_ret != duckdb_miniz::MZ_STREAM_END) {
				FormatException("Failed to decompress GZIP block", mz_ret);
			}
			mz_inflateEnd(&stream);

			// Update indices
			compressed_data += GZIP_FOOTER_SIZE + stream.total_in;
			compressed_size -= GZIP_FOOTER_SIZE + stream.total_in;
			out_data += stream.total_out;
			out_size -= stream.total_out;

			ResetStreamInternal();
		}
	}

	static size_t MaxCompressedLength(size_t input_size) {
		return duckdb_miniz::mz_compressBound(input_size) + GZIP_HEADER_MINSIZE + GZIP_FOOTER_SIZE;
	}

	static void InitializeGZIPHeader(unsigned char *gzip_header) {
		memset(gzip_header, 0, GZIP_HEADER_MINSIZE);
		gzip_header[0] = 0x1F;
		gzip_header[1] = 0x8B;
		gzip_header[2] = GZIP_COMPRESSION_DEFLATE;
		gzip_header[3] = 0;
		gzip_header[4] = 0;
		gzip_header[5] = 0;
		gzip_header[6] = 0;
		gzip_header[7] = 0;
		gzip_header[8] = 0;
		gzip_header[9] = 0xFF;
	}

	static void InitializeGZIPFooter(unsigned char *gzip_footer, duckdb_miniz::mz_ulong crc, idx_t uncompressed_size) {
		gzip_footer[0] = crc & 0xFF;
		gzip_footer[1] = (crc >> 8) & 0xFF;
		gzip_footer[2] = (crc >> 16) & 0xFF;
		gzip_footer[3] = (crc >> 24) & 0xFF;
		gzip_footer[4] = uncompressed_size & 0xFF;
		gzip_footer[5] = (uncompressed_size >> 8) & 0xFF;
		gzip_footer[6] = (uncompressed_size >> 16) & 0xFF;
		gzip_footer[7] = (uncompressed_size >> 24) & 0xFF;
	}

	void Compress(const char *uncompressed_data, size_t uncompressed_size, char *out_data, size_t *out_size) {
		auto mz_ret =
		    mz_deflateInit2(&stream, duckdb_miniz::MZ_DEFAULT_LEVEL, MZ_DEFLATED, -MZ_DEFAULT_WINDOW_BITS, 1, 0);
		if (mz_ret != duckdb_miniz::MZ_OK) {
			FormatException("Failed to initialize miniz", mz_ret);
		}
		type = MiniZStreamType::MINIZ_TYPE_DEFLATE;

		auto gzip_header = reinterpret_cast<unsigned char *>(out_data);
		InitializeGZIPHeader(gzip_header);

		auto gzip_body = gzip_header + GZIP_HEADER_MINSIZE;

		stream.next_in = reinterpret_cast<const unsigned char *>(uncompressed_data);
		stream.avail_in = static_cast<unsigned int>(uncompressed_size);
		stream.next_out = gzip_body;
		stream.avail_out = static_cast<unsigned int>(*out_size - GZIP_HEADER_MINSIZE);

		mz_ret = mz_deflate(&stream, duckdb_miniz::MZ_FINISH);
		if (mz_ret != duckdb_miniz::MZ_OK && mz_ret != duckdb_miniz::MZ_STREAM_END) {
			FormatException("Failed to compress GZIP block", mz_ret);
		}
		auto gzip_footer = gzip_body + stream.total_out;
		auto crc = duckdb_miniz::mz_crc32(MZ_CRC32_INIT, reinterpret_cast<const unsigned char *>(uncompressed_data),
		                                  uncompressed_size);
		InitializeGZIPFooter(gzip_footer, crc, uncompressed_size);

		*out_size = stream.total_out + GZIP_HEADER_MINSIZE + GZIP_FOOTER_SIZE;
	}

private:
	void ResetStreamInternal() {
		memset(&stream, 0, sizeof(duckdb_miniz::mz_stream));
	}

private:
	duckdb_miniz::mz_stream stream;
	MiniZStreamType type;
};

} // namespace duckdb