/* 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 <stdarg.h>
#include <stdio.h>
#include <stdlib.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);
}

int main(int argc, char *argv[]) {
	if (argc < 2) return 1;
	const char *exePath = argv[1];
	const char *name = strrchr(exePath, '/');
	name = (name ? &name[1] : exePath);
	char icoPath[256];
	snprintf(
		icoPath, sizeof(icoPath), "%.*s.ICO", (int)strcspn(name, "."), name
	);

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

	// Read the EXE header
	if (SDL_ReadU8(exe) != 'M' || SDL_ReadU8(exe) != 'Z') {
		errx(1, "%s: invalid MZ signature", exePath);
	}
	SDL_RWseek(exe, 0x3C, RW_SEEK_SET);
	Uint16 neOffset = SDL_ReadLE16(exe);
	SDL_RWseek(exe, neOffset, RW_SEEK_SET);
	Uint8 sig[2] = { SDL_ReadU8(exe), SDL_ReadU8(exe) };
	if (sig[0] == 'P' && sig[1] == 'E') {
		errx(1, "%s: not implemented for PE", exePath);
	}
	if (sig[0] != 'N' || sig[1] != 'E') {
		errx(1, "%s: invalid NE/PE signature", exePath);
	}

	// Read the NE header
	SDL_RWseek(exe, neOffset + 0x24, RW_SEEK_SET);
	Uint16 resourceTableOffset = neOffset + SDL_ReadLE16(exe);

	// Find the RT_GROUP_ICON resource
	SDL_RWseek(exe, resourceTableOffset, RW_SEEK_SET);
	Uint16 alignmentShift = SDL_ReadLE16(exe);
	for (;;) {
		Uint16 typeID = SDL_ReadLE16(exe);
		Uint16 count = SDL_ReadLE16(exe);
		SDL_ReadLE32(exe); // reserved
		if (!typeID) {
			errx(1, "%s: no group icon resource", exePath);
		}
		if (typeID == 0x800E) break;
		SDL_RWseek(exe, 0x0C * count, RW_SEEK_CUR);
	}

	// Copy the ICONHEADER
	// https://devblogs.microsoft.com/oldnewthing/20101018-00/?p=12513
	SDL_RWseek(exe, (Uint32)SDL_ReadLE16(exe) << alignmentShift, RW_SEEK_SET);
	SDL_WriteLE16(ico, SDL_ReadLE16(exe)); // reserved
	SDL_WriteLE16(ico, SDL_ReadLE16(exe)); // type
	Uint16 iconDirCount = SDL_ReadLE16(exe);
	SDL_WriteLE16(ico, iconDirCount);

	// Copy ICONDIRENTRY's. The entries in the RT_GROUP_ICON resource differ
	// from those in ICO in their last field. In ICO, it is a 32-bit offset
	// of the icon data in the ICO file. In the resource, it is a 16-bit
	// resource ID of an RT_ICON resource. Here we assume the resources are
	// in the same order as the entries referencing them. We translate the
	// resource IDs into data offsets by simply tallying up the data lengths
	// so far. The data offsets need to be aligned to alignmentShift to match
	// up with the lengths of the icon resources.
	Uint32 iconDirLength = 0x06 + 0x10 * iconDirCount;
	Uint32 dataOffset = iconDirLength;
	dataOffset += -dataOffset & ((1 << alignmentShift) - 1);
	for (Uint16 i = 0; i < iconDirCount; ++i) {
		SDL_WriteU8(ico, SDL_ReadU8(exe)); // width
		SDL_WriteU8(ico, SDL_ReadU8(exe)); // height
		SDL_WriteU8(ico, SDL_ReadU8(exe)); // colors
		SDL_WriteU8(ico, SDL_ReadU8(exe)); // reserved
		SDL_WriteLE16(ico, SDL_ReadLE16(exe)); // color planes
		SDL_WriteLE16(ico, SDL_ReadLE16(exe)); // bits per pixel
		Uint32 dataLength = SDL_ReadLE32(exe);
		SDL_ReadLE16(exe); // resource ID, assumed sequential
		SDL_WriteLE32(ico, dataLength);
		SDL_WriteLE32(ico, dataOffset);
		dataOffset += dataLength;
		dataOffset += -dataOffset & ((1 << alignmentShift) - 1);
	}
	// Padding before the first aligned data offset
	for (Uint8 i = 0; i < (-iconDirLength & ((1 << alignmentShift) - 1)); ++i) {
		SDL_WriteU8(ico, 0);
	}

	// Find the RT_ICON resources
	SDL_RWseek(exe, resourceTableOffset + 0x02, RW_SEEK_SET);
	Uint16 resourceCount;
	for (;;) {
		Uint16 typeID = SDL_ReadLE16(exe);
		resourceCount = SDL_ReadLE16(exe);
		SDL_ReadLE32(exe); // reserved
		if (!typeID) {
			errx(1, "%s: no icon resources", exePath);
		}
		if (typeID == 0x8003) break;
		SDL_RWseek(exe, 0x0C * resourceCount, RW_SEEK_CUR);
	}

	// Copy all icon data sequentially
	Uint8 *data = NULL;
	for (Uint16 i = 0; i < resourceCount; ++i) {
		Uint32 dataOffset = (Uint32)SDL_ReadLE16(exe) << alignmentShift;
		Uint32 dataLength = (Uint32)SDL_ReadLE16(exe) << alignmentShift;
		SDL_ReadLE16(exe); // flags
		SDL_ReadLE16(exe); // resource ID
		SDL_ReadLE32(exe); // reserved
		Sint64 nextResource = SDL_RWtell(exe);
		if (nextResource < 0) err(1, "%s", exePath);

		SDL_RWseek(exe, dataOffset, RW_SEEK_SET);
		data = realloc(data, dataLength);
		if (!data) errx(1, "malloc failed");
		size_t num = SDL_RWread(exe, data, dataLength, 1);
		if (!num) err(1, "%s", exePath);
		num = SDL_RWwrite(ico, data, dataLength, 1);
		if (!num) err(1, "%s", icoPath);

		SDL_RWseek(exe, nextResource, RW_SEEK_SET);
	}
	free(data);

	SDL_RWclose(ico);
	SDL_RWclose(exe);
}