#include <stdint.h>
#include <stdio.h>
#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

static int grayscale(char **argv);
static int indexed(char **argv);

int
main(int argc, char **argv) {
	if (argc != 3 && argc != 4) {
		fprintf(stderr, "usage: %s name image.png [palette.png]\n",
		        argv[0]);
		return 1;
	}

	if (argc == 3)
		return grayscale(argv);
	return indexed(argv);
}

static int
grayscale(char **argv)
{
	int img_w, img_h, img_col_chans;
	uint8_t *const img_data =
	    stbi_load(argv[2], &img_w, &img_h, &img_col_chans, 1);

	if (img_data == NULL) {
		fprintf(stderr, "error: failed to load image\n");
		return 1;
	}
	if (img_col_chans != 1) {
		fprintf(stderr, "error: image isn't grayscale\n");
		stbi_image_free(img_data);
		return 1;
	}
	if (img_h * img_w == 0) {
		fprintf(stderr, "error: image size is 0\n");
		stbi_image_free(img_data);
		return 1;
	}

	printf("#include \"px.h\"\nconst PxSpr %s={%d,%d,(uint8_t[]){",
	       argv[1], img_w, img_h);
	for (int i = 0; i < img_w * img_h; i++)
		printf("%u,", (unsigned)img_data[i]);
	printf("}};\n");

	stbi_image_free(img_data);
	return 0;
}

static int
indexed(char **argv)
{
	/* read palette */
	int pal_w, pal_h, pal_col_chans;
	uint8_t *const pal_data =
	    stbi_load(argv[3], &pal_w, &pal_h, &pal_col_chans, 3);

	if (pal_data == NULL) {
		fprintf(stderr, "error: failed to load palette\n");
		return 1;
	}
	if (pal_col_chans != 3) {
		fprintf(stderr, "error: palette isn't RGB\n");
		stbi_image_free(pal_data);
		return 1;
	}
	if (pal_w * pal_h == 0) {
		fprintf(stderr, "error: palette size is 0\n");
		stbi_image_free(pal_data);
		return 1;
	}
	if (pal_w * pal_h >= 256) {
		fprintf(stderr, "error: palette size is greater than 255\n");
		stbi_image_free(pal_data);
		return 1;
	}

	/* store unpacked palette */
	const int pal_size = pal_w * pal_h;
	uint8_t pal[256 * 3] = {0};
	for (int i = 0; i < pal_size * 3; i++)
		pal[i] = pal_data[i];
	stbi_image_free(pal_data);

	/* read image */
	int img_w, img_h, img_col_chans;
	uint8_t *const img_data =
	    stbi_load(argv[2], &img_w, &img_h, &img_col_chans, 4);

	if (img_data == NULL) {
		fprintf(stderr, "error: failed to load image\n");
		return 1;
	}
	if (img_col_chans != 4) {
		fprintf(stderr, "error: image isn't RBGA\n");
		stbi_image_free(img_data);
		return 1;
	}
	if (img_h * img_w == 0) {
		fprintf(stderr, "error: image size is 0\n");
		stbi_image_free(img_data);
		return 1;
	}

	printf("#include \"px.h\"\nconst PxSpr %s={%d,%d,(uint8_t[]){",
	       argv[1], img_w, img_h);
	for (int i = 0; i < img_w * img_h; i++) {
		int col = -1;
		for (int k = 0; k < pal_size; k++)
			if (memcmp(pal + k * 3, img_data + i * 4, 3) == 0) {
				col = k;
				break;
			}
		if (img_col_chans == 4 && img_data[i * 4 + 3] == 0)
			col = 0;
		if (col == -1) {
			fprintf(stderr,
			        "color at pixel %d:%d not found in palette\n",
			        i % img_w, i / img_w);
			stbi_image_free(img_data);
			return 1;
		}
		printf("%d,", col);
	}
	printf("}};\n");

	stbi_image_free(img_data);
	return 0;
}