#include <iostream>
#include <iomanip>
#include <string>

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "stb_image_write.h"

#include "main_colour_extractor.h"
#include "complimentary_colour_extractor.h"

using namespace std;

void print_pixel(ColourRGB rgb)
{
	uint8_t r = static_cast<uint8_t>(round(rgb.r * 255));
	uint8_t g = static_cast<uint8_t>(round(rgb.g * 255));
	uint8_t b = static_cast<uint8_t>(round(rgb.b * 255));

	cout << "#" << std::hex << (int)r << (int)g << (int)b;
	cout << std::dec << " - " << (int)r << " " << (int)g << " " << (int)b;
}

void usage(char *name)
{
	cout << "Usage: " << name << " [input image]" << endl;
	cout << "Supported image formats: JPG, PNG, TGA, BMP, GIF" << endl;
	cout << "Image must have 3 or more channels" << endl;
}

int main(int argc, char **argv)
{
	int x_size;
	int y_size;
	int channels;
	int error;
	unsigned char *image = nullptr;
	ColourRGB pixel;
	ColourRGB pixel2;
	ColourHSL pixel_hsl;
	MainColourExtractor main_extractor(36, 10, 10);
	ComplimentaryColourExtractor comp_extractor({0, 0, 0}, 36, 10, 10);

	if (argc < 2)
	{
		cout << "Not enough arguments" << endl;
		usage(argv[0]);
		return -1;
	}

	error = stbi_info(argv[1], &x_size, &y_size, &channels);
	if (1 != error)
	{
		cout << "Image type is not supported" << endl;
		usage(argv[0]);
		return -1;
	}

	if (channels < 3)
	{
		cout << "This image is in grayscale (it has no colour)" << endl;
		usage(argv[0]);
		return -1;
	}

	image = stbi_load(argv[1], &x_size, &y_size, &channels, 0);
	if (nullptr == image)
	{
		cout << "Could not open image" << endl;
		return -1;
	}

	for (int i = 0; i < x_size * y_size * channels; i += channels)
	{
		pixel.r = image[i + 0];
		pixel.g = image[i + 1];
		pixel.b = image[i + 2];

		pixel.r /= 255;
		pixel.g /= 255;
		pixel.b /= 255;

		if (channels > 3)
		{
			// TODO: Handle alpha here
		}

		pixel_hsl = RGBtoHSL(pixel);
		main_extractor.add_pixel(pixel_hsl);
		comp_extractor.add_pixel(pixel_hsl);
	}
	stbi_image_free(image);

	pixel_hsl = main_extractor.extract_colour();
	comp_extractor.set_main_colour(pixel_hsl);
	pixel = HSLtoRGB(pixel_hsl);
	pixel2 = HSLtoRGB(comp_extractor.extract_colour());
	
	cout << "Main - ";
	print_pixel(pixel);
	cout << endl;
	cout << "Comp - ";
	print_pixel(pixel2);
	cout << endl;

	x_size = 16;
	y_size = 16;
	channels = 3;
	image = new unsigned char[x_size * y_size * channels];
	if (nullptr == image)
	{
		cout << "Could not allocate space for image" << endl;
		return -1;
	}

	for (int i = 0; i < x_size * y_size; ++i)
	{
		image[i * channels + 0] = static_cast<unsigned char>(round(pixel.r * 255));
		image[i * channels + 1] = static_cast<unsigned char>(round(pixel.g * 255));
		image[i * channels + 2] = static_cast<unsigned char>(round(pixel.b * 255));
	}

	stbi_write_png("0 - main.png", x_size, y_size, channels, image, x_size * channels);

	for (int i = 0; i < x_size * y_size; ++i)
	{
		image[i * channels + 0] = static_cast<unsigned char>(round(pixel2.r * 255));
		image[i * channels + 1] = static_cast<unsigned char>(round(pixel2.g * 255));
		image[i * channels + 2] = static_cast<unsigned char>(round(pixel2.b * 255));
	}

	stbi_write_png("1 - comp.png", x_size, y_size, channels, image, x_size * channels);

	delete[] image;

	return 0;
}