#include "fpk.h"

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

#define FPK_HEADER_FIELD_ENTRY_COUNT 4
#define FPK_HEADER_FIELD_DIR_OFFSET 12
#define FPK_HEADER_FIELD_STRTAB_OFFSET 20

#define FPK_DIR_ENTRY_FIELD_PATH_OFFSET 0
#define FPK_DIR_ENTRY_FIELD_DATA_OFFSET 8
#define FPK_DIR_ENTRY_FIELD_DATA_SIZE 16
#define FPK_DIR_ENTRY_FIELD_METHOD 24

struct fpk {
	FILE *file;
	uint64_t entry_count;
	fpk_entry_t *entries;
	char *strtab;
};

struct fpk_writer_entry {
	char *path;
	uint64_t data_offset;
	uint64_t data_size;
	uint8_t method;
};

struct fpk_writer {
	FILE *file;
	struct fpk_writer_entry *entries;
	uint64_t count;
	uint64_t capacity;
	uint64_t cursor;
};

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

static uint64_t get_u64le(const unsigned char *p)
{
	return (uint64_t)get_u32le(p) | (uint64_t)get_u32le(p + 4) << 32;
}

static void put_u32le(unsigned char *p, uint32_t v)
{
	p[0] = (unsigned char)(v);
	p[1] = (unsigned char)(v >> 8);
	p[2] = (unsigned char)(v >> 16);
	p[3] = (unsigned char)(v >> 24);
}

static void put_u64le(unsigned char *p, uint64_t v)
{
	put_u32le(p, (uint32_t)v);
	put_u32le(p + 4, (uint32_t)(v >> 32));
}

static char *dup_path(const char *path)
{
	size_t len = strlen(path) + 1;
	char *copy = malloc(len);

	if (copy != NULL)
		memcpy(copy, path, len);

	return copy;
}

struct fpk *fpk_open(const char *path)
{
	struct fpk *fpk;
	unsigned char header[FPK_HEADER_SIZE];
	uint64_t dir_offset, strtab_offset, strtab_size;
	uint64_t i;

	fpk = malloc(sizeof(*fpk));
	if (fpk == NULL)
		goto out;
	fpk->file = NULL;
	fpk->entries = NULL;
	fpk->strtab = NULL;

	fpk->file = fopen(path, "rb");
	if (fpk->file == NULL)
		goto fail;

	if (fread(header, 1, FPK_HEADER_SIZE, fpk->file) != FPK_HEADER_SIZE)
		goto fail;
	if (memcmp(header, FPK_MAGIC, FPK_MAGIC_SIZE) != 0)
		goto fail;

	fpk->entry_count = get_u64le(header + FPK_HEADER_FIELD_ENTRY_COUNT);
	dir_offset = get_u64le(header + FPK_HEADER_FIELD_DIR_OFFSET);
	strtab_offset = get_u64le(header + FPK_HEADER_FIELD_STRTAB_OFFSET);
	if (dir_offset < strtab_offset)
		goto fail;
	strtab_size = dir_offset - strtab_offset;

	fpk->strtab = malloc((size_t)strtab_size);
	if (fpk->strtab == NULL)
		goto fail;
	if (fseek(fpk->file, (long)strtab_offset, SEEK_SET) != 0)
		goto fail;
	if (fread(fpk->strtab, 1, (size_t)strtab_size, fpk->file) !=
	    strtab_size)
		goto fail;

	fpk->entries = malloc((size_t)fpk->entry_count * sizeof(fpk_entry_t));
	if (fpk->entries == NULL)
		goto fail;
	if (fseek(fpk->file, (long)dir_offset, SEEK_SET) != 0)
		goto fail;
	for (i = 0; i < fpk->entry_count; i++) {
		unsigned char rec[FPK_DIR_ENTRY_SIZE];
		uint64_t path_offset;

		if (fread(rec, 1, FPK_DIR_ENTRY_SIZE, fpk->file) !=
		    FPK_DIR_ENTRY_SIZE)
			goto fail;

		path_offset = get_u64le(rec + FPK_DIR_ENTRY_FIELD_PATH_OFFSET);
		if (path_offset >= strtab_size)
			goto fail;

		fpk->entries[i].path = fpk->strtab + path_offset;
		fpk->entries[i].data_offset =
			get_u64le(rec + FPK_DIR_ENTRY_FIELD_DATA_OFFSET);
		fpk->entries[i].data_size =
			get_u64le(rec + FPK_DIR_ENTRY_FIELD_DATA_SIZE);
		fpk->entries[i].method = rec[FPK_DIR_ENTRY_FIELD_METHOD];
	}
	goto out;

fail:
	fpk_close(fpk);
	fpk = NULL;
out:
	return fpk;
}

void fpk_close(struct fpk *fpk)
{
	if (fpk == NULL)
		return;
	if (fpk->file != NULL)
		fclose(fpk->file);
	free(fpk->entries);
	free(fpk->strtab);
	free(fpk);
}

fpk_entry_t *fpk_find(struct fpk *fpk, const char *path)
{
	uint64_t i;

	for (i = 0; i < fpk->entry_count; i++) {
		if (strcmp(fpk->entries[i].path, path) == 0)
			return &fpk->entries[i];
	}
	return NULL;
}

void *fpk_read(struct fpk *fpk, const fpk_entry_t *entry, size_t *out_size)
{
	void *buf;

	buf = malloc((size_t)entry->data_size);
	if (buf == NULL)
		goto out;

	if (fseek(fpk->file, (long)entry->data_offset, SEEK_SET) != 0 ||
	    fread(buf, 1, (size_t)entry->data_size, fpk->file) !=
		    entry->data_size) {
		free(buf);
		buf = NULL;
		goto out;
	}

	*out_size = (size_t)entry->data_size;
out:
	return buf;
}

uint64_t fpk_entry_count(const struct fpk *fpk)
{
	return fpk->entry_count;
}

const fpk_entry_t *fpk_entry_at(const struct fpk *fpk, uint64_t index)
{
	if (index >= fpk->entry_count)
		return NULL;
	return &fpk->entries[index];
}

#define FPK_WRITER_INITIAL_CAPACITY 16

struct fpk_writer *fpk_writer_open(const char *path)
{
	struct fpk_writer *w;
	unsigned char placeholder[FPK_HEADER_SIZE];

	w = malloc(sizeof(*w));
	if (w == NULL)
		goto out;
	w->entries = NULL;
	w->count = 0;
	w->capacity = 0;

	w->file = fopen(path, "wb");
	if (w->file == NULL)
		goto fail_free;

	memset(placeholder, 0, sizeof(placeholder));
	if (fwrite(placeholder, 1, FPK_HEADER_SIZE, w->file) !=
	    FPK_HEADER_SIZE)
		goto fail_close;
	w->cursor = FPK_HEADER_SIZE;
	goto out;

fail_close:
	fclose(w->file);
fail_free:
	free(w);
	w = NULL;
out:
	return w;
}

int fpk_writer_add(struct fpk_writer *w, const char *path, const void *data,
		    size_t size)
{
	struct fpk_writer_entry *entry;
	char *path_copy = NULL;
	int ret = 0;

	if (w->count == w->capacity) {
		uint64_t new_capacity = w->capacity == 0
						 ? FPK_WRITER_INITIAL_CAPACITY
						 : w->capacity * 2;
		struct fpk_writer_entry *new_entries = realloc(
			w->entries, (size_t)new_capacity * sizeof(*new_entries));

		if (new_entries == NULL) {
			ret = -1;
			goto out;
		}
		w->entries = new_entries;
		w->capacity = new_capacity;
	}

	path_copy = dup_path(path);
	if (path_copy == NULL) {
		ret = -1;
		goto out;
	}

	if (size > 0 && fwrite(data, 1, size, w->file) != size) {
		ret = -1;
		goto out;
	}

	entry = &w->entries[w->count];
	entry->path = path_copy;
	entry->data_offset = w->cursor;
	entry->data_size = size;
	entry->method = FPK_METHOD_STORE;

	w->cursor += size;
	w->count++;

out:
	if (ret != 0)
		free(path_copy);
	return ret;
}

int fpk_writer_finalize(struct fpk_writer *w)
{
	uint64_t strtab_offset, dir_offset, path_offset;
	unsigned char header[FPK_HEADER_SIZE];
	unsigned char rec[FPK_DIR_ENTRY_SIZE];
	uint64_t i;
	int ret = 0;

	strtab_offset = w->cursor;
	if (fseek(w->file, (long)strtab_offset, SEEK_SET) != 0) {
		ret = -1;
		goto out;
	}

	path_offset = 0;
	for (i = 0; i < w->count && ret == 0; i++) {
		size_t len = strlen(w->entries[i].path) + 1;

		if (fwrite(w->entries[i].path, 1, len, w->file) != len)
			ret = -1;
		else
			path_offset += len;
	}
	if (ret != 0)
		goto out;
	dir_offset = strtab_offset + path_offset;

	path_offset = 0;
	for (i = 0; i < w->count && ret == 0; i++) {
		put_u64le(rec + FPK_DIR_ENTRY_FIELD_PATH_OFFSET, path_offset);
		put_u64le(rec + FPK_DIR_ENTRY_FIELD_DATA_OFFSET,
			  w->entries[i].data_offset);
		put_u64le(rec + FPK_DIR_ENTRY_FIELD_DATA_SIZE,
			  w->entries[i].data_size);
		rec[FPK_DIR_ENTRY_FIELD_METHOD] = w->entries[i].method;

		if (fwrite(rec, 1, FPK_DIR_ENTRY_SIZE, w->file) !=
		    FPK_DIR_ENTRY_SIZE)
			ret = -1;
		else
			path_offset += strlen(w->entries[i].path) + 1;
	}
	if (ret != 0)
		goto out;

	memcpy(header, FPK_MAGIC, FPK_MAGIC_SIZE);
	put_u64le(header + FPK_HEADER_FIELD_ENTRY_COUNT, w->count);
	put_u64le(header + FPK_HEADER_FIELD_DIR_OFFSET, dir_offset);
	put_u64le(header + FPK_HEADER_FIELD_STRTAB_OFFSET, strtab_offset);

	if (fseek(w->file, 0, SEEK_SET) != 0 ||
	    fwrite(header, 1, FPK_HEADER_SIZE, w->file) != FPK_HEADER_SIZE ||
	    fflush(w->file) != 0)
		ret = -1;

out:
	return ret;
}

void fpk_writer_close(struct fpk_writer *w)
{
	uint64_t i;

	if (w == NULL)
		return;
	if (w->file != NULL)
		fclose(w->file);
	for (i = 0; i < w->count; i++)
		free(w->entries[i].path);
	free(w->entries);
	free(w);
}