#include <stdio.h>
#include <stdlib.h>
#include <math.h>

#include "plat.h"
#include "str.h"
#include "vid_enums.h"

typedef struct {
	float r, g, b;
} vec3;

typedef struct {
	unsigned char r, g, b, a;
} Colour;

typedef struct {
	Colour* pixels;
	int w, h;
} Image;

typedef struct {
	vec3 start, end;
	unsigned char indices[4 * 4];
} Block;

typedef struct {
	float start, end;
	unsigned char indices[4 * 4];
} Block_BC4;

const char* format_names[] = {
#define x(n) #n,
	texture_format_xmacro()
#undef x
};

unsigned encode_endpoint(const vec3* v) {
	return
		((unsigned)(v->r * 31.9999f) << 11) |
		((unsigned)(v->g * 63.9999f) << 5) |
		((unsigned)(v->b * 31.9999f));
}

unsigned encode_endpoint_bc4(float v) {
	return (unsigned)(v * 255.0f);
}

void get_block(Block* block, Colour* pixels, int stride) {
	int x, y, i;
	vec3 cols[4 * 4];
	vec3 palette[4];
	vec3 s, e;
	unsigned es, ee;
	float cd;
	for (y = 0; y < 4; y++) {
		for (x = 0; x < 4; x++) {
			Colour pixel = pixels[x + y * stride];
			vec3 vec;
			vec.r = (float)pixel.r / 255.0f;
			vec.g = (float)pixel.g / 255.0f;
			vec.b = (float)pixel.b / 255.0f;
			cols[x + y * 4] = vec;
		}
	}
	cd = -1.0f;
	for (i = 0; i < 4 * 4; i++) {
		int j;
		vec3* a = &cols[i];
		for (j = i + 1; j < 4 * 4; j++) {
			vec3* b = &cols[j];
			float d =
				(b->r - a->r) * (b->r - a->r) +
				(b->g - a->g) * (b->g - a->g) +
				(b->b - a->b) * (b->b - a->b);
			if (d > cd) {
				cd = d;
				s = *a;
				e = *b;
			}
		}
	}
	es = encode_endpoint(&s);
	ee = encode_endpoint(&e);
	if (es == ee) {
		block->start = s;
		block->end.r = 0.0f;
		block->end.g = 0.0f;
		block->end.b = 0.0f;
		for (i = 0; i < 16; i++) {
			block->indices[i] = 0;
		}
		return;
	}
	if (es < ee) {
		/* todo alpha */
		vec3 t = s;
		s = e;
		e = t;
	}
	palette[0] = s;
	palette[1] = e;
	palette[2].r = 0.666666f*palette[0].r + 0.333333f*palette[1].r;
	palette[2].g = 0.666666f*palette[0].g + 0.333333f*palette[1].g;
	palette[2].b = 0.666666f*palette[0].b + 0.333333f*palette[1].b;
	palette[3].r = 0.333333f*palette[0].r + 0.666666f*palette[1].r;
	palette[3].g = 0.333333f*palette[0].g + 0.666666f*palette[1].g;
	palette[3].b = 0.333333f*palette[0].b + 0.666666f*palette[1].b;
	for (i = 0; i < 4 * 4; i++) {
		vec3* col = &cols[i];
		int closest = 0, j;
		float cd = 99999.0f;
		for (j = 0; j < 4; j++) {
			vec3* p = &palette[j];
			float d =
				(p->r - col->r) * (p->r - col->r) +
				(p->g - col->g) * (p->g - col->g) +
				(p->b - col->b) * (p->b - col->b);
			if (d < cd) {
				cd = d;
				closest = j;
			}
		}
		block->indices[i] = closest;
	}
	block->start = s;
	block->end = e;
}

void compress_block(const Block* block, FILE* f) {
	int i;
	unsigned start = encode_endpoint(&block->start);
	unsigned end   = encode_endpoint(&block->end);
	unsigned indices = 0;
	fwrite(&start, 1, 2, f);
	fwrite(&end,   1, 2, f);
	for (i = 0; i < 4 * 4; i++) {
		indices |= block->indices[i] << (i * 2);
	}
	fwrite(&indices, 1, 4, f);
}

void get_block_bc4(
	Block_BC4* block,
	Colour* pixels,
	int stride,
	int channel
) {
	int x, y;
	float palette[8];
	palette[0] = 0.08f;
	palette[1] = 0.92f;
	float biggest = -1.0f;
	float smallest = 1.0f;
	for (y = 0; y < 4; y++) {
		for (x = 0; x < 4; x++) {
			uint8_t* pixel = (uint8_t*)&pixels[x + y * stride];
			float v = (float)pixel[channel] / 255.0f;
			if (v > biggest)  biggest = v;
			if (v < smallest) smallest = v;
			if (v > palette[0]) palette[0] = v;
			if (v < palette[1]) palette[1] = v;
		}
	}
	if (biggest == smallest) {
		int i, idx = 1;
		block->end = biggest;
		if (palette[0] > 0.9f) {
			idx = 0;
			block->start = smallest;
			block->end = 0.0f;
		} else
			block->start = 1.0f;
		assert(block->start > block->end);
		for (i = 0; i < 16; i++)
			block->indices[i] = idx;
	}
	if (biggest > 0.92f && smallest < 0.08f) {
		float t = palette[1];
		palette[1] = palette[0];
		palette[0] = t;
		palette[2] = (4.0f * palette[0] + 1.0f * palette[1]) / 5.0f;
		palette[3] = (3.0f * palette[0] + 2.0f * palette[1]) / 5.0f;
		palette[4] = (2.0f * palette[0] + 3.0f * palette[1]) / 5.0f;
		palette[5] = (1.0f * palette[0] + 4.0f * palette[1]) / 5.0f;
		palette[6] = 0.0f;
		palette[7] = 1.0f;
	} else {
		palette[2] = (6.0f * palette[0] + 1.0f * palette[1]) / 7.0f;
		palette[3] = (5.0f * palette[0] + 2.0f * palette[1]) / 7.0f;
		palette[4] = (4.0f * palette[0] + 3.0f * palette[1]) / 7.0f;
		palette[5] = (3.0f * palette[0] + 4.0f * palette[1]) / 7.0f;
		palette[6] = (2.0f * palette[0] + 5.0f * palette[1]) / 7.0f;
		palette[7] = (1.0f * palette[0] + 6.0f * palette[1]) / 7.0f;
	}
	block->start = palette[0];
	block->end = palette[1];
	for (y = 0; y < 4; y++) {
		for (x = 0; x < 4; x++) {
			uint8_t* pixel = (uint8_t*)&pixels[x + y * stride];
			float v = (float)pixel[channel] / 255.0f;
			float d = 100.0f;
			int i;
			for (i = 0; i < 8; i++) {
				float d2 = (v - palette[i]);
				d2 = d2 < 0.0f? -d2: d2;
				if (d2 < d) {
					d = d2;
					block->indices[x + y * 4] = i;
				}
			}
		}
	}
}

void compress_block_bc4(const Block_BC4* block, FILE* f) {
	int i;
	unsigned start = encode_endpoint_bc4(block->start);
	unsigned end   = encode_endpoint_bc4(block->end);
	uint64_t indices = 0;
	fwrite(&start, 1, 1, f);
	fwrite(&end,   1, 1, f);
	for (i = 0; i < 4 * 4; i++) {
		indices |= ((uint64_t)block->indices[i]) << (i * 3);
	} /* 48 bits of indices */
	fwrite(&indices, 1, 6, f);
}

void compress_bc1(Colour* pixels, int w, int h, FILE* f) {
	int x, y, cw, ch;
	Block block;
	cw = w / 4;
	ch = h / 4;
	for (y = 0; y < ch; y++) {
		for (x = 0; x < cw; x++) {
			get_block(&block, &pixels[x * 4 + y * 4 * w], w);
			compress_block(&block, f);
		}
	}
}

void compress_bc4(Colour* pixels, int w, int h, FILE* f) {
	int x, y, cw = w / 4, ch = h / 4;
	Block_BC4 block;
	for (y = 0; y < ch; y++) {
		for (x = 0; x < cw; x++) {
			get_block_bc4(&block, &pixels[x * 4 + y * 4 * w], w, 0);
			compress_block_bc4(&block, f);
		}
	}
}

void compress_bc5(Colour* pixels, int w, int h, FILE* f) {
	int x, y, cw = w / 4, ch = h / 4;
	Block_BC4 a, b;
	for (y = 0; y < ch; y++) {
		for (x = 0; x < cw; x++) {
			Colour* p = &pixels[x * 4 + y * 4 * w];
			get_block_bc4(&a, p, w, 0);
			get_block_bc4(&b, p, w, 1);
			compress_block_bc4(&a, f);
			compress_block_bc4(&b, f);
		}
	}
}

void write_header(FILE* f, int w, int h, Texture_Format fmt) {
	fwrite("TXTR", 1, 4, f);
	fwrite(&w, 1, 4, f);
	fwrite(&h, 1, 4, f);
	fwrite(&fmt, 1, 4, f);
}

void convert(Image* image, Texture_Format target, FILE* f) {
	write_header(f, image->w, image->h, target);
	switch (target) {
		case texture_format_bc1:
			compress_bc1(image->pixels, image->w, image->h, f);
			break;
		case texture_format_bc4:
			compress_bc4(image->pixels, image->w, image->h, f);
			break;
		case texture_format_bc5:
			compress_bc5(image->pixels, image->w, image->h, f);
			break;
		default:
			print_err("Unsupported target format.\n");
			pbreak(40);
	}
}

Texture_Format texture_format_from_string(const char* s) {
	int i, e = sizeof format_names / sizeof *format_names;
	for (i = 0; i < e; i++)
		if (string_equal(s, format_names[i]))
			return (Texture_Format)i;
	print_err("Invalid texture format %s\n", s);
	pbreak(10);
	return (Texture_Format)0;
}

int proc_bitmap(FILE* infile, FILE* outfile, Texture_Format target) {
	unsigned bmp_offset;
	int bmp_w, bmp_h, s, x, y;
	unsigned short bmp_bits;
	Colour* buffer;
	Image img;
	Colour pixel;

	fseek(infile, 10, SEEK_SET);
	fread(&bmp_offset, 4, 1, infile);
	fseek(infile, 18, SEEK_SET);
	fread(&bmp_w,    4, 1, infile);
	fread(&bmp_h,    4, 1, infile);
	fseek(infile, 28, SEEK_SET);
	fread(&bmp_bits, 2, 1, infile);

	if (bmp_bits != 32) {
		print_err("Bitmap must have 32 bit pixels. Instead has %d bit pixels.\n", bmp_bits);
		return 4;
	}

	if (bmp_w % 4 != 0 || bmp_h % 4 != 0) {
		print_err("Bitmap must have a size divisible by four.\n");
		return 5;
	}

	fseek(infile, bmp_offset, SEEK_SET);
	s = bmp_w * bmp_h;
	buffer = malloc(s * 4);
	/* Flip & read. */
	for (y = 0; y < bmp_h; y++) {
		for (x = 0; x < bmp_w; x++) {
			unsigned char t;
			fread(&pixel, 1, 4, infile);
			t = pixel.r;
			pixel.r = pixel.b;
			pixel.b = t;
			if (pixel.a == 0) {
				pixel.r = 0;
				pixel.g = 0;
				pixel.b = 0;
			}
			buffer[x + (bmp_h - y - 1) * bmp_w] = pixel;
		}
	}
	img.pixels = buffer;
	img.w = bmp_w;
	img.h = bmp_h;
	convert(&img, target, outfile);
	fclose(outfile);
	return 0;
}

typedef struct {
	float r, g, b;
} Hdr_Pixel;

Hdr_Pixel hdr_pixel(uint8_t* bytes) {
	Hdr_Pixel p;
	float f = ldexp(1.0f, (int)bytes[3] - (128 + 8));
	p.r = ((float)bytes[0] + 0.5f) * f;
	p.g = ((float)bytes[1] + 0.5f) * f;
	p.b = ((float)bytes[2] + 0.5f) * f;
	return p;
}

/* from open xr */
uint16_t f2h(uint32_t i) {
	int s =  (i >> 16) & 0x00008000;
	int e = ((i >> 23) & 0x000000ff) - (127 - 15);
	int m =   i        & 0x007fffff;
	if (e <= 0)
	{
		if (e < -10)
			return s;
		m = m | 0x00800000;
		int t = 14 - e;
		int a = (1 << (t - 1)) - 1;
		int b = (m >> t) & 1;
		m = (m + a + b) >> t;
		return s | m;
	}
	else if (e == 0xff - (127 - 15))
	{
		if (m == 0)
			return s | 0x7c00;
		else
		{
			m >>= 13;
			return s | 0x7c00 | m | (m == 0);
		}
	}
	else
	{
		m = m + 0x00000fff + ((m >> 13) & 1);
		if (m & 0x00800000)
		{
			m =  0;
			e += 1;
		}
		if (e > 30)
			return s | 0x7c00;
		return s | (e << 10) | (m >> 13);
	}
}

void write_halves(FILE* f, int c, float* pixels) {
	int i;
	for (i = 0; i < c; i++) {
		uint16_t h = f2h(*(uint32_t*)&pixels[i]);
		fwrite(&h, 2, 1, f);
	}
}

int proc_hdr(FILE* infile, FILE* outfile, int half) {
	char buf[256];
	char* tok;
	int w, h, i;
	Hdr_Pixel* pixels;
	uint8_t* row;
	int valid = 0;
	fseek(infile, 0, SEEK_SET);
	while (fgets(buf, sizeof buf, infile)) {
		if (string_equal(buf, "FORMAT=32-bit_rle_rgbe\n"))
			valid = 1;
		if (buf[0] == '\n') break;
	}
	if (!valid) {
		print_err("Format unsupported.\n");
		return 10;
	}
	tok = fgets(buf, sizeof buf, infile);
	if (tok[0] != '-' || tok[1] != 'Y') {
		print_err("Format unsupported.\n");
		return 11;
	}
	for (tok = tok + 2; *tok == ' '; tok++);
	h = (int)strtol(tok, &tok, 10);
	for (; *tok == ' '; tok++);
	if (tok[0] != '+' || tok[1] != 'X') {
		print_err("Format unsupported.\n");
		return 11;
	}
	for (tok = tok + 2; *tok == ' '; tok++);
	w = (int)strtol(tok, &tok, 10);
	if (w < 8 || w > 0x7ffff) {
		/* todo */
		print_err("Format unsupported.\n");
		return 12;
	}
	pixels = malloc(w * h * sizeof *pixels);
	row = malloc(w * 4);
	for (i = 0; i < h; i++) {
		int a, b, c, d, j;
		a = fgetc(infile);
		b = fgetc(infile);
		c = fgetc(infile);
		d = fgetc(infile);
		if (d == EOF) {
			print_err("Unexpected end of image file.\n");
			return 13;
		}
		if (b != 2 || (c & 0x80)) {
			int j, shift = 0;
			row[0] = a;
			row[1] = b;
			row[2] = c;
			row[3] = d;
			for (j = 1; j < w; j++) {
				uint8_t* p = &row[j * 4];
				p[0] = fgetc(infile);
				p[1] = fgetc(infile);
				p[2] = fgetc(infile);
				p[3] = fgetc(infile);
				if (
					p[0] == 1 &&
					p[1] == 1 &&
					p[2] == 1
				) {
					int k;
					for (k = p[3] << shift; k > 0; k--) {
						p[0] = p[-4];
						p[1] = p[-3];
						p[2] = p[-2];
						p[3] = p[-1];
						p += 4;
						j++;
					}
					shift += 8;
				} else
					shift = 0;
			}
			goto convert;
		}
		if (((c << 8) | d) != w) {
			print("%d, %d\n", i, ((c << 8) | d));
			print_err("Corrupt run.\n");
			return 14;
		}
		for (j = 0; j < 4; j++) {
			int k;
			for (k = 0; k < w;) {
				int d1 = fgetc(infile), d2;
				if (d1 == EOF) {
					print_err("Unexpected end of image file.\n");
					return 14;
				}
				if (d1 > 128) {
					d1 &= 127;
					d2 = fgetc(infile);
					for (; d1--; k++)
						row[k * 4 + j] = (uint8_t)d2;
				} else
					for (; d1--; k++)
						row[k * 4 + j] = (uint8_t)fgetc(infile);
			}
		}
		convert:
		for (j = 0; j < w; j++) {
			pixels[j + i * w] = hdr_pixel(&row[j * 4]);
		}
	}
	if (half) {
		write_header(outfile, w, h, texture_format_rgb16f);
		write_halves(outfile, w * h * 3, (float*)pixels);
	} else {
		write_header(outfile, w, h, texture_format_rgb32f);
		fwrite(pixels, sizeof *pixels, w * h, outfile);
	}
	fclose(outfile);
	return 0;
}

int main(int argc, const char** argv) {
	FILE* outfile, * infile;
	char magic[11];
	Texture_Format target;

	if (argc < 4) {
		print_err("Usage: %s infile outfile format.\n", argv[0]);
		return 1;
	}

	infile = fopen(argv[1], "rb");
	if (!infile) {
		print_err("Failed to open %s.\n", argv[1]);
		return 2;
	}
	target = texture_format_from_string(argv[3]);

	outfile = fopen(argv[2], "wb");
	if (!outfile) {
		print_err("Failed to open %s.\n", argv[2]);
		return 6;
	}

	fread(magic, 1, 10, infile); magic[10] = 0;
	if (magic[0] == 'B' && magic[1] == 'M') {
		return proc_bitmap(infile, outfile, target);
	}
	if (string_equal(magic, "#?RADIANCE")) {
		int half = target == texture_format_rgb16f;
		if (!half && target != texture_format_rgb32f) {
			print_err("Unsupported target format.\n");
			return 40;
		}
		return proc_hdr(infile, outfile, half);
	}
	print_err("Invalid image file.\n");
	return 80;
}