/* SPDX-License-Identifier: MIT */
/*
 * konpack.c
 *
 * Copyright (c) 2026 KonAki
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */


#include "konpack.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#define KON_PACK_MAGIC_ "KPAK"
#define KON_PACK_VERSION_ 1u
#define KON_PACK_HEADER_SIZE_ 36u
#define KON_PACK_FLAG_ENCRYPTED_ 1u
#define KON_PACK_TABLE_INDEX_ 0xFFFFFFFFu

static const char *kon_packErrorText_ = "no error";

const char *kon_packError(void) {
	return kon_packErrorText_;
}

/*** private helper ***/

static uint32_t kon_load32_(const uint8_t *p) {
	return (uint32_t)p[0] | ((uint32_t)p[1] << 8) | ((uint32_t)p[2] << 16) | ((uint32_t)p[3] << 24);
}

static void kon_store32_(uint8_t *p, uint32_t v) {
	p[0] = (uint8_t)(v & 0xFF);
	p[1] = (uint8_t)((v >> 8) & 0xFF);
	p[2] = (uint8_t)((v >> 16) & 0xFF);
	p[3] = (uint8_t)((v >> 24) & 0xFF);
}

/*** crc32 ***/

uint32_t kon_crc32(const uint8_t *data, size_t length) {
	static uint32_t table[256];
	static int ready;

	if (!ready) {
		for (uint32_t i = 0; i < 256; i++) {
			uint32_t c = i;
			for (int k = 0; k < 8; k++) c = (c & 1u) ? (0xEDB88320u ^ (c >> 1)) : (c >> 1);
			table[i] = c;
		}
		ready = 1;
	}

	uint32_t crc = 0xFFFFFFFFu;
	for (size_t i = 0; i < length; i++) crc = table[(crc ^ data[i]) & 0xFF] ^ (crc >> 8);
	return crc ^ 0xFFFFFFFFu;
}

/*** chacha20 ***/

#define KON_ROTL32_(v, n) (((v) << (n)) | ((v) >> (32 - (n))))

#define KON_QUARTER_(a, b, c, d) \
	a += b; d ^= a; d = KON_ROTL32_(d, 16); \
	c += d; b ^= c; b = KON_ROTL32_(b, 12); \
	a += b; d ^= a; d = KON_ROTL32_(d, 8); \
	c += d; b ^= c; b = KON_ROTL32_(b, 7);

static void kon_chachaBlock_(const uint32_t state[16], uint8_t out[64]) {
	uint32_t x[16];

	for (int i = 0; i < 16; i++) x[i] = state[i];

	for (int round = 0; round < 10; round++) {
		KON_QUARTER_(x[0], x[4], x[8],  x[12])
		KON_QUARTER_(x[1], x[5], x[9],  x[13])
		KON_QUARTER_(x[2], x[6], x[10], x[14])
		KON_QUARTER_(x[3], x[7], x[11], x[15])
		KON_QUARTER_(x[0], x[5], x[10], x[15])
		KON_QUARTER_(x[1], x[6], x[11], x[12])
		KON_QUARTER_(x[2], x[7], x[8],  x[13])
		KON_QUARTER_(x[3], x[4], x[9],  x[14])
	}

	for (int i = 0; i < 16; i++) kon_store32_(&out[i * 4], x[i] + state[i]);
}

void kon_chacha20Xor(const uint8_t key[KON_PACK_KEY_SIZE], const uint8_t nonce[12], uint32_t counter, const uint8_t *in, uint8_t *out, size_t length) {
	uint32_t state[16];
	uint8_t block[64];

	state[0] = 0x61707865u;
	state[1] = 0x3320646Eu;
	state[2] = 0x79622D32u;
	state[3] = 0x6B206574u;
	for (int i = 0; i < 8; i++) state[4 + i] = kon_load32_(&key[i * 4]);
	state[12] = counter;
	for (int i = 0; i < 3; i++) state[13 + i] = kon_load32_(&nonce[i * 4]);

	while (length > 0) {
		size_t n = length < 64 ? length : 64;

		kon_chachaBlock_(state, block);
		for (size_t i = 0; i < n; i++) out[i] = (uint8_t)(in[i] ^ block[i]);

		in += n;
		out += n;
		length -= n;
		state[12]++;
	}
}

static void kon_fileNonce_(const uint8_t packNonce[8], uint32_t index, uint8_t out[12]) {
	memcpy(out, packNonce, 8);
	kon_store32_(&out[8], index);
}

/*** reading ***/

typedef struct kon_packEntry {
	const char *name;
	uint32_t offset, size, crc;
} kon_packEntry_t;

struct kon_pack {
	FILE *file;
	uint8_t key[KON_PACK_KEY_SIZE];
	uint8_t nonce[8];
	int encrypted;
	int count;
	kon_packEntry_t *entries;
	char *names;
};

void kon_packClose(kon_pack_t *pack) {
	if (!pack) return;

	if (pack->file) fclose(pack->file);
	free(pack->entries);
	free(pack->names);
	free(pack);
}

kon_pack_t *kon_packOpen(const char *path, const uint8_t *key) {
	uint8_t header[KON_PACK_HEADER_SIZE_];
	uint8_t *table = NULL;
	kon_pack_t *pack = NULL;
	uint32_t flags, count, tableOffset, tableSize, tableCrc;
	uint64_t fileSize;
	long end;

	if (!path) {
		kon_packErrorText_ = "no path";
		return NULL;
	}

	pack = calloc(1, sizeof(kon_pack_t));
	if (!pack) {
		kon_packErrorText_ = "out of memory";
		return NULL;
	}

	pack->file = fopen(path, "rb");
	if (!pack->file) {
		kon_packErrorText_ = "could not open the pack file";
		goto fail;
	}

	if (fread(header, 1, sizeof(header), pack->file) != sizeof(header)) {
		kon_packErrorText_ = "pack is too small";
		goto fail;
	}
	if (memcmp(header, KON_PACK_MAGIC_, 4) != 0) {
		kon_packErrorText_ = "not a pack file";
		goto fail;
	}
	if (kon_load32_(&header[4]) != KON_PACK_VERSION_) {
		kon_packErrorText_ = "unsupported pack version";
		goto fail;
	}

	flags = kon_load32_(&header[8]);
	count = kon_load32_(&header[12]);
	tableOffset = kon_load32_(&header[16]);
	tableSize = kon_load32_(&header[20]);
	tableCrc = kon_load32_(&header[24]);
	memcpy(pack->nonce, &header[28], 8);

	pack->encrypted = (flags & KON_PACK_FLAG_ENCRYPTED_) != 0;
	if (pack->encrypted) {
		if (!key) {
			kon_packErrorText_ = "pack is encrypted and no key was given";
			goto fail;
		}
		memcpy(pack->key, key, KON_PACK_KEY_SIZE);
	}

	if (fseek(pack->file, 0, SEEK_END) != 0 || (end = ftell(pack->file)) < 0) {
		kon_packErrorText_ = "could not get the pack size";
		goto fail;
	}
	fileSize = (uint64_t)end;

	if (tableOffset < KON_PACK_HEADER_SIZE_ || (uint64_t)tableOffset + tableSize > fileSize ||
		(uint64_t)count * 15u > tableSize || count > 0x7FFFFFFFu) {
		kon_packErrorText_ = "pack table is damaged";
		goto fail;
	}

	table = malloc(tableSize ? tableSize : 1);
	pack->entries = calloc(count ? count : 1, sizeof(kon_packEntry_t));
	pack->names = malloc((size_t)tableSize + count + 1);
	if (!table || !pack->entries || !pack->names) {
		kon_packErrorText_ = "out of memory";
		goto fail;
	}

	if (fseek(pack->file, (long)tableOffset, SEEK_SET) != 0 || fread(table, 1, tableSize, pack->file) != tableSize) {
		kon_packErrorText_ = "could not read the pack table";
		goto fail;
	}

	if (pack->encrypted) {
		uint8_t nonce[12];
		kon_fileNonce_(pack->nonce, KON_PACK_TABLE_INDEX_, nonce);
		kon_chacha20Xor(pack->key, nonce, 0, table, table, tableSize);
	}

	if (kon_crc32(table, tableSize) != tableCrc) {
		kon_packErrorText_ = pack->encrypted ? "wrong key or damaged pack" : "pack table is damaged";
		goto fail;
	}

	{
		size_t at = 0, nameAt = 0;

		for (uint32_t i = 0; i < count; i++) {
			size_t nameLength;

			if (at + 2 > tableSize) goto damaged;
			nameLength = (size_t)table[at] | ((size_t)table[at + 1] << 8);
			at += 2;
			if (nameLength == 0 || at + nameLength + 12 > tableSize) goto damaged;

			memcpy(&pack->names[nameAt], &table[at], nameLength);
			pack->names[nameAt + nameLength] = '\0';
			if (strlen(&pack->names[nameAt]) != nameLength) goto damaged;
			pack->entries[i].name = &pack->names[nameAt];
			nameAt += nameLength + 1;
			at += nameLength;

			pack->entries[i].offset = kon_load32_(&table[at]);
			pack->entries[i].size = kon_load32_(&table[at + 4]);
			pack->entries[i].crc = kon_load32_(&table[at + 8]);
			at += 12;

			if (pack->entries[i].offset < KON_PACK_HEADER_SIZE_ ||
				(uint64_t)pack->entries[i].offset + pack->entries[i].size > tableOffset) goto damaged;
			if (i > 0 && strcmp(pack->entries[i - 1].name, pack->entries[i].name) >= 0) goto damaged;
		}
		if (at != tableSize) goto damaged;
	}

	pack->count = (int)count;
	free(table);
	return pack;

damaged:
	kon_packErrorText_ = "pack table is damaged";
fail:
	free(table);
	kon_packClose(pack);
	return NULL;
}

int kon_packCount(const kon_pack_t *pack) {
	return pack ? pack->count : 0;
}

const char *kon_packName(const kon_pack_t *pack, int index) {
	if (!pack || index < 0 || index >= pack->count) return NULL;
	return pack->entries[index].name;
}

uint32_t kon_packSize(const kon_pack_t *pack, int index) {
	if (!pack || index < 0 || index >= pack->count) return 0;
	return pack->entries[index].size;
}

int kon_packFind(const kon_pack_t *pack, const char *name) {
	int low = 0, high;

	if (!pack || !name) return -1;
	high = pack->count - 1;

	while (low <= high) {
		int mid = low + (high - low) / 2;
		int cmp = strcmp(name, pack->entries[mid].name);

		if (cmp == 0) return mid;
		if (cmp < 0) high = mid - 1;
		else low = mid + 1;
	}
	return -1;
}

uint8_t *kon_packRead(kon_pack_t *pack, int index, size_t *size) {
	kon_packEntry_t *entry;
	uint8_t *data;

	if (size) *size = 0;
	if (!pack || index < 0 || index >= pack->count) {
		kon_packErrorText_ = "no such file in the pack";
		return NULL;
	}
	entry = &pack->entries[index];

	data = malloc((size_t)entry->size + 1);
	if (!data) {
		kon_packErrorText_ = "out of memory";
		return NULL;
	}

	if (fseek(pack->file, (long)entry->offset, SEEK_SET) != 0 || fread(data, 1, entry->size, pack->file) != entry->size) {
		kon_packErrorText_ = "could not read from the pack";
		free(data);
		return NULL;
	}

	if (pack->encrypted) {
		uint8_t nonce[12];
		kon_fileNonce_(pack->nonce, (uint32_t)index, nonce);
		kon_chacha20Xor(pack->key, nonce, 0, data, data, entry->size);
	}

	if (kon_crc32(data, entry->size) != entry->crc) {
		kon_packErrorText_ = "file is damaged";
		free(data);
		return NULL;
	}

	data[entry->size] = 0;
	if (size) *size = entry->size;
	return data;
}

/*** writing ***/

#ifdef KONPACK_WRITER

typedef struct kon_packSource_ {
	const char *path;
	char *name;
} kon_packSource_t;

static int kon_packSourceCompare_(const void *a, const void *b) {
	return strcmp(((const kon_packSource_t *)a)->name, ((const kon_packSource_t *)b)->name);
}

static void kon_packMakeNonce_(uint8_t nonce[8]) {
	FILE *f = fopen("/dev/urandom", "rb");

	if (f) {
		size_t got = fread(nonce, 1, 8, f);
		fclose(f);
		if (got == 8) return;
	}

	/* nothing better around, mix what changes between runs. only has to never repeat, it is not secret */
	{
		uint64_t h = 1469598103934665603ull;
		uint64_t parts[3];
		int local;

		parts[0] = (uint64_t)time(NULL);
		parts[1] = (uint64_t)clock();
		parts[2] = (uint64_t)(size_t)&local;

		for (int i = 0; i < 3; i++) {
			for (int b = 0; b < 8; b++) {
				h ^= (parts[i] >> (b * 8)) & 0xFF;
				h *= 1099511628211ull;
			}
		}
		for (int i = 0; i < 8; i++) nonce[i] = (uint8_t)(h >> (i * 8));
	}
}

static char *kon_packMakeName_(const char *path, const char *baseDir) {
	size_t baseLength = baseDir ? strlen(baseDir) : 0;
	const char *rest = path;
	char *name;

	if (baseLength > 0 && strncmp(path, baseDir, baseLength) == 0) {
		rest = path + baseLength;
		while (*rest == '/' || *rest == '\\') rest++;
	}
	while (rest[0] == '.' && (rest[1] == '/' || rest[1] == '\\')) rest += 2;

	name = malloc(strlen(rest) + 1);
	if (!name) return NULL;

	strcpy(name, rest);
	for (char *c = name; *c; c++) {
		if (*c == '\\') *c = '/';
	}
	return name;
}

int kon_packWrite(const char *outPath, const uint8_t *key, const char *baseDir, const char *const *files, int count) {
	kon_packSource_t *sources = NULL;
	FILE *out = NULL;
	uint8_t *table = NULL;
	uint32_t *offsets = NULL, *sizes = NULL, *crcs = NULL;
	uint8_t header[KON_PACK_HEADER_SIZE_];
	uint8_t packNonce[8];
	uint64_t position = KON_PACK_HEADER_SIZE_;
	size_t tableSize = 0, at = 0;
	int result = -1, made = 0;

	if (!outPath || !files || count < 0) {
		kon_packErrorText_ = "bad arguments";
		return -1;
	}

	sources = calloc(count ? (size_t)count : 1, sizeof(kon_packSource_t));
	offsets = calloc(count ? (size_t)count : 1, sizeof(uint32_t));
	sizes = calloc(count ? (size_t)count : 1, sizeof(uint32_t));
	crcs = calloc(count ? (size_t)count : 1, sizeof(uint32_t));
	if (!sources || !offsets || !sizes || !crcs) {
		kon_packErrorText_ = "out of memory";
		goto done;
	}

	for (int i = 0; i < count; i++) {
		sources[i].path = files[i];
		sources[i].name = kon_packMakeName_(files[i], baseDir);
		made++;
		if (!sources[i].name || sources[i].name[0] == '\0' || strlen(sources[i].name) > 65535) {
			kon_packErrorText_ = "bad file name";
			goto done;
		}
	}

	qsort(sources, (size_t)count, sizeof(kon_packSource_t), kon_packSourceCompare_);
	for (int i = 1; i < count; i++) {
		if (strcmp(sources[i - 1].name, sources[i].name) == 0) {
			kon_packErrorText_ = "two files have the same name in the pack";
			goto done;
		}
	}

	kon_packMakeNonce_(packNonce);

	out = fopen(outPath, "wb");
	if (!out) {
		kon_packErrorText_ = "could not create the pack file";
		goto done;
	}

	memset(header, 0, sizeof(header));
	if (fwrite(header, 1, sizeof(header), out) != sizeof(header)) {
		kon_packErrorText_ = "could not write the pack";
		goto done;
	}

	for (int i = 0; i < count; i++) {
		FILE *in = fopen(sources[i].path, "rb");
		uint8_t *data = NULL;
		long length;

		if (!in) {
			kon_packErrorText_ = "could not open an input file";
			goto done;
		}
		if (fseek(in, 0, SEEK_END) != 0 || (length = ftell(in)) < 0 || fseek(in, 0, SEEK_SET) != 0) {
			kon_packErrorText_ = "could not get an input file size";
			fclose(in);
			goto done;
		}
		if ((uint64_t)length > 0xFFFFFFFFu || position + (uint64_t)length > 0xFFFFFFFFu) {
			kon_packErrorText_ = "pack would be over 4 GB";
			fclose(in);
			goto done;
		}

		data = malloc((size_t)length ? (size_t)length : 1);
		if (!data || fread(data, 1, (size_t)length, in) != (size_t)length) {
			kon_packErrorText_ = "could not read an input file";
			free(data);
			fclose(in);
			goto done;
		}
		fclose(in);

		offsets[i] = (uint32_t)position;
		sizes[i] = (uint32_t)length;
		crcs[i] = kon_crc32(data, (size_t)length);

		if (key) {
			uint8_t nonce[12];
			kon_fileNonce_(packNonce, (uint32_t)i, nonce);
			kon_chacha20Xor(key, nonce, 0, data, data, (size_t)length);
		}

		if (fwrite(data, 1, (size_t)length, out) != (size_t)length) {
			kon_packErrorText_ = "could not write the pack";
			free(data);
			goto done;
		}
		free(data);
		position += (uint64_t)length;
	}

	for (int i = 0; i < count; i++) tableSize += 2 + strlen(sources[i].name) + 12;
	if (position + tableSize > 0xFFFFFFFFu) {
		kon_packErrorText_ = "pack would be over 4 GB";
		goto done;
	}

	table = malloc(tableSize ? tableSize : 1);
	if (!table) {
		kon_packErrorText_ = "out of memory";
		goto done;
	}
	for (int i = 0; i < count; i++) {
		size_t nameLength = strlen(sources[i].name);

		table[at] = (uint8_t)(nameLength & 0xFF);
		table[at + 1] = (uint8_t)((nameLength >> 8) & 0xFF);
		memcpy(&table[at + 2], sources[i].name, nameLength);
		at += 2 + nameLength;
		kon_store32_(&table[at], offsets[i]);
		kon_store32_(&table[at + 4], sizes[i]);
		kon_store32_(&table[at + 8], crcs[i]);
		at += 12;
	}

	memcpy(header, KON_PACK_MAGIC_, 4);
	kon_store32_(&header[4], KON_PACK_VERSION_);
	kon_store32_(&header[8], key ? KON_PACK_FLAG_ENCRYPTED_ : 0u);
	kon_store32_(&header[12], (uint32_t)count);
	kon_store32_(&header[16], (uint32_t)position);
	kon_store32_(&header[20], (uint32_t)tableSize);
	kon_store32_(&header[24], kon_crc32(table, tableSize));
	memcpy(&header[28], packNonce, 8);

	if (key) {
		uint8_t nonce[12];
		kon_fileNonce_(packNonce, KON_PACK_TABLE_INDEX_, nonce);
		kon_chacha20Xor(key, nonce, 0, table, table, tableSize);
	}

	if (fwrite(table, 1, tableSize, out) != tableSize || fseek(out, 0, SEEK_SET) != 0 ||
		fwrite(header, 1, sizeof(header), out) != sizeof(header)) {
		kon_packErrorText_ = "could not write the pack";
		goto done;
	}

	result = 0;

done:
	if (out && fclose(out) != 0 && result == 0) {
		kon_packErrorText_ = "could not write the pack";
		result = -1;
	}
	if (sources) {
		for (int i = 0; i < made; i++) free(sources[i].name);
	}
	free(sources);
	free(offsets);
	free(sizes);
	free(crcs);
	free(table);
	return result;
}

#endif /* KONPACK_WRITER */