/*
 * 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 <cstdint>
#include <fstream>
#include <stdint.h>
#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION

#include <argparse/argparse.hpp>
#include "header.hpp"
#include "Timer.h"
#include "fastlz.h"

namespace fs = std::filesystem;

int main(int argc, char **argv)
{
	srand(time(0));
	argparse::ArgumentParser program("sillyimage", "0.14");
	Options opt;

	// Register argument
	{
		program.add_argument("-i", "--input")
			.help("input files")
			.default_value<std::string>("")
			.required()
			.nargs(1)
			.store_into(opt.path_input);

		program.add_argument("-o", "--out")
			.help("output directory")
			.required()
			.store_into(opt.path_output);

		program.add_argument("-f", "--format")
			.help("texture format { rgb256, rgb16, indexed4, indexed16, indexed256, indexed32a8, indexed8a32 }")
			.default_value<std::string>("rgb16")
			.choices("rgb256", "rgb16", "indexed4", "indexed16", "indexed256", "indexed32a8", "indexed8a32")
			.nargs(1)
			.store_into(opt.format);

		program.add_argument("-pf", "--palette-format")
			.help("palette format { rgb16, rgb256 }")
			.default_value<std::string>("rgb16")
			.choices("rgb16", "rgb256")
			.nargs(1)
			.store_into(opt.palette_format);

		// program.add_argument("-c", "--compression")
		// 	.help("compression { none, fastlz }")
		// 	.default_value<std::string>("fastlz")
		// 	.choices("none", "fastlz")
		// 	.nargs(1)
		// 	.store_into(opt.compression);

		program.add_argument("-be", "--big-endian")
			.help("enable big endian mode")
			.default_value<bool>(false)
			.nargs(1)
			.store_into(opt.enable_be);
	}

	// Parse argument
	{
		try
		{
			program.parse_args(argc, argv);

			// input list
			if (opt.path_input.empty() || !fs::exists(opt.path_input)) 
				throw std::runtime_error(std::format("file doesnt exsist: \"{}\"", opt.path_input));

			// output directory
			if (opt.path_output.empty() || !fs::is_directory(opt.path_output))
				throw std::runtime_error(std::format("invalid output directory: \"{}\"", opt.path_output));

		}
		catch(const std::exception &e)
		{
			std::cerr << std::format("Error: {}\n\n", e.what()) << program << "\n";
			return 1;
		}
	}

	// Convert
	try {

		if (opt.format == "rgb256")
		{
			Image img;
			img.Load(opt.path_input);

			uint32_t img255_size = img.width*img.height;
			uint32_t img255[img255_size];
			for (int i=0; i<img255_size; i++)
			{
				img255[i] = (static_cast<uint32_t>(img.data[i*4+3]) << 24) |  // A
										(static_cast<uint32_t>(img.data[i*4+0]) << 16) |  // R
										(static_cast<uint32_t>(img.data[i*4+1]) << 8)  |  // G
										 static_cast<uint32_t>(img.data[i*4+2]);          // B

				if (opt.enable_be)
					img255[i] = to_be_int32(img255[i]);
			}

			Metadata meta;
			meta.format = Format::Format_RGB_256;
			meta.width = img.width;
			meta.height = img.height;
			meta.palette_count = 0;
			meta.original_size = img.data.size();

			uint8_t compress[uint64_t(meta.original_size*1.5)];
			meta.compress_size = fastlz_compress_level(1, img.data.data(), meta.original_size, compress);

			Palette palette(0);

			WriteOutput(opt, meta, palette, compress);
		}

		if (opt.format == "rgb16")
		{
			Image img;
			img.Load(opt.path_input);
			std::vector<uint16_t> img16(img.width*img.height);

			// convert into RGB16 color
			for (int i=0; i<img.width*img.height; i++)
				img16[i] = RGB255_RGB16(img.data[i*4], img.data[i*4+1], img.data[i*4+2], img.data[i*4+3]);

			Metadata meta;
			meta.format = Format::Format_RGB_16;
			meta.width = img.width;
			meta.height = img.height;
			meta.palette_count = 0;
			meta.original_size = img16.size()*sizeof(uint16_t);

			uint8_t compress[uint64_t(meta.original_size*1.5)];
			meta.compress_size = fastlz_compress_level(2, img16.data(), meta.original_size, compress);

			Palette palette(0);

			WriteOutput(opt, meta, palette, compress);
		}

		if (opt.format == "indexed4" ||
				opt.format == "indexed16" ||
				opt.format == "indexed256" ||
				opt.format == "indexed32a8" ||
				opt.format == "indexed8a32")
		{
			uint16_t numcol;
			uint8_t format;

			if (opt.format == "indexed4")
			{
				numcol = 4;
				format = Format_INDEXED_4;
			}
			else if (opt.format == "indexed16")
			{
				numcol = 16;
				format = Format_INDEXED_16;
			}
			else if (opt.format == "indexed256")
			{
				numcol = 256;
				format = Format_INDEXED_256;
			}
			else if (opt.format == "indexed32a8")
			{
				numcol = 32;
				format = Format_INDEXED_32A8;
			}
			else if (opt.format == "indexed8a32")
			{
				numcol = 8;
				format = Format_INDEXED_8A32;
			}

			Palette palette;
			Image image;
			image.Load(opt.path_input);

			ImageMapped remap = Quantize(image, palette, numcol, false, format);
			
			Metadata meta;
			meta.format = format;
			meta.width = image.width;
			meta.height = image.height;
			meta.palette_count = palette.size();
			meta.original_size = remap.size();

			uint8_t compress[uint64_t(meta.original_size*1.5)];
			meta.compress_size = fastlz_compress_level(2, remap.data(), meta.original_size, compress);

			WriteOutput(opt, meta, palette, compress);
		}
	}
	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;
	}

	std::cout  << "done\n";

	return 0;
}

// // ----

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

	int size = width*height*4;
	data.reserve(size);
	for (int i = 0; i < size; ++i)
		data.push_back(ptr_img[i]);

	stbi_image_free(ptr_img);
}

void WriteOutput(const Options &opt, const Metadata &meta, const Palette &palette, void *image_buffer)
{
	std::filesystem::path input_file(opt.path_input);
	std::filesystem::path output_dir(opt.path_output);
	std::string basename = input_file.stem();
	
	std::string output_file = output_dir / std::format("{}.sillyimg", basename);
	std::ofstream out(output_file, std::ios::binary);

	out.write(reinterpret_cast<const char*>(&meta.header), sizeof(uint64_t));
	out.write(reinterpret_cast<const char*>(&meta.version), sizeof(int8_t));
	out.write(reinterpret_cast<const char*>(&meta.format), sizeof(int8_t));

	int8_t isBE = opt.enable_be;
	int8_t palette_format;
	if (opt.palette_format == "rgb16") palette_format = 1;
	if (opt.palette_format == "rgb256") palette_format = 2;

	out.write(reinterpret_cast<const char*>(&isBE), sizeof(int8_t));
	out.write(reinterpret_cast<const char*>(&palette_format), sizeof(int8_t));

	if (!opt.enable_be)
	{
		out.write(reinterpret_cast<const char*>(&meta.width), sizeof(int16_t));
		out.write(reinterpret_cast<const char*>(&meta.height), sizeof(int16_t));
		out.write(reinterpret_cast<const char*>(&meta.palette_count), sizeof(int16_t));
		out.write(reinterpret_cast<const char*>(&meta.original_size), sizeof(int32_t));
		out.write(reinterpret_cast<const char*>(&meta.compress_size), sizeof(int32_t));
	}
	else
	{
		uint16_t be_width = to_be_int16(meta.width);
		uint16_t be_height = to_be_int16(meta.height);
		uint16_t be_palette_count = to_be_int16(meta.palette_count);
		uint32_t be_original_size = to_be_int32(meta.original_size);
		uint32_t be_compress_size = to_be_int32(meta.compress_size);

		out.write(reinterpret_cast<const char*>(&be_width), sizeof(uint16_t));
		out.write(reinterpret_cast<const char*>(&be_height), sizeof(uint16_t));
		out.write(reinterpret_cast<const char*>(&be_palette_count), sizeof(uint16_t));
		out.write(reinterpret_cast<const char*>(&be_original_size), sizeof(uint32_t));
		out.write(reinterpret_cast<const char*>(&be_compress_size), sizeof(uint32_t));
	}

	if (!palette.empty())
	{
		if (palette_format == 1)
		{
			uint16_t pal16[meta.palette_count];
			for (int i=0; i<meta.palette_count; i++)
				pal16[i] = RGB255_RGB16(palette[i].r, palette[i].g, palette[i].b, palette[i].a);

			out.write(reinterpret_cast<const char*>(pal16), sizeof(uint16_t)*meta.palette_count);
		}
		else if (palette_format == 2)
		{
			uint32_t pal256[meta.palette_count];
			for (int i=0; i<meta.palette_count; i++)
			{
				pal256[i] = (static_cast<uint32_t>(palette[i].a) << 24) |  // A
										(static_cast<uint32_t>(palette[i].r) << 16) |  // R
										(static_cast<uint32_t>(palette[i].g) << 8)  |  // G
										 static_cast<uint32_t>(palette[i].b);          // B
				
				if (opt.enable_be) pal256[i] = to_be_int32(pal256[i]);
			}

			out.write(reinterpret_cast<const char*>(pal256), sizeof(uint32_t)*meta.palette_count);
		}
	}

	out.write(reinterpret_cast<const char*>(image_buffer), meta.compress_size);
	out.close();
}