/* Copyright (C) 2022  June McEnroe <june@causal.agency>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero 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 Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <SDL_rwops.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <string.h>

static void errx(int eval, const char *fmt, ...) {
	va_list ap;
	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	putc('\n', stderr);
	exit(eval);
}
static void err(int eval, const char *fmt, ...) {
	va_list ap;
	va_start(ap, fmt);
	vfprintf(stderr, fmt, ap);
	va_end(ap);
	fprintf(stderr, ": %s\n", SDL_GetError());
	exit(eval);
}

static Uint32 crcTable(Uint8 n) {
	static Uint32 table[256];
	if (table[1]) return table[n];
	for (int i = 0; i < 256; ++i) {
		table[i] = i;
		for (int j = 0; j < 8; ++j) {
			table[i] = (table[i] >> 1) ^ (table[i] & 1 ? 0xEDB88320 : 0);
		}
	}
	return table[n];
}

static Uint32 pngCRC;
static void pngWrite(SDL_RWops *rw, const Uint8 *ptr, Uint32 len) {
	if (!SDL_RWwrite(rw, ptr, len, 1)) err(1, "pngWrite");
	for (Uint32 i = 0; i < len; ++i) {
		pngCRC = crcTable(pngCRC ^ ptr[i]) ^ (pngCRC >> 8);
	}
}
static void pngUint32(SDL_RWops *rw, Uint32 n) {
	Uint8 bytes[4] = { n >> 24, n >> 16, n >> 8, n  };
	pngWrite(rw, bytes, sizeof(bytes));
}
static void pngChunk(SDL_RWops *rw, const char *type, Uint32 len) {
	pngUint32(rw, len);
	pngCRC = ~0;
	pngWrite(rw, (const Uint8 *)type, 4);
}

// Write out image data as DEFLATE literal blocks
static void pngData(SDL_RWops *rw, const Uint8 *data, Uint32 len) {
	Uint32 adler1 = 1, adler2 = 0;
	for (Uint32 i = 0; i < len; ++i) {
		adler1 = (adler1 + data[i]) % 65521;
		adler2 = (adler1 + adler2) % 65521;
	}
	Uint32 zlen = 2 + 5 * ((len + 0xFFFE) / 0xFFFF) + len + 4;
	pngChunk(rw, "IDAT", zlen);
	Uint8 zlib[2] = { 0x08, 0x1D };
	pngWrite(rw, zlib, sizeof(zlib));
	for (; len > 0xFFFF; data += 0xFFFF, len -= 0xFFFF) {
		Uint8 block[5] = { 0x00, 0xFF, 0xFF, 0x00, 0x00 };
		pngWrite(rw, block, sizeof(block));
		pngWrite(rw, data, 0xFFFF);
	}
	Uint8 block[5] = { 0x01, len, len >> 8, ~len, ~len >> 8 };
	pngWrite(rw, block, sizeof(block));
	pngWrite(rw, data, len);
	pngUint32(rw, adler2 << 16 | adler1);
	pngUint32(rw, ~pngCRC);
}

int main(int argc, char *argv[]) {
	if (argc < 4) return 1;
	const char *icoPath = argv[1];
	const char *pngPath = argv[3];
	Uint32 scale = strtoul(argv[2], NULL, 10);
	if (!scale) return 1;

	SDL_RWops *ico = SDL_RWFromFile(icoPath, "rb");
	if (!ico) err(1, "%s", icoPath);

	// Read the ICONHEADER
	// https://devblogs.microsoft.com/oldnewthing/20101018-00/?p=12513
	SDL_ReadLE16(ico); // reserved
	if (SDL_ReadLE16(ico) != 0x0001) {
		errx(1, "%s: not an ICO", icoPath);
	}
	Uint16 count = SDL_ReadLE16(ico);

	// Find the best icon
	Uint32 bestScore = 0;
	struct Icon {
		Uint8 width;
		Uint8 height;
		Uint8 colors;
		Uint16 planes;
		Uint16 bpp;
		Uint32 length;
		Uint32 offset;
	} bestIcon = {0};
	for (Uint16 i = 0; i < count; ++i) {
		struct Icon icon = {0};
		icon.width = SDL_ReadU8(ico);
		icon.height = SDL_ReadU8(ico);
		icon.colors = SDL_ReadU8(ico);
		SDL_ReadU8(ico); // reserved
		icon.planes = SDL_ReadLE16(ico);
		icon.bpp = SDL_ReadLE16(ico); // some ICOs lie about this...
		icon.length = SDL_ReadLE32(ico);
		icon.offset = SDL_ReadLE32(ico);
		Uint32 score = icon.width * icon.height * icon.colors;
		if (score > bestScore) {
			bestScore = score;
			bestIcon = icon;
		}
	}
	if (!bestScore) errx(1, "%s: no icons found", icoPath);
	if (bestIcon.planes != 1) {
		errx(1, "%s: non-indexed bitmap not implemented", icoPath);
	}

	SDL_RWops *png = SDL_RWFromFile(pngPath, "wb");
	if (!png) err(1, "%s", pngPath);

	// Write the PNG header
	SDL_RWwrite(png, "\x89PNG\r\n\x1A\n", 8, 1);
	pngChunk(png, "IHDR", 13);
	pngUint32(png, bestIcon.width * scale);
	pngUint32(png, bestIcon.height * scale);
	Uint8 depth = 8;
	Uint8 color = 3; // indexed
	Uint8 zero[3] = {0};
	pngWrite(png, &depth, 1);
	pngWrite(png, &color, 1);
	pngWrite(png, zero, 3);
	pngUint32(png, ~pngCRC);

	// Read the real bpp and palette count from the DIB header
	SDL_RWseek(ico, bestIcon.offset, RW_SEEK_SET);
	Uint32 dibHeaderLength = SDL_ReadLE32(ico);
	if (dibHeaderLength != 0x28) {
		errx(
			1, "%s: unrecognized DIB header length %u",
			icoPath, dibHeaderLength
		);
	}
	SDL_RWseek(ico, 0x0A, RW_SEEK_CUR);
	bestIcon.bpp = SDL_ReadLE16(ico); // the truth
	SDL_RWseek(ico, 0x10, RW_SEEK_CUR);
	Uint32 paletteCount = SDL_ReadLE32(ico);
	if (!paletteCount) paletteCount = 1 << bestIcon.bpp;

	// Copy the palette, reserving index 0 for transparency. We assume at most
	// 255 colors will be used by the icon.
	Uint8 *palette = calloc(1 + paletteCount, 3);
	if (!palette) errx(1, "calloc failed");
	SDL_RWseek(ico, bestIcon.offset + dibHeaderLength, RW_SEEK_SET);
	for (Uint32 i = 0; i < paletteCount; ++i) {
		palette[3 * (i + 1) + 2] = SDL_ReadU8(ico);
		palette[3 * (i + 1) + 1] = SDL_ReadU8(ico);
		palette[3 * (i + 1) + 0] = SDL_ReadU8(ico);
		SDL_ReadU8(ico);
	}
	pngChunk(png, "PLTE", 3 * (1 + paletteCount));
	pngWrite(png, palette, 3 * (1 + paletteCount));
	pngUint32(png, ~pngCRC);
	pngChunk(png, "tRNS", 1);
	pngWrite(png, zero, 1);
	pngUint32(png, ~pngCRC);
	free(palette);

	Uint32 pngStride = 1 + bestIcon.width;
	Uint8 *lines = calloc(bestIcon.height, pngStride);
	if (!lines) errx(1, "calloc failed");

	// Copy and expand the icon pixels to 8 bpp, adding 1 for the transparency
	// index.
	Uint32 icoPPByte = 8 / bestIcon.bpp;
	Uint32 icoStride = bestIcon.width / icoPPByte;
	icoStride += -icoStride & 0x03;
	for (Uint32 y = bestIcon.height - 1; y < bestIcon.height; --y) {
		Uint8 *ptr = &lines[y * pngStride + 1];
		Uint8 *end = &lines[(y + 1) * pngStride];
		for (Uint32 x = 0; x < icoStride; ++x) {
			Uint8 byte = SDL_ReadU8(ico);
			for (Uint8 i = icoPPByte - 1; i < icoPPByte; --i) {
				if (ptr == end) break;
				*ptr = (byte >> (i * bestIcon.bpp)) & ((1 << bestIcon.bpp) - 1);
				*ptr++ += 1;
			}
		}
	}

	// Read the 1 bpp transparency mask, setting PNG pixels to index 0.
	// https://devblogs.microsoft.com/oldnewthing/20101019-00/?p=12503
	icoPPByte = 8;
	icoStride = bestIcon.width / icoPPByte;
	icoStride += -icoStride & 0x03;
	for (Uint32 y = bestIcon.height - 1; y < bestIcon.height; --y) {
		Uint8 *ptr = &lines[y * pngStride + 1];
		Uint8 *end = &lines[(y + 1) * pngStride];
		for (Uint32 x = 0; x < icoStride; ++x) {
			Uint8 byte = SDL_ReadU8(ico);
			for (Uint8 i = icoPPByte - 1; i < icoPPByte; --i) {
				if (ptr == end) break;
				if ((byte >> i) & 1) *ptr = 0;
				ptr++;
			}
		}
	}

	// Upscale with nearest neighbor
	Uint32 scaledWidth = bestIcon.width * scale;
	Uint32 scaledHeight = bestIcon.height * scale;
	Uint32 scaledStride = 1 + scaledWidth;
	Uint8 *scaled = calloc(scaledHeight, scaledStride);
	for (Uint32 y = 0; y < scaledHeight; ++y) {
		for (Uint32 x = 0; x < scaledWidth; ++x) {
			scaled[y * scaledStride + 1 + x] =
				lines[(y / scale) * pngStride + 1 + (x / scale)];
		}
	}
	pngData(png, scaled, scaledHeight * scaledStride);
	pngChunk(png, "IEND", 0);
	pngUint32(png, ~pngCRC);

	free(scaled);
	free(lines);
	SDL_RWclose(png);
	SDL_RWclose(ico);
}