/*
 * Copyright (C) 2024 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 <algorithm>
#include <cassert>
#include <cstdint>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <stdexcept>
#include <format>
#include <iostream>
#include <vector>

#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION

#include "argh.h"
#include "libimagequant.h"
#include "header.hpp"
#include "external/Timer.h"

namespace fs = std::filesystem;

static void imagegen_job(uint8_t *image_output, uint8_t *image_indexed, const liq_palette *pal, int start, int end)
{
	for (int i=start; i<end; i++)
	{
		auto &col = pal->entries[image_indexed[i]];
		image_output[i*4+0] = col.r;
		image_output[i*4+1] = col.g;
		image_output[i*4+2] = col.b;
		image_output[i*4+3] = col.a;
	}
}

int main(int argc, char **argv)
{
	try {
		srand(time(0));

		Params params;
		params.Parse(argc, argv);
		params.Convert();
	}
	catch(const std::exception &e)
	{
		std::cerr << std::format("Exception caught: {}\n", e.what());
		return 1;
	}
	catch (...)
	{
		std::cerr << std::format("An unknown exception occurred\n");
		return 1;
	}

	return 0;
}

void Params::Parse(int argc, char **argv)
{
	argh::parser cmdl({
		"-f", "--format",
		"-o", "--out",
		"-c", "--compres",
		"-p", "--palette",
		"-pid", "--palette-id",
	});

	cmdl.parse(argc, argv);

	if (cmdl.size() == 1) PrintHelp();

	verbose = cmdl[{ "-v", "--verbose" }];
	paletteOnly = cmdl[{ "-po", "--palette-only" }];
	cmdl({"-pid", "--palette-id"}, rand() % 256) >> paletteId;

	// file input checking
	{
		std::stringstream ss(cmdl(1, "").str());
		std::string token;

		while (getline(ss, token, ','))
			input_list.push_back(token);

		for (auto &i : input_list)
			if (i.empty() || !fs::exists(i)) 
				throw std::runtime_error(std::format("file doesnt exsist: {}", i.string()));
	}

	// output directory checking
	{
		output_dir = cmdl({"-o", "--out"}, "").str();

		if (output_dir.empty() || !fs::is_directory(output_dir))
			throw std::runtime_error(std::format("invalid output directory: {}", output_dir.string()));
	}

	// format checking
	{
		std::string temp = cmdl({"-f", "--format"}, "rgb16").str();

		if (temp == "rgb256" || temp == "rgb255" || temp == "rgb")
			format = Format_RGB_256;

		else if (temp == "rgb16" || temp == "rgb15")
			format = Format_RGB_16;

		else if (temp == "indexed4" || temp == "4")
			format = Format_INDEXED_4;

		else if (temp == "indexed16" || temp == "16")
			format = Format_INDEXED_16;

		else if (temp == "indexed256" || temp == "256")
			format = Format_INDEXED_256;

		else if (temp == "indexed32a3" || temp == "i5a3")
			format = Format_INDEXED_32A3;

		else if (temp == "indexed8a5" || temp == "i3a5")
			format = Format_INDEXED_8A5;

		else throw std::runtime_error(std::format("Invalid format: {}", temp));
	}

	// compress checking
	{
		std::string temp = cmdl({"-c", "--compress"}, "none").str();

		if (temp == "none")
			compress = Compress_NONE;

		else if (temp == "lzsswram" || temp == "lzss")
			compress = Compress_LZSS_WRAM;

		else if (temp == "lzssvram")
			compress = Compress_LZSS_VRAM;

		else if (temp == "gzip")
			compress = Compress_GZIP;

		else throw std::runtime_error(std::format("Invalid compress: {}", temp));
	}

	// num check
	{
		cmdl({"-n", "--num"}, 0) >> limitColor;

		if (limitColor > 0)
		{
			switch (format) {
				case Format_RGB_256: limitColor = std::clamp(limitColor, 0, 256); break;
				case Format_RGB_16: limitColor = std::clamp(limitColor, 0, 16); break;
				case Format_INDEXED_4: limitColor = std::clamp(limitColor, 0, 4); break;
				case Format_INDEXED_16: limitColor = std::clamp(limitColor, 0, 16); break;
				case Format_INDEXED_256: limitColor = std::clamp(limitColor, 0, 256); break;
				case Format_INDEXED_32A3: limitColor = std::clamp(limitColor, 0, 32); break;
				case Format_INDEXED_8A5: limitColor = std::clamp(limitColor, 0, 8); break;
				case Format_PALETTE_16: break;
			}
		}
	}

	Print("\n==============\n");
	Print("input:\n");
	for (auto &i : input_list)
		Print(std::format("  {}\n", i.string()));
	Print(std::format("output dir: {}\n", output_dir.string()));
	
	Print(std::format("format: {}\n", static_cast<int>(format)));
	Print(std::format("compress: {}\n", static_cast<int>(compress)));
	Print(std::format("paletteId: {}\n", static_cast<int>(paletteId)));
	
	if (limitColor > 0)
		Print(std::format("limitColor: {}\n", limitColor));

	Print("==============\n\n");
}

void Params::Convert()
{

	switch (format) {

		case Format_RGB_256:
		{
			for (int i=0; i<input_list.size(); i++)
			{
				Image img;
				img.Load(input_list[i]);

				WriteImage(
					(output_dir/input_list[i].stem()).string() + "_img.bin",
					format, img.width, img.height, compress, 0,
					img.data, img.width*img.height*4);
			}

			break;
		}

		case Format_RGB_16:
		{
			for (int i=0; i<input_list.size(); i++)
			{
				Image img;
				img.Load(input_list[i]);
				uint16_t *img16 = new uint16_t[img.width*img.height];

				for (int i2=0; i2<img.width*img.height; i2++)
					img16[i2*2] = RGB16(img.data[i2*4+3], img.data[i2*4+0], img.data[i2*4+1], img.data[i2*4+2]);

				WriteImage(
					output_dir / std::format("{}_img.bin", input_list[i].stem().string()),
					format, img.width, img.height, compress, 0,
					img16, img.width*img.height*2);

				delete[] img16;
			}

			break;
		}

		// TODO: wrap indexed into function to reduce code repetition
		case Format_INDEXED_4:
		{
			std::vector<Color> palette;
			std::vector<Image> images(input_list.size());
			std::vector<IndexedImage> indexes(input_list.size());

			for (int i=0; i<images.size(); i++)
				images[i].Load(input_list[i]);
			
			GetIndexed(4, images, palette, indexes);
			WritePalette(output_dir / std::format("{}_pal.bin", paletteId), paletteId, palette);

			
			for (int i=0; i<indexes.size(); i++)
			{
				uint32_t length;
				BitPacking(2, indexes[i], length);

				WriteImage(
					output_dir / std::format("{}_img.bin", input_list[i].stem().string()),
					format, indexes[i].width, indexes[i].height, compress, paletteId,
					indexes[i].data, length);
			}

			break;
		}

		case Format_INDEXED_16:
		{
			std::vector<Color> palette;
			std::vector<Image> images(input_list.size());
			std::vector<IndexedImage> indexes(input_list.size());

			for (int i=0; i<images.size(); i++)
				images[i].Load(input_list[i]);

			GetIndexed(16, images, palette, indexes);
			WritePalette(output_dir / std::format("{}_pal.bin", paletteId), paletteId, palette);

			for (int i=0; i<indexes.size(); i++)
			{
				uint32_t length;
				BitPacking(4, indexes[i], length);

				WriteImage(
					output_dir / std::format("{}_img.bin", input_list[i].stem().string()),
					format, indexes[i].width, indexes[i].height, compress, paletteId,
					indexes[i].data, length);
			}

			break;
		}

		case Format_INDEXED_256:
		{
			std::vector<Color> palette;
			std::vector<Image> images(input_list.size());
			std::vector<IndexedImage> indexes(input_list.size());

			for (int i=0; i<images.size(); i++)
				images[i].Load(input_list[i]);

			GetIndexed(256, images, palette, indexes);
			WritePalette(output_dir / std::format("{}_pal.bin", paletteId), paletteId, palette);

			for (int i=0; i<indexes.size(); i++)
				WriteImage(
					output_dir / std::format("{}_img.bin", input_list[i].stem().string()),
					format, indexes[i].width, indexes[i].height, compress, paletteId,
					indexes[i].data, indexes[i].width*indexes[i].height);

			break;
		}

		case Format_INDEXED_32A3:
		case Format_INDEXED_8A5:
		case Format_PALETTE_16:

		throw std::runtime_error("currently not implemented yet!");
		break;
	}	
}

void Params::PrintHelp()
{
	// std::cout << "quantization <input path> [output path]" << "\n\n";
}

void Params::Print(const std::string &msg)
{
	if (verbose) std::cout << msg;
}

// ----

void WritePalette(const std::string &path, const uint8_t paletteId, std::vector<Color> &palette)
{
	Metadata meta;

	meta.format = Format_PALETTE_16;
	meta.paletteId = paletteId;
	meta.width = palette.size();
	meta.height = palette.size();
	meta.compression = Compress_NONE;
	meta.length = palette.size()*2;

	std::ofstream out(path, std::ios::binary);
	out.write(reinterpret_cast<const char*>(&meta), sizeof(meta));

	for (const auto &i : palette)
	{
		uint16_t rgb16 = RGB16(i.a, i.r, i.g, i.b);
		out.write(reinterpret_cast<const char*>(&rgb16), sizeof(rgb16));
	}

	out.close();
}

void WriteImage(const std::string &path, const Format &format, const uint16_t width, const uint16_t height, const Compress &compress, const uint8_t paletteId, void *buffer, const uint32_t length)
{
	Metadata meta;

	meta.format = format;
	meta.paletteId = paletteId;
	meta.width = width;
	meta.height = height;
	meta.compression = Compress_NONE;
	meta.length = length;

	std::ofstream out(path, std::ios::binary);
	out.write(reinterpret_cast<const char*>(&meta), sizeof(meta));
	out.write(reinterpret_cast<const char*>(buffer), length);
	out.close();
}


// ----

uint16_t RGB16(uint8_t a, uint8_t r, uint8_t g, uint8_t b)
{
	uint16_t a1 = a & 0x01;
	uint16_t r5 = (r >> 3) & 0x1F;
	uint16_t g5 = (g >> 3) & 0x1F;
	uint16_t b5 = (b >> 3) & 0x1F;

	return (a1 << 15) | (r5) | (g5 << 5) | (b5 << 10);
}

void GetIndexed(const int numColor, std::vector<Image> &images, std::vector<Color> &palette, std::vector<IndexedImage> &indexes)
{
	liq_attr *liq_attr;
	liq_result *liq_result;
	liq_histogram *liq_histogram;
	const liq_palette *liq_palette;

	liq_attr = liq_attr_create();
	liq_histogram = liq_histogram_create(liq_attr);

	liq_set_speed(liq_attr, 1);
	liq_set_max_colors(liq_attr, numColor);

	std::vector<liq_image*> liq_images(images.size());

	for (int i=0; i<images.size(); i++)
	{
		liq_images[i] = liq_image_create_rgba(liq_attr, images[i].data, images[i].width, images[i].height, 0); 
		liq_histogram_add_image(liq_histogram, liq_attr, liq_images[i]);
	}

	auto liq_err = liq_histogram_quantize(liq_histogram, liq_attr, &liq_result);
	liq_set_dithering_level(liq_result, 0);
	liq_palette = liq_get_palette(liq_result);
	assert(liq_err == LIQ_OK);

	palette.resize(liq_palette->count);
	for(int i=0; i<liq_palette->count; i++)
		palette[i] = {
			liq_palette->entries[i].a,
			liq_palette->entries[i].r,
			liq_palette->entries[i].g,
			liq_palette->entries[i].b };

	for (int i=0; i<images.size(); i++)
	{
		indexes[i].Load(images[i].width, images[i].height);
		liq_err = liq_write_remapped_image(liq_result, liq_images[i], indexes[i].data, indexes[i].width*indexes[i].height);
		if (liq_err != LIQ_OK) throw std::runtime_error(std::format("{} {} {}", static_cast<int>(liq_err), indexes[i].width, indexes[i].height));
	}

	for (auto &i : liq_images)
		liq_image_destroy(i);
	liq_attr_destroy(liq_attr);
	liq_histogram_destroy(liq_histogram);
	liq_result_destroy(liq_result);
}

void BitPacking(const uint8_t bpp, IndexedImage &img, uint32_t &length)
{
	length = ((img.width*img.height) * bpp) / 8;
	uint8_t *buffer = new uint8_t[length]();
	uint32_t bufferPos = 0;
	uint8_t bitpos = 0;

	for (int i2=0; i2<(img.width*img.height); i2++)
	{
		buffer[bufferPos] |= img.data[i2] << bitpos;
		bitpos += bpp;

		if (bitpos == 8)
		{
			bitpos = 0;
			bufferPos++;
		}
	}

	delete[] img.data;
	img.data = buffer;
}

// ----

Image::Image(const std::filesystem::path &path)
{
	Load(path);
}

Image::~Image()
{
	stbi_image_free(data);
}

void Image::Load(const std::filesystem::path &path)
{
	data = stbi_load(path.c_str(), &width, &height, &comp, 4);
	if (data == nullptr) throw std::runtime_error(stbi_failure_reason());
}

// ----

IndexedImage::IndexedImage(int width, int height)
	: width(width), height(height), data(nullptr)
{
	Load(width, height);
}

IndexedImage::~IndexedImage()
{
	delete[] data; 
}

void IndexedImage::Load(int width, int height)
{
	this->width = width;
	this->height = height;

	data = new uint8_t[width*height];
	if (data == nullptr) throw std::runtime_error("not enough memory");
}