/*
  Copyright (c) 2021 Alex Diener
  
  This software is provided 'as-is', without any express or implied
  warranty. In no event will the authors be held liable for any damages
  arising from the use of this software.
  
  Permission is granted to anyone to use this software for any purpose,
  including commercial applications, and to alter it and redistribute it
  freely, subject to the following restrictions:
  
  1. The origin of this software must not be misrepresented; you must not
     claim that you wrote the original software. If you use this software
     in a product, an acknowledgment in the product documentation would be
     appreciated but is not required.
  2. Altered source versions must be plainly marked as such, and must not be
     misrepresented as being the original software.
  3. This notice may not be removed or altered from any source distribution.
  
  Alex Diener alex@ludobloom.com
*/

#include "utilities/FileBundle.h"
#include "utilities/IOUtilities.h"
#include <limits.h>
#include <stdint.h>
#include <string.h>

FileBundle * FileBundle_create(void) {
	FileBundle * bundle = malloc(sizeof(*bundle));
	bundle->entryCount = 0;
	bundle->entries = NULL;
	bundle->file = NULL;
	bundle->data = NULL;
	bundle->dataSize = 0;
	bundle->dataOwned = false;
	return bundle;
}

void FileBundle_dispose(FileBundle * bundle) {
	for (unsigned int entryIndex = 0; entryIndex < bundle->entryCount; entryIndex++) {
		free(bundle->entries[entryIndex].identifier);
		if (bundle->entries[entryIndex].dataOwned) {
			free(bundle->entries[entryIndex].data);
		}
	}
	free(bundle->entries);
	if (bundle->file != NULL) {
		fclose(bundle->file);
	}
	if (bundle->dataOwned) {
		free(bundle->data);
	}
	free(bundle);
}

unsigned int FileBundle_getFileCount(FileBundle * bundle) {
	return bundle->entryCount;
}

const char * FileBundle_getFileIdentifier(FileBundle * bundle, unsigned int fileIndex) {
	if (fileIndex < bundle->entryCount) {
		return bundle->entries[fileIndex].identifier;
	}
	return NULL;
}

unsigned int FileBundle_getFileIndex(FileBundle * bundle, const char * fileIdentifier) {
	for (unsigned int entryIndex = 0; entryIndex < bundle->entryCount; entryIndex++) {
		if (!strcmp(bundle->entries[entryIndex].identifier, fileIdentifier)) {
			return entryIndex;
		}
	}
	return UINT_MAX;
}

FileBundle_fileType FileBundle_getFileType(FileBundle * bundle, const char * fileIdentifier) {
	return FileBundle_getFileTypeAtIndex(bundle, FileBundle_getFileIndex(bundle, fileIdentifier));
}

FileBundle_fileType FileBundle_getFileTypeAtIndex(FileBundle * bundle, unsigned int fileIndex) {
	if (fileIndex < bundle->entryCount) {
		return bundle->entries[fileIndex].type;
	}
	return FILE_TYPE_INVALID;
}

unsigned int FileBundle_getFileCountOfType(FileBundle * bundle, FileBundle_fileType type) {
	unsigned int count = 0;
	for (unsigned int entryIndex = 0; entryIndex < bundle->entryCount; entryIndex++) {
		if (bundle->entries[entryIndex].type == type) {
			count++;
		}
	}
	return count;
}

unsigned int FileBundle_getNextFileIndexOfType(FileBundle * bundle, FileBundle_fileType type, unsigned int startIndex) {
	for (; startIndex < bundle->entryCount; startIndex++) {
		if (bundle->entries[startIndex].type == type) {
			return startIndex;
		}
	}
	return UINT_MAX;
}

const void * FileBundle_getFile(FileBundle * bundle, const char * fileIdentifier, uint32_t * outSize) {
	return FileBundle_getFileAtIndex(bundle, FileBundle_getFileIndex(bundle, fileIdentifier), outSize);
}

const void * FileBundle_getFileAtIndex(FileBundle * bundle, unsigned int fileIndex, uint32_t * outSize) {
	if (fileIndex < bundle->entryCount) {
		if (outSize != NULL) {
			*outSize = bundle->entries[fileIndex].size;
		}
		if (bundle->entries[fileIndex].data == NULL) {
			if (bundle->data != NULL) {
				return bundle->data + bundle->entries[fileIndex].offset;
			}
			if (bundle->file != NULL) {
				bundle->entries[fileIndex].data = malloc(bundle->entries[fileIndex].size);
				fseek(bundle->file, bundle->entries[fileIndex].offset, SEEK_SET);
				if (fread(bundle->entries[fileIndex].data, 1, bundle->entries[fileIndex].size, bundle->file) < bundle->entries[fileIndex].size) {
					free(bundle->entries[fileIndex].data);
					bundle->entries[fileIndex].data = NULL;
				}
				bundle->entries[fileIndex].dataOwned = true;
			}
		}
		return bundle->entries[fileIndex].data;
	}
	return NULL;
}

void FileBundle_addFile(FileBundle * bundle, const char * fileIdentifier, FileBundle_fileType type, void * data, uint32_t size, bool takeOwnership, bool copy) {
	if (type == FILE_TYPE_INVALID) {
#ifdef DEBUG
		fprintf(stderr, "Warning: Attempt to add a file of type FILE_TYPE_INVALID to FileBundle %p\n", bundle);
#endif
		return;
	}
	bundle->entries = realloc(bundle->entries, (bundle->entryCount + 1) * sizeof(*bundle->entries));
	if (strlen(fileIdentifier) > FILE_BUNDLE_IDENTIFIER_MAX) {
		bundle->entries[bundle->entryCount].identifier = malloc(FILE_BUNDLE_IDENTIFIER_MAX + 1);
		memcpy(bundle->entries[bundle->entryCount].identifier, fileIdentifier, FILE_BUNDLE_IDENTIFIER_MAX);
		bundle->entries[bundle->entryCount].identifier[FILE_BUNDLE_IDENTIFIER_MAX] = 0;
#ifdef DEBUG
		fprintf(stderr, "Warning: File identifier \"%s\" is too long for FileBundle; it will be truncated to \"%s\"\n", fileIdentifier, bundle->entries[bundle->entryCount].identifier);
#endif
		
	} else {
		bundle->entries[bundle->entryCount].identifier = strdup(fileIdentifier);
	}
	bundle->entries[bundle->entryCount].type = type;
	if (copy) {
		bundle->entries[bundle->entryCount].data = malloc(size);
		memcpy(bundle->entries[bundle->entryCount].data, data, size);
	} else {
		bundle->entries[bundle->entryCount].data = data;
	}
	bundle->entries[bundle->entryCount].size = size;
	bundle->entries[bundle->entryCount].dataOwned = takeOwnership;
	bundle->entryCount++;
}

bool FileBundle_removeFile(FileBundle * bundle, const char * fileIdentifier) {
	return FileBundle_removeFileAtIndex(bundle, FileBundle_getFileIndex(bundle, fileIdentifier));
}

bool FileBundle_removeFileAtIndex(FileBundle * bundle, unsigned int fileIndex) {
	if (fileIndex < bundle->entryCount) {
		free(bundle->entries[fileIndex].identifier);
		if (bundle->entries[fileIndex].dataOwned) {
			free(bundle->entries[fileIndex].data);
		}
		bundle->entryCount--;
		for (; fileIndex < bundle->entryCount; fileIndex++) {
			bundle->entries[fileIndex] = bundle->entries[fileIndex + 1];
		}
		return true;
	}
	return false;
}

void FileBundle_replaceFileContents(FileBundle * bundle, const char * fileIdentifier, void * newData, uint32_t newSize, bool takeOwnership, bool copy) {
	FileBundle_replaceFileContentsAtIndex(bundle, FileBundle_getFileIndex(bundle, fileIdentifier), newData, newSize, takeOwnership, copy);
}

void FileBundle_replaceFileContentsAtIndex(FileBundle * bundle, unsigned int fileIndex, void * newData, uint32_t newSize, bool takeOwnership, bool copy) {
	if (fileIndex < bundle->entryCount) {
		if (bundle->entries[fileIndex].dataOwned) {
			free(bundle->entries[fileIndex].data);
		}
		if (copy) {
			bundle->entries[fileIndex].data = malloc(newSize);
			memcpy(bundle->entries[fileIndex].data, newData, newSize);
		} else {
			bundle->entries[fileIndex].data = newData;
		}
		bundle->entries[fileIndex].size = newSize;
		bundle->entries[fileIndex].dataOwned = takeOwnership;
	}
}

void FileBundle_setFileType(FileBundle * bundle, const char * fileIdentifier, FileBundle_fileType newType) {
	FileBundle_setFileTypeAtIndex(bundle, FileBundle_getFileIndex(bundle, fileIdentifier), newType);
}

void FileBundle_setFileTypeAtIndex(FileBundle * bundle, unsigned int fileIndex, FileBundle_fileType newType) {
	if (newType == FILE_TYPE_INVALID) {
#ifdef DEBUG
		fprintf(stderr, "Warning: Attempt to set a file's type to FILE_TYPE_INVALID in FileBundle %p\n", bundle);
#endif
		return;
	}
	if (fileIndex < bundle->entryCount) {
		bundle->entries[fileIndex].type = newType;
	}
}

void FileBundle_renameFile(FileBundle * bundle, const char * fileIdentifier, const char * newFileIdentifier) {
	FileBundle_renameFileAtIndex(bundle, FileBundle_getFileIndex(bundle, fileIdentifier), newFileIdentifier);
}

void FileBundle_renameFileAtIndex(FileBundle * bundle, unsigned int fileIndex, const char * newFileIdentifier) {
	if (fileIndex < bundle->entryCount) {
		free(bundle->entries[fileIndex].identifier);
		
		if (strlen(newFileIdentifier) > FILE_BUNDLE_IDENTIFIER_MAX) {
			bundle->entries[fileIndex].identifier = malloc(FILE_BUNDLE_IDENTIFIER_MAX + 1);
			memcpy(bundle->entries[fileIndex].identifier, newFileIdentifier, FILE_BUNDLE_IDENTIFIER_MAX);
			bundle->entries[fileIndex].identifier[FILE_BUNDLE_IDENTIFIER_MAX] = 0;
#ifdef DEBUG
			fprintf(stderr, "Warning: File identifier \"%s\" is too long for FileBundle; it will be truncated to \"%s\"\n", newFileIdentifier, bundle->entries[fileIndex].identifier);
#endif
			
		} else {
			bundle->entries[fileIndex].identifier = strdup(newFileIdentifier);
		}
	}
}

unsigned int FileBundle_reorderFile(FileBundle * bundle, const char * fileIdentifier, unsigned int newFileIndex) {
	return FileBundle_reorderFileAtIndex(bundle, FileBundle_getFileIndex(bundle, fileIdentifier), newFileIndex);
}

unsigned int FileBundle_reorderFileAtIndex(FileBundle * bundle, unsigned int fileIndex, unsigned int newFileIndex) {
	if (fileIndex >= bundle->entryCount) {
		return UINT_MAX;
	}
	if (newFileIndex >= bundle->entryCount) {
		newFileIndex = bundle->entryCount - 1;
	}
	struct FileBundle_entry swap = bundle->entries[fileIndex];
	int direction = newFileIndex < fileIndex ? -1 : 1;
	for (unsigned int swapIndex = fileIndex; swapIndex != newFileIndex; swapIndex += direction) {
		bundle->entries[swapIndex] = bundle->entries[swapIndex + direction];
	}
	bundle->entries[newFileIndex] = swap;
	return newFileIndex;
}

static uint32_t unpackUInt32(const unsigned char * bytes) {
	return (uint32_t) bytes[0] | (uint32_t) bytes[1] << 8 | (uint32_t) bytes[2] << 16 | (uint32_t) bytes[3] << 24;
}

static bool readFunctionMemread(void * context, size_t size, void * outData) {
	return memread(context, size, outData);
}

static bool readFunctionFread(void * context, size_t size, void * outData) {
	return fread(outData, 1, size, context) == size;
}

static bool parseBundleHeader(FileBundle * bundle, bool (* readFunction)(void * context, size_t size, void * outData), void * context) {
	unsigned char word[4];
	if (!readFunction(context, 4, word)) {
		return false;
	}
	if (memcmp(word, FILE_BUNDLE_FORMAT_SIGNATURE, 4)) {
		return false;
	}
	if (!readFunction(context, 4, word)) {
		return false;
	}
	uint32_t version = unpackUInt32(word);
	if (version > FILE_BUNDLE_FORMAT_VERSION) {
		return false;
	}
	if (!readFunction(context, 4, word)) { // header size
		return false;
	}
	if (!readFunction(context, 4, word)) {
		return false;
	}
	bundle->entryCount = unpackUInt32(word);
	bundle->entries = calloc(bundle->entryCount, sizeof(*bundle->entries));
	for (unsigned int entryIndex = 0; entryIndex < bundle->entryCount; entryIndex++) {
		if (!readFunction(context, 4, word)) {
			return false;
		}
		bundle->entries[entryIndex].type = unpackUInt32(word);
		if (!readFunction(context, 4, word)) {
			return false;
		}
		bundle->entries[entryIndex].offset = unpackUInt32(word);
		if (!readFunction(context, 4, word)) {
			return false;
		}
		bundle->entries[entryIndex].size = unpackUInt32(word);
		uint8_t identifierLength = 0;
		if (!readFunction(context, 1, &identifierLength)) {
			return false;
		}
		bundle->entries[entryIndex].identifier = malloc(identifierLength + 1);
		if (!readFunction(context, identifierLength, bundle->entries[entryIndex].identifier)) {
			return false;
		}
		bundle->entries[entryIndex].identifier[identifierLength] = 0;
	}
	return true;
}

FileBundle * FileBundle_loadData(const void * data, size_t size) {
	FileBundle * bundle = FileBundle_create();
	struct memreadContext context = memreadContextInit(data, size);
	if (!parseBundleHeader(bundle, readFunctionMemread, &context)) {
		FileBundle_dispose(bundle);
		return NULL;
	}
	
	for (unsigned int entryIndex = 0; entryIndex < bundle->entryCount; entryIndex++) {
		if ((size_t) bundle->entries[entryIndex].offset + bundle->entries[entryIndex].size > size) {
			FileBundle_dispose(bundle);
			return NULL;
		}
		bundle->entries[entryIndex].data = malloc(bundle->entries[entryIndex].size);
		memcpy(bundle->entries[entryIndex].data, data + bundle->entries[entryIndex].offset, bundle->entries[entryIndex].size);
		bundle->entries[entryIndex].dataOwned = true;
	}
	return bundle;
}

FileBundle * FileBundle_loadFile(const char * filePath) {
	size_t size = 0;
	void * data = readFileSimple(filePath, &size);
	if (data != NULL) {
		FileBundle * bundle = FileBundle_loadData(data, size);
		free(data);
		return bundle;
	}
	return NULL;
}

FileBundle * FileBundle_referenceDataNonImmediate(void * data, size_t size, bool takeOwnership, bool copy) {
	FileBundle * bundle = FileBundle_create();
	if (copy) {
		bundle->data = malloc(size);
		memcpy(bundle->data, data, size);
	} else {
		bundle->data = data;
	}
	bundle->dataSize = size;
	bundle->dataOwned = takeOwnership;
	
	struct memreadContext context = memreadContextInit(data, size);
	if (!parseBundleHeader(bundle, readFunctionMemread, &context)) {
		FileBundle_dispose(bundle);
		return NULL;
	}
	for (unsigned int entryIndex = 0; entryIndex < bundle->entryCount; entryIndex++) {
		if ((size_t) bundle->entries[entryIndex].offset + bundle->entries[entryIndex].size > size) {
			FileBundle_dispose(bundle);
			return NULL;
		}
	}
	return bundle;
}

FileBundle * FileBundle_referenceFileNonImmediate(const char * filePath) {
	FILE * file = fopen(filePath, "rb");
	if (file == NULL) {
		return NULL;
	}
	FileBundle * bundle = FileBundle_create();
	bundle->file = file;
	
	if (!parseBundleHeader(bundle, readFunctionFread, file)) {
		FileBundle_dispose(bundle);
		return NULL;
	}
	
	fseek(file, 0, SEEK_END);
	size_t fileLength = ftell(file);
	for (unsigned int entryIndex = 0; entryIndex < bundle->entryCount; entryIndex++) {
		if ((size_t) bundle->entries[entryIndex].offset + bundle->entries[entryIndex].size > fileLength) {
			FileBundle_dispose(bundle);
			return NULL;
		}
	}
	return bundle;
}

static void packUInt32(char * outBytes, uint32_t value) {
	outBytes[0] = value & 0xFF;
	outBytes[1] = value >> 8 & 0xFF;
	outBytes[2] = value >> 16 & 0xFF;
	outBytes[3] = value >> 24;
}

void * FileBundle_writeData(FileBundle * bundle, size_t * outSize) {
	char packedUInt32[4];
	uint32_t headerSize = 16;
	uint32_t totalSize = headerSize;
	for (unsigned int entryIndex = 0; entryIndex < bundle->entryCount; entryIndex++) {
		uint32_t entrySize = 12 + strlen(bundle->entries[entryIndex].identifier) + 1;
		headerSize += entrySize;
		totalSize += entrySize + bundle->entries[entryIndex].size;
	}
	struct memwriteContext context = memwriteContextInit(malloc(totalSize), 0, totalSize, false);
	memwrite(&context, 4, FILE_BUNDLE_FORMAT_SIGNATURE);
	packUInt32(packedUInt32, FILE_BUNDLE_FORMAT_VERSION);
	memwrite(&context, 4, packedUInt32);
	packUInt32(packedUInt32, headerSize);
	memwrite(&context, 4, packedUInt32);
	packUInt32(packedUInt32, bundle->entryCount);
	memwrite(&context, 4, packedUInt32);
	uint32_t dataOffset = headerSize;
	for (unsigned int entryIndex = 0; entryIndex < bundle->entryCount; entryIndex++) {
		packUInt32(packedUInt32, bundle->entries[entryIndex].type);
		memwrite(&context, 4, packedUInt32);
		packUInt32(packedUInt32, dataOffset);
		memwrite(&context, 4, packedUInt32);
		packUInt32(packedUInt32, bundle->entries[entryIndex].size);
		memwrite(&context, 4, packedUInt32);
		uint8_t length = strlen(bundle->entries[entryIndex].identifier);
		memwrite(&context, 1, &length);
		memwrite(&context, length, bundle->entries[entryIndex].identifier);
		dataOffset += bundle->entries[entryIndex].size;
	}
	for (unsigned int entryIndex = 0; entryIndex < bundle->entryCount; entryIndex++) {
		memwrite(&context, bundle->entries[entryIndex].size, FileBundle_getFileAtIndex(bundle, entryIndex, NULL));
	}
	if (outSize != NULL) {
		*outSize = totalSize;
	}
	return context.data;
}

bool FileBundle_writeFile(FileBundle * bundle, const char * filePath) {
	size_t size = 0;
	void * data = FileBundle_writeData(bundle, &size);
	bool success = writeFileSimple(filePath, data, size);
	free(data);
	return success;
}

bool FileBundle_writeFileAtomic(FileBundle * bundle, const char * filePath) {
	size_t size = 0;
	void * data = FileBundle_writeData(bundle, &size);
	bool success = writeFileAtomic(filePath, data, size);
	free(data);
	return success;
}