#include <cstring>
#include <cstdio>
#include <cstdlib>
#include <limits>
#include <encodings/rzstd.h>

#include "Common/Data/Encoding/Compression.h"

#include "Common/Log.h"
#include "Common/Data/Format/ZIMLoad.h"
#include "Common/Math/math_util.h"
#include "Common/File/VFS/VFS.h"

int LoadZIMPtr(const uint8_t *zim, size_t datasize, int *width, int *height, int *flags, uint8_t **image) {
	if (datasize < 16) {
		ERROR_LOG(Log::IO, "ZIM file is too small");
		return 0;
	}
	if (zim[0] != 'Z' || zim[1] != 'I' || zim[2] != 'M' || zim[3] != 'G') {
		ERROR_LOG(Log::IO, "Not a ZIM file");
		return 0;
	}
	memcpy(width, zim + 4, 4);
	memcpy(height, zim + 8, 4);
	memcpy(flags, zim + 12, 4);

	if (*width <= 0 || *height <= 0) {
		ERROR_LOG(Log::IO, "Invalid ZIM dimensions %i x %i", *width, *height);
		return 0;
	}

	int num_levels = 1;
	size_t image_data_size[ZIM_MAX_MIP_LEVELS]{};
	if (*flags & ZIM_HAS_MIPS) {
		num_levels = log2i(*width < *height ? *width : *height) + 1;
	}
	if (num_levels > ZIM_MAX_MIP_LEVELS) {
		ERROR_LOG(Log::IO, "Too many ZIM mip levels: %i", num_levels);
		return 0;
	}

	size_t total_data_size = 0;
	for (int i = 0; i < num_levels; i++) {
		if (i > 0) {
			width[i] = width[i-1] / 2;
			height[i] = height[i-1] / 2;
		}

		const size_t level_width = (size_t)width[i];
		const size_t level_height = (size_t)height[i];
		size_t bytes_per_pixel;
		switch (*flags & ZIM_FORMAT_MASK) {
		case ZIM_RGBA8888:
			bytes_per_pixel = 4;
			break;
		case ZIM_RGBA4444:
		case ZIM_RGB565:
			bytes_per_pixel = 2;
			break;
		default:
			ERROR_LOG(Log::IO, "Invalid ZIM format %i", *flags & ZIM_FORMAT_MASK);
			return 0;
		}
		if (level_width > std::numeric_limits<size_t>::max() / level_height ||
			level_width * level_height > std::numeric_limits<size_t>::max() / bytes_per_pixel) {
			ERROR_LOG(Log::IO, "ZIM image size is too large");
			return 0;
		}
		image_data_size[i] = level_width * level_height * bytes_per_pixel;
		if (total_data_size > std::numeric_limits<size_t>::max() - image_data_size[i]) {
			ERROR_LOG(Log::IO, "ZIM mip data size is too large");
			return 0;
		}
		total_data_size += image_data_size[i];
	}

	if (total_data_size == 0) {
		ERROR_LOG(Log::IO, "Invalid ZIM data size 0");
		return 0;
	}

	const size_t payload_size = datasize - 16;
	if (!(*flags & (ZIM_ZLIB_COMPRESSED | ZIM_ZSTD_COMPRESSED)) && payload_size != total_data_size) {
		ERROR_LOG(Log::IO, "Wrong size data in ZIM: %i vs %i", (int)payload_size, (int)total_data_size);
		return 0;
	}
	if ((*flags & ZIM_ZLIB_COMPRESSED) &&
		(total_data_size > (size_t)std::numeric_limits<long>::max() || payload_size > (size_t)std::numeric_limits<long>::max())) {
		ERROR_LOG(Log::IO, "ZIM zlib data is too large");
		return 0;
	}

	image[0] = (uint8_t *)malloc(total_data_size);
	if (!image[0]) {
		ERROR_LOG(Log::IO, "Failed to allocate %lld bytes for ZIM image", (long long)total_data_size);
		return 0;
	}
	for (int i = 1; i < num_levels; i++) {
		image[i] = image[i-1] + image_data_size[i-1];
	}

	if (*flags & ZIM_ZLIB_COMPRESSED) {
		const int64_t outlen = InflateBuffer(15, zim + 16, (size_t)payload_size, *image, (size_t)total_data_size);
		if (outlen < 0) {
			ERROR_LOG(Log::IO, "ZIM zlib format decompression failed");
			free(*image);
			*image = 0;
			return 0;
		}
		if (outlen != total_data_size) {
			// Shouldn't happen if return value was Z_OK.
			ERROR_LOG(Log::IO, "Wrong size data in ZIM: %i vs %i", (int)outlen, (int)total_data_size);
		}
	} else if (*flags & ZIM_ZSTD_COMPRESSED) {
		size_t outlen = 0;
		if (rzstd_decode(*image, total_data_size, zim + 16, payload_size, &outlen) != RZSTD_PROCESS_END || outlen != (size_t)total_data_size) {
			ERROR_LOG(Log::IO, "ZIM zstd format decompression failed: %lld", (long long)outlen);
			free(*image);
			*image = 0;
			return 0;
		}
	} else {
		memcpy(*image, zim + 16, payload_size);
	}
	return num_levels;
}

int LoadZIM(const char *filename, int *width, int *height, int *format, uint8_t **image) {
	size_t size;
	uint8_t *buffer = g_VFS.ReadFile(filename, &size);
	if (!buffer) {
		ERROR_LOG(Log::IO, "Couldn't read data for '%s'", filename);
		return 0;
	}

	int retval = LoadZIMPtr(buffer, size, width, height, format, image);
	if (!retval) {
		ERROR_LOG(Log::IO, "Not a valid ZIM file: %s (size: %lld bytes)", filename, (long long)size);
	}
	delete [] buffer;
	return retval;
}