#include <libgen.h>
#include <png.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include "libpdtex/pdtex.h"

/**
 * tex2png - Convert a PD texture to PNG images.
 *
 * Each PD texture, herein called a pdtex, can contain up to 7 images at
 * different levels of detail (LOD). Each image can use a different pixel format
 * and compression algorithm.
 *
 * tex2png is called with two arguments:
 * 1) a path to a single pdtex, which is a binary slice of the PD ROM, and
 * 2) a path to a directory where the PNGs will be created.
 *
 * PNG filenames are generated by combining the original filename with the LOD.
 * For example, if the pdtex's basename is foo.bin then the PNG files may be
 * named foo-0.png, foo-1.png and so on. If a pdtex doesn't contain an image for
 * a particular LOD then no PNG will be created for that LOD.
 *
 * The N64's pixel formats are more granular than what PNG supports, so pixel
 * format are converted where necessary. For example, a pdtex using the 2 byte
 * 5/5/5/1 (R/G/B/A) format is converted to regular 4 byte RGBA for PNG.
 */

static int usage(void)
{
	fprintf(stderr, "Usage: tex2png <infile> <outdir>\n");
	fprintf(stderr, "eg. tex2png 0000.bin pngs/\n");

	return 1;
}

/**
 * Return a filename for writing a LOD PNG.
 *
 * eg. if infile is "foo/0000.bin", dir is "pngs" and lod is 6,
 * return "pngs/0000-6.png"
 */
static void make_lod_filename(char *infile, char *dir, int lod, char *dst)
{
	char *base = basename(infile);
	char *ptr;

	ptr = strrchr(base, '.');

	if (ptr) {
		ptr[0] = '\0';
	}

	snprintf(dst, 1024, "%s/%s-%d.png", dir, base, lod);
}

static int tex_format_to_png_format(enum pd_format pdformat)
{
	switch (pdformat) {
	case PDFORMAT_RGBA32:
	case PDFORMAT_RGBA16:
		return PNG_FORMAT_RGBA;
	case PDFORMAT_RGB24:
	case PDFORMAT_RGB15:
		return PNG_FORMAT_RGB;
		return PNG_FORMAT_RGB;
	case PDFORMAT_IA16:
	case PDFORMAT_IA8:
	case PDFORMAT_IA4:
		return PNG_FORMAT_GA;
	case PDFORMAT_I8:
	case PDFORMAT_I4:
		return PNG_FORMAT_GRAY;
	case PDFORMAT_RGBA16_CI8:
	case PDFORMAT_RGBA16_CI4:
		return PNG_FORMAT_RGBA | PNG_FORMAT_FLAG_COLORMAP;
	case PDFORMAT_IA16_CI8:
		return PNG_FORMAT_GA | PNG_FORMAT_FLAG_COLORMAP;
	case PDFORMAT_IA4_CI4:
		return PNG_FORMAT_GA | PNG_FORMAT_FLAG_COLORMAP;
	}

	fprintf(stderr, "Unrecognised format 0x%02x\n", pdformat);
	return 0;
}

/**
 * Convert RGBA 5/5/5/1 to RGBA 8/8/8/8.
 */
static void rgba16_to_rgba32(uint8_t *in, png_bytep out, int numpixels)
{
	int outpos = 0;
	int i;

	for (i = 0; i < numpixels; i++) {
		uint16_t rgba16 = in[i * 2] << 8 | in[i * 2 + 1];
		uint8_t r = (rgba16 >> 11) & 0x1f;
		uint8_t g = (rgba16 >> 6) & 0x1f;
		uint8_t b = (rgba16 >> 1) & 0x1f;
		out[outpos + 0] = r * 8;
		out[outpos + 1] = g * 8;
		out[outpos + 2] = b * 8;
		out[outpos + 3] = (rgba16 & 1) * 255;
		outpos += 4;
	}
}

/**
 * Convert RGB 5/5/5 to RGB 8/8/8.
 */
static void rgb15_to_rgb24(uint8_t *in, png_bytep out, int numpixels)
{
	int outpos = 0;
	int i;

	for (i = 0; i < numpixels; i++) {
		uint16_t rgb15 = in[i * 2] << 8 | in[i * 2 + 1];
		uint8_t r = (rgb15 >> 11) & 0x1f;
		uint8_t g = (rgb15 >> 6) & 0x1f;
		uint8_t b = (rgb15 >> 1) & 0x1f;
		out[outpos + 0] = r * 8;
		out[outpos + 1] = g * 8;
		out[outpos + 2] = b * 8;
		outpos += 3;
	}
}

/**
 * Convert IA 4/4 to IA 8/8.
 */
static void ia8_to_ia16(uint8_t *in, png_bytep out, int numpixels)
{
	int outpos = 0;
	int i;

	for (i = 0; i < numpixels; i++) {
		out[outpos + 0] = ((in[i] >> 4) & 0xf) * 16;
		out[outpos + 1] = (in[i] & 0xf) * 16;
		outpos += 2;
	}
}

/**
 * Convert IA 3/1 to IA 8/8.
 */
static void ia4_to_ia16(uint8_t *in, png_bytep out, int numpixels)
{
	int outpos = 0;
	int i;

	for (i = 0; i < (numpixels + 1) / 2; i++) {
		out[outpos + 0] = ((in[i] >> 1) & 7) * 32;
		out[outpos + 1] = (in[i] & 1) * 255;
		outpos += 2;
	}
}

/**
 * Convert I 4 to I 8.
 */
static void i4_to_i8(uint8_t *in, png_bytep out, int numpixels)
{
	int i;

	for (i = 0; i < numpixels; i++) {
		out[i] = in[i] * 16;
	}
}

/**
 * Convert a pdtex image to PNG and write it to disk.
 *
 * To do this, build the PNG using libpng and convert the pixel data into a
 * format that's supported by PNG.
 */
static bool write_image(struct pd_image *image, struct pd_tex *tex, char *outfile)
{
	png_image png;
	png_bytep buffer;
	uint8_t *palette = NULL;

	memset(&png, 0, sizeof(png));

	png.version = PNG_IMAGE_VERSION;
	png.width = image->width;
	png.height = image->height;
	png.format = tex_format_to_png_format(image->format);

	if (png.format | PNG_FORMAT_FLAG_COLORMAP) {
		png.colormap_entries = tex->numcolours;
		palette = malloc(PNG_IMAGE_COLORMAP_SIZE(png));
	}

	buffer = malloc(PNG_IMAGE_SIZE(png));

	if (!buffer) {
		fprintf(stderr, "Unable to allocate memory for PNG buffer\n");
		return false;
	}

	switch (image->format) {
	case PDFORMAT_RGBA16:
		rgba16_to_rgba32(image->pixels, buffer, image->width * image->height);
		break;
	case PDFORMAT_RGB15:
		rgb15_to_rgb24(image->pixels, buffer, image->width * image->height);
		break;
	case PDFORMAT_IA8:
		ia8_to_ia16(image->pixels, buffer, image->width * image->height);
		break;
	case PDFORMAT_IA4:
		ia4_to_ia16(image->pixels, buffer, image->width * image->height);
		break;
	case PDFORMAT_I4:
		i4_to_i8(image->pixels, buffer, image->width * image->height);
		break;
	case PDFORMAT_RGBA16_CI8:
		rgba16_to_rgba32(tex->palette, palette, tex->numcolours);
		memcpy(buffer, image->pixels, PNG_IMAGE_SIZE(png));
		break;
	case PDFORMAT_RGBA16_CI4:
		rgba16_to_rgba32(tex->palette, palette, tex->numcolours);
		memcpy(buffer, image->pixels, PNG_IMAGE_SIZE(png));
		break;
	case PDFORMAT_IA16_CI8:
		memcpy(palette, tex->palette, tex->numcolours * 2);
		memcpy(buffer, image->pixels, PNG_IMAGE_SIZE(png));
		break;
	case PDFORMAT_IA4_CI4:
		ia8_to_ia16(tex->palette, palette, tex->numcolours);
		memcpy(buffer, image->pixels, PNG_IMAGE_SIZE(png));
		break;
	default:
		memcpy(buffer, image->pixels, PNG_IMAGE_SIZE(png));
		break;
	}

	if (!png_image_write_to_file(&png, outfile, 0, buffer, 0, palette)) {
		fprintf(stderr, "Unable to write PNG to '%s': %s\n", outfile, png.message);
		return false;
	}

	return true;
}

int main(int argc, char **argv)
{
	if (argc != 3) {
		return usage();
	}

	struct pd_tex *tex = pdtex_allocate();

	if (!tex) {
		fprintf(stderr, "Unable to allocate memory for texture\n");
		return 1;
	}

	if (pdtex_read(tex, argv[1]) != 0) {
		fprintf(stderr, "Unable to read texture at '%s'\n", argv[1]);
		return 1;
	}

	for (int i = 0; i < PDTEX_MAX_IMAGES; i++) {
		if (tex->images[i].exists) {
			char filename[1024];
			make_lod_filename(argv[1], argv[2], i, filename);

			if (!write_image(&tex->images[i], tex, filename)) {
				return 1;
			}
		}
	}

	return 0;
}
