/*
 * Copyright (C) 2025 sillysagiri
 *
 * This program is free software: you can redistribute it and/or modify it under the terms of the GNU General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.
 */
 
#include <cmath>
#include <cstdio>
#include <libimagequant.h>
#include <stdexcept>
#include "header.hpp"

ImageMapped Quantize(const Image &image, Palette &palette_output, int num_colors, bool dither, uint8_t format)
{
	liq_attr *attr = liq_attr_create();

	liq_set_max_colors(attr, num_colors);
	liq_set_speed(attr, 1);
	liq_set_quality(attr, 0, 100);

	liq_image *liq_img = liq_image_create_rgba(attr, image.data.data(), image.width, image.height, 0);

	liq_result *result;
	if (LIQ_OK != liq_image_quantize(liq_img, attr, &result))
		throw std::runtime_error("failed to quantize histogram");

	if (dither) liq_set_dithering_level(result, 1.0f);
	else liq_set_dithering_level(result, 0.0f);

	ImageMapped output(image.width*image.height);

	if (LIQ_OK != liq_write_remapped_image(result, liq_img, output.data(), output.size()))
		throw std::runtime_error("failed to remap image");

	const liq_palette *pal = liq_get_palette(result);
	palette_output.resize(pal->count);
	for (int i=0; i<pal->count; i++)
	{
		palette_output[i].r = pal->entries[i].r;
		palette_output[i].g = pal->entries[i].g;
		palette_output[i].b = pal->entries[i].b;
		palette_output[i].a = pal->entries[i].a;
	}

	liq_image_destroy(liq_img);
	liq_attr_destroy(attr);
	liq_result_destroy(result);

	// ---

	if (format == Format_INDEXED_4)
	{
		ImageMapped packed;
		packed.reserve((output.size()+3) / 4);

		for (int chunk_start=0; chunk_start<output.size(); chunk_start+=4) {
			uint8_t current_byte = 0;

			for (int offset=0; offset<4 && (chunk_start+offset)<output.size(); ++offset) {
				current_byte |= (output[chunk_start + offset] & 0b11) << (offset * 2);
			}

			packed.push_back(current_byte);
		}
		
		return packed;
	}

	else if (format == Format_INDEXED_16)
	{
		ImageMapped packed;
		packed.reserve((output.size()+3) / 4);

		for (int offset=0; offset<output.size(); offset+=2) {
			uint8_t current_byte = 0;

			current_byte |= (output[offset] & 0b1111);
			if (offset+1 <output.size())
				current_byte |= (output[offset+1] & 0b1111) << 4;

			packed.push_back(current_byte);
		}

		return packed;
	}

	else if (format == Format_INDEXED_8A32)
	{
		ImageMapped packed;
		packed.reserve((output.size()+3) / 4);

		for (int offset=0; offset<output.size(); offset++) {
			uint8_t current_byte = 0;

			uint8_t alpha = image.data[offset*4+3] >> 3;
			current_byte |= (output[offset] & 0b111) | (alpha << 3);
			packed.push_back(current_byte);
		}

		return packed;
	}

	else if (format == Format_INDEXED_32A8)
	{
		ImageMapped packed(output.size());
		packed.reserve((output.size()+3) / 4);

		for (int offset=0; offset<output.size(); offset++) {
			uint8_t current_byte = 0;

			uint8_t alpha = image.data[offset*4+3] >> 5;
			current_byte |= (output[offset] & 0b11111) | (alpha << 5);
			packed.push_back(current_byte);
		}

		return packed;
	}

	return output;
}