/* Copyright (C) 2019  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 <errno.h>
#include <stdlib.h>
#include <string.h>

#include "cards.h"

static SDL_Surface *
dibSurface(SDL_RWops *rw, Uint32 offset, Uint32 length) {
	Uint32 bitmapLength = 0x0E + length;

	Uint8 *buffer = malloc(bitmapLength);
	if (!buffer) {
		SDL_SetError("malloc error: %s", strerror(errno));
		return NULL;
	}

	SDL_RWops *bitmap = SDL_RWFromMem(buffer, bitmapLength);
	if (!bitmap) {
		free(buffer);
		return NULL;
	}

	if (SDL_RWseek(rw, offset, RW_SEEK_SET) < 0) goto fail;
	Uint32 dibHeaderLength = SDL_ReadLE32(rw);
	Uint32 bitmapDataOffset = 0x0E + dibHeaderLength;

	if (dibHeaderLength == 0x0C) {
		if (SDL_RWseek(rw, 0x06, RW_SEEK_CUR) < 0) goto fail;
		Uint16 bitsPerPixel = SDL_ReadLE16(rw);
		bitmapDataOffset += 3 * (1 << bitsPerPixel);

	} else if (dibHeaderLength == 0x28) {
		if (SDL_RWseek(rw, 0x0A, RW_SEEK_CUR) < 0) goto fail;
		Uint16 bitsPerPixel = SDL_ReadLE16(rw);
		if (SDL_RWseek(rw, 0x10, RW_SEEK_CUR) < 0) goto fail;
		Uint32 paletteLength = SDL_ReadLE32(rw);

		if (!paletteLength && bitsPerPixel < 16) {
			paletteLength = 1 << bitsPerPixel;
		}
		bitmapDataOffset += 4 * paletteLength;

	} else {
		SDL_SetError("unrecognized DIB header length %u", dibHeaderLength);
		goto fail;
	}

	SDL_WriteU8(bitmap, 'B');
	SDL_WriteU8(bitmap, 'M');
	SDL_WriteLE32(bitmap, bitmapLength);
	SDL_WriteLE16(bitmap, 0); // reserved
	SDL_WriteLE16(bitmap, 0); // reserved
	SDL_WriteLE32(bitmap, bitmapDataOffset);

	if (SDL_RWseek(rw, offset, RW_SEEK_SET) < 0) goto fail;
	if (SDL_RWread(rw, &buffer[SDL_RWtell(bitmap)], length, 1) < 1) goto fail;

	SDL_RWseek(bitmap, 0, RW_SEEK_SET);
	SDL_Surface *surface = SDL_LoadBMP_RW(bitmap, 1);
	free(buffer);
	return surface;

fail:
	SDL_RWclose(bitmap);
	free(buffer);
	return NULL;
}

// exefmt.txt
static int
loadNE(
	SDL_Surface **surfaces, size_t count,
	SDL_RWops *rw, Uint16 neOffset, Uint32 idSkip
) {
	if (SDL_RWseek(rw, neOffset + 0x24, RW_SEEK_SET) < 0) return -1;
	Uint16 resourceTableOffset = neOffset + SDL_ReadLE16(rw);

	if (SDL_RWseek(rw, resourceTableOffset, RW_SEEK_SET) < 0) return -1;
	Uint16 alignmentShift = SDL_ReadLE16(rw);

	Uint16 resourceCount;
	for (;;) {
		Uint16 typeID = SDL_ReadLE16(rw);
		resourceCount = SDL_ReadLE16(rw);
		SDL_ReadLE32(rw); // reserved

		if (!typeID) {
			SDL_SetError("no bitmap resources");
			return -1;
		}
		if (typeID == 0x8002) break;

		if (SDL_RWseek(rw, 0x0C * resourceCount, RW_SEEK_CUR) < 0) return -1;
	}

	for (Uint16 i = 0; i < resourceCount; ++i) {
		Uint16 dataOffset = SDL_ReadLE16(rw);
		Uint16 dataLength = SDL_ReadLE16(rw);
		/* Uint16 flags = */ SDL_ReadLE16(rw);
		Uint16 resourceID = SDL_ReadLE16(rw);
		SDL_ReadLE32(rw); // reserved

		resourceID &= 0x7FFF;
		if (resourceID >= idSkip) resourceID -= idSkip;
		if (resourceID >= count) continue;

		Sint64 nextResource = SDL_RWtell(rw);
		if (nextResource < 0) return -1;

		surfaces[resourceID] = dibSurface(
			rw,
			(Uint32)dataOffset << alignmentShift,
			(Uint32)dataLength << alignmentShift
		);
		if (!surfaces[resourceID]) return -1;

		if (SDL_RWseek(rw, nextResource, RW_SEEK_SET) < 0) return -1;
	}

	return 0;
}

// <https://docs.microsoft.com/en-us/windows/desktop/Debug/pe-format>
static int
loadPE(
	SDL_Surface **surfaces, size_t count,
	SDL_RWops *rw, Uint16 peOffset, Uint32 idSkip
) {
	if (SDL_RWseek(rw, peOffset + 0x04 + 0x02, RW_SEEK_SET) < 0) return -1;
	Uint16 sectionCount = SDL_ReadLE16(rw);

	if (SDL_RWseek(rw, peOffset + 0x04 + 0x10, RW_SEEK_SET) < 0) return -1;
	Uint16 optionalHeaderLength = SDL_ReadLE16(rw);
	Uint16 sectionTableOffset = peOffset + 0x04 + 0x14 + optionalHeaderLength;

	Uint32 resourceTableOffset = 0;
	Uint32 resourceTableVirtual = 0;
	for (Uint16 i = 0; i < sectionCount; ++i) {
		Uint16 headerOffset = sectionTableOffset + 0x28 * i;

		if (SDL_RWseek(rw, headerOffset, RW_SEEK_SET) < 0) return -1;
		char name[8];
		if (!SDL_RWread(rw, name, sizeof(name), 1)) return -1;
		if (strncmp(".rsrc", name, sizeof(name))) continue;

		if (SDL_RWseek(rw, headerOffset + 0x0C, RW_SEEK_SET) < 0) return -1;
		resourceTableVirtual = SDL_ReadLE32(rw);

		if (SDL_RWseek(rw, headerOffset + 0x14, RW_SEEK_SET) < 0) return -1;
		resourceTableOffset = SDL_ReadLE32(rw);

		break;
	}
	if (!resourceTableOffset) {
		SDL_SetError("no resource table");
		return -1;
	}

	if (SDL_RWseek(rw, resourceTableOffset + 0x0C, RW_SEEK_SET) < 0) return -1;
	Uint16 typeNameCount = SDL_ReadLE16(rw);
	Uint16 typeIDCount = SDL_ReadLE16(rw);

	if (SDL_RWseek(rw, 0x08 * typeNameCount, RW_SEEK_CUR) < 0) return -1;

	Uint32 bitmapTableOffset = 0;
	for (Uint16 i = 0; i < typeIDCount; ++i) {
		Uint32 typeID = SDL_ReadLE32(rw);
		Uint32 subdirOffset = SDL_ReadLE32(rw);
		if (typeID != 0x02) continue;
		if (!(subdirOffset & (1 << 31))) {
			SDL_SetError("bitmap type entry does not point to table");
			return -1;
		}
		bitmapTableOffset = resourceTableOffset + (subdirOffset & ~(1 << 31));
		break;
	}

	if (SDL_RWseek(rw, bitmapTableOffset + 0x0C, RW_SEEK_SET) < 0) return -1;
	Uint16 nameNameCount = SDL_ReadLE16(rw);
	Uint16 nameIDCount = SDL_ReadLE16(rw);

	if (SDL_RWseek(rw, 0x08 * nameNameCount, RW_SEEK_CUR) < 0) return -1;

	for (Uint16 i = 0; i < nameIDCount; ++i) {
		Uint32 nameID = SDL_ReadLE32(rw);
		if (nameID >= idSkip) nameID -= idSkip;
		if (nameID >= count) continue;

		Uint32 subdirOffset = SDL_ReadLE32(rw);
		if (!(subdirOffset & (1 << 31))) {
			SDL_SetError("bitmap name entry does not point to table");
			return -1;
		}

		Sint64 nextName = SDL_RWtell(rw);
		if (nextName < 0) return -1;

		Uint32 langTableOffset =
			resourceTableOffset + (subdirOffset & ~(1 << 31));

		if (SDL_RWseek(rw, langTableOffset + 0x0C, RW_SEEK_SET) < 0) return -1;
		Uint16 langNameCount = SDL_ReadLE16(rw);
		Uint16 langIDCount = SDL_ReadLE16(rw);
		if (langNameCount != 0 || langIDCount != 1) {
			SDL_SetError("language table contains more than one entry");
			return -1;
		}

		/* Uint32 langID = */ SDL_ReadLE32(rw);
		Uint32 dataEntryOffset = SDL_ReadLE32(rw);
		if (dataEntryOffset & (1 << 31)) {
			SDL_SetError("language entry does not point to data");
			return -1;
		}
		dataEntryOffset += resourceTableOffset;

		if (SDL_RWseek(rw, dataEntryOffset, RW_SEEK_SET) < 0) return -1;
		Uint32 dataVirtual = SDL_ReadLE32(rw);
		Uint32 dataLength = SDL_ReadLE32(rw);
		Uint32 dataOffset =
			dataVirtual - (resourceTableVirtual - resourceTableOffset);

		surfaces[nameID] = dibSurface(rw, dataOffset, dataLength);
		if (!surfaces[nameID]) return -1;

		if (SDL_RWseek(rw, nextName, RW_SEEK_SET) < 0) return -1;
	}

	return 0;
}

static int
loadEXE(SDL_Surface **surfaces, size_t count, SDL_RWops *rw, Uint32 idSkip) {
	if (SDL_RWseek(rw, 0, RW_SEEK_SET) < 0) return -1;
	if (SDL_ReadU8(rw) != 'M' || SDL_ReadU8(rw) != 'Z') {
		SDL_SetError("invalid MZ signature");
		return -1;
	}

	if (SDL_RWseek(rw, 0x3C, RW_SEEK_SET) < 0) return -1;
	Uint16 offset = SDL_ReadLE16(rw);

	if (SDL_RWseek(rw, offset, RW_SEEK_SET) < 0) return -1;
	Uint8 sig[2] = { SDL_ReadU8(rw), SDL_ReadU8(rw) };

	if (sig[0] == 'N' && sig[1] == 'E') {
		return loadNE(surfaces, count, rw, offset, idSkip);
	} else if (sig[0] == 'P' && sig[1] == 'E') {
		return loadPE(surfaces, count, rw, offset, idSkip);
	} else {
		SDL_SetError("invalid NE/PE signature");
		return -1;
	}
}

static int setColorKey(SDL_Surface **surfaces, size_t count) {
	size_t i = Cards_Empty;
	if (i < count) {
		if (SDL_SetColorKey(surfaces[i], SDL_TRUE, 1) < 0) return -1;
	}
	for (i = Cards_X; i <= Cards_O; ++i) {
		if (i >= count) break;
		if (SDL_SetColorKey(surfaces[i], SDL_TRUE, 12) < 0) return -1;
	}
	return 0;
}

static int setAlphaCorners(SDL_Surface **surfaces, size_t count) {
	SDL_Surface *alpha = NULL;
	for (size_t i = 0; i < count; ++i) {
		if (!surfaces[i]) continue;

		alpha = SDL_ConvertSurfaceFormat(surfaces[i], SDL_PIXELFORMAT_RGBA32, 0);
		if (!alpha) return -1;

		if (SDL_SetSurfaceBlendMode(alpha, SDL_BLENDMODE_BLEND) < 0) goto fail;

		SDL_Rect rects[8] = {
			{ 0, 0, 2, 1 },
			{ 0, 1, 1, 1 },
			{ Cards_CardWidth - 2, 0, 2, 1 },
			{ Cards_CardWidth - 1, 1, 1, 1 },
			{ 0, Cards_CardHeight - 1, 2, 1 },
			{ 0, Cards_CardHeight - 2, 1, 1 },
			{ Cards_CardWidth - 2, Cards_CardHeight - 1, 2, 1 },
			{ Cards_CardWidth - 1, Cards_CardHeight - 2, 1, 1 },
		};
		Uint32 trans = SDL_MapRGBA(alpha->format, 0x00, 0x00, 0x00, 0x00);
		if (SDL_FillRects(alpha, rects, 8, trans) < 0) goto fail;

		SDL_FreeSurface(surfaces[i]);
		surfaces[i] = alpha;
	}
	return 0;

fail:
	SDL_FreeSurface(alpha);
	return -1;
}

static int setBlackBorders(SDL_Surface **surfaces, size_t count) {
	for (size_t i = Cards_Diamond + Cards_A; i <= Cards_Heart + Cards_K; ++i) {
		if (i >= count) break;
		if (!surfaces[i]) continue;
		SDL_Rect rects[8] = {
			{ 2, 0, Cards_CardWidth - 4, 1 },
			{ 2, Cards_CardHeight - 1, Cards_CardWidth - 4, 1 },
			{ 0, 2, 1, Cards_CardHeight - 4 },
			{ Cards_CardWidth - 1, 2, 1, Cards_CardHeight - 4 },
			{ 1, 1, 1, 1 },
			{ Cards_CardWidth - 2, 1, 1, 1 },
			{ 1, Cards_CardHeight - 2, 1, 1 },
			{ Cards_CardWidth - 2, Cards_CardHeight - 2, 1, 1 },
		};
		Uint32 black = SDL_MapRGB(surfaces[i]->format, 0x00, 0x00, 0x00);
		if (SDL_FillRects(surfaces[i], rects, 8, black) < 0) return -1;
	}
	return 0;
}

static int
checkRange(SDL_Surface **surfaces, size_t count, size_t a, size_t b) {
	for (size_t i = a; i <= b; ++i) {
		if (i >= count) break;
		if (surfaces[i]) continue;
		SDL_SetError("missing resource %zu", i);
		return -1;
	}
	return 0;
}

int
Cards_LoadCards(
	SDL_Surface *surfaces[], size_t count,
	SDL_RWops *rw, enum Cards_Flag flags
) {
	memset(surfaces, 0, sizeof(*surfaces) * count);
	if (loadEXE(surfaces, count, rw, 0) < 0) return -1;
	if (checkRange(surfaces, count, Cards_A, Cards_Back12) < 0) return -1;
	if (checkRange(surfaces, count, Cards_X, Cards_O) < 0) return -1;
	if (flags & Cards_ColorKey) {
		if (setColorKey(surfaces, count) < 0) return -1;
	}
	if (flags & Cards_AlphaCorners) {
		if (setAlphaCorners(surfaces, count) < 0) return -1;
	}
	if (flags & Cards_BlackBorders) {
		if (setBlackBorders(surfaces, count) < 0) return -1;
	}
	return 0;
}

int
Cards_LoadFreeCell(
	SDL_Surface *surfaces[], size_t count,
	SDL_RWops *rw, enum Cards_Flag flags
) {
	memset(surfaces, 0, sizeof(*surfaces) * count);
	if (loadEXE(surfaces, count, rw, 402) < 0) return -1;
	if (checkRange(surfaces, count, Cards_KingRight, Cards_KingWin) < 0) {
		return -1;
	}
	if (flags & Cards_ColorKey) {
		for (size_t i = Cards_KingRight; i <= Cards_KingWin; ++i) {
			if (i >= count) break;
			if (SDL_SetColorKey(surfaces[i], SDL_TRUE, 2) < 0) return -1;
		}
	}
	return 0;
}

static int invertPalette(SDL_Surface *surface) {
	const SDL_Palette *palette = surface->format->palette;
	SDL_Palette *invert = SDL_AllocPalette(palette->ncolors);
	if (!invert) return -1;
	for (int i = 0; i < invert->ncolors; ++i) {
		invert->colors[i].r = ~palette->colors[i].r;
		invert->colors[i].g = ~palette->colors[i].g;
		invert->colors[i].b = ~palette->colors[i].b;
		invert->colors[i].a =  palette->colors[i].a;
	}
	if (SDL_SetSurfacePalette(surface, invert) < 0) return -1;
	SDL_FreePalette(invert);
	return 0;
}

static int invertPixels(SDL_Surface *surface) {
	if (SDL_LockSurface(surface) < 0) return -1;
	Uint8 *pixels = surface->pixels;
	for (int y = 0; y < surface->h; ++y) {
		Uint32 *row = (Uint32 *)&pixels[y * surface->pitch];
		for (int x = 0; x < surface->w; ++x) {
			Uint32 color = ~row[x] & ~surface->format->Amask;
			Uint32 alpha = row[x] & surface->format->Amask;
			row[x] = color | alpha;
		}
	}
	SDL_UnlockSurface(surface);
	return 0;
}

int Cards_InvertSurface(SDL_Surface *surface) {
	if (surface->format->palette) {
		if (invertPalette(surface) < 0) return -1;
	} else if (surface->format->BytesPerPixel == 4) {
		if (invertPixels(surface) < 0) return -1;
	} else {
		SDL_SetError("cannot invert surface format");
		return -1;
	}
	return 0;
}