/*
 *  upkg: tool for manipulating Unreal Tournament packages.
 *  Copyright © 2009-2011 Nick Bowler
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU 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 General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

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

#include "upkg.h"
#include "pack.h"

#define MIN(a, b) ((a) < (b) ? (a) : (b))

/*
 * Print a message and execute some statement(s) if the expression evaluates
 * to zero.  Intended to help verify that assumed constraints on the file
 * format actually are not violated.
 */
#define format_assert(expr, body) do { \
	if (!(expr)) { \
		fprintf(stderr, "%s: %d: %s: format assertion failed: %s\n", \
		                __FILE__, __LINE__, __func__, #expr); \
		body; \
	} \
} while (0)

struct upkg_name {
	unsigned long flags;
	char *name;
};

struct upkg_export_priv {
	struct upkg_export pub;

	long class, super;
	unsigned long size, offset;
};

struct upkg_import {
	const char *class_package, *class_name, *object_name;
	long package;
};

struct upkg_priv {
	struct upkg pub;

	const struct upkg_file_ops *fops;
	int (*dtor)(void *handle);
	void *f;

	struct upkg_file *last_file;

	struct upkg_name        *names;
	struct upkg_export_priv *exports;
	struct upkg_import      *imports;

	unsigned long name_offset, export_offset, import_offset;
	unsigned char guid[16];
};

/* Default I/O operations for ordinary files. */
static size_t file_read(void *buf, size_t size, void *handle)
{
	return fread(buf, 1, size, (FILE *)handle);
}

static int file_seek(void *handle, long offset, int whence)
{
	return fseek((FILE *)handle, offset, whence);
}

static long file_tell(void *handle)
{
	return ftell((FILE *)handle);
}

static int file_eof(void *handle)
{
	return feof((FILE *)handle);
}

static int file_close(void *handle)
{
	return fclose((FILE *)handle);
}

const struct upkg_file_ops upkg_default_fops = {
	.read = file_read,
	.seek = file_seek,
	.tell = file_tell,
	.eof  = file_eof,
};

/*
 * Decode the compact index format from the upkg.  This format is fucked.
 * Stores the result in *val and returns the number of input bytes read (or 0
 * if the input is invalid, in which case *val is undefined).
 */
size_t upkg_decode_index(long *val, const unsigned char *bytes, size_t n)
{
	*val = 0;

	for (size_t i = 0; i < MIN(n, 5); i++) {
		/*
		 * Least significant bytes are first, so we need to do this
		 * nonsense.
		 */
		long tmp = bytes[i] & (i == 0 ? 0x3f : 0x7f);

		if (i > 0) tmp <<= 6;
		if (i > 1) tmp <<= 7*(i-1);
		*val += tmp;

		if (!(bytes[i] & (i == 0 ? 0x40 : 0x80))) {
			if (bytes[0] & 0x80)
				*val = -*val;
			return i+1;
		}
	}

	/* Error */
	return 0;
}

static struct upkg_priv *init_upkg(unsigned char hdr[static UPKG_HDR_SIZE])
{
	struct upkg_priv *pkg;

	pkg = malloc(sizeof *pkg);
	if (!pkg)
		return NULL;

	*pkg = (struct upkg_priv) {
		.pub = {
			.version      = unpack_16_le(hdr+4),
			.license      = unpack_16_le(hdr+6),
			.flags        = unpack_32_le(hdr+8),
			.name_count   = unpack_32_le(hdr+12),
			.export_count = unpack_32_le(hdr+20),
			.import_count = unpack_32_le(hdr+28),
		},

		.name_offset   = unpack_32_le(hdr+16),
		.export_offset = unpack_32_le(hdr+24),
		.import_offset = unpack_32_le(hdr+32),
	};

	return pkg;
}

static int pkg_init_guid(struct upkg_priv *pkg)
{
	const struct upkg_file_ops *fops = pkg->fops;
	size_t rc;

	if (pkg->pub.version < 68) {
		unsigned long heritage_count, heritage_offset;
		unsigned char buf[8];

		rc = fops->read(buf, sizeof buf, pkg->f);
		if (rc < 8)
			return -1;

		heritage_count  = unpack_32_le(buf+0);
		heritage_offset = unpack_32_le(buf+4);

		if (heritage_count == 0)
			return -1;
		if (fops->seek(pkg->f, heritage_offset, SEEK_SET) != 0)
			return -1;
	}

	rc = fops->read(pkg->pub.guid, 16, pkg->f);
	if (rc < 16)
		return -1;

	return 0;
}

static int pkg_init_names(struct upkg_priv *pkg)
{
	const struct upkg_file_ops *fops = pkg->fops;
	void *f = pkg->f;

	size_t rc, len, nbuf = 0;
	unsigned long index = 0;
	unsigned char buf[512];

	if (fops->seek(f, pkg->name_offset, SEEK_SET) != 0)
		return -1;

	pkg->names = malloc(pkg->pub.name_count * sizeof *pkg->names);
	if (!pkg->names)
		return -1;

	while (index < pkg->pub.name_count) {
		struct upkg_name *name = &pkg->names[index];

		/* Read some data into buffer. */
		if (!fops->eof(pkg->f)) {
			rc = fops->read(buf+nbuf, sizeof buf-nbuf, f);
			if (rc == 0 && nbuf == 0)
				goto err;
			nbuf += rc;
		}

		if (pkg->pub.version >= 64) {
			len = buf[0];
			if (nbuf <= len + 4 || buf[len])
				goto err;
			name->name = malloc(len);
			if (!name->name)
				goto err;
			memcpy(name->name, buf+1, len);
			name->flags = unpack_32_le(buf+len+1);
			len += 4;

			nbuf -= len + 1;
			memmove(buf, buf+len+1, nbuf);
			index++;
		} else {
			unsigned char *c = memchr(buf, 0, nbuf);
			if (!c || nbuf <= c - buf + 5)
				goto err;
			len = c - buf + 1;
			name->name = malloc(len);
			if (!name->name)
				goto err;
			memcpy(name->name, buf, len);
			name->flags = unpack_32_le(buf+len);
			len += 4;

			nbuf -= len;
			memmove(buf, buf+len, nbuf);
			index++;
		}
	}

	return 0;
err:
	for (unsigned i = 0; i < index; i++)
		free(pkg->names[i].name);
	free(pkg->names);
	return -1;
}

static int pkg_init_exports(struct upkg_priv *pkg)
{
	const struct upkg_file_ops *fops = pkg->fops;
	void *f = pkg->f;

	size_t rc, len, nbuf = 0;
	unsigned long index = 0;
	unsigned char buf[512];

	if (fops->seek(f, pkg->export_offset, SEEK_SET) != 0)
		return -1;

	pkg->exports = malloc(pkg->pub.export_count * sizeof *pkg->exports);
	if (!pkg->exports)
		return -1;

	while (index < pkg->pub.export_count) {
		struct upkg_export_priv *export = &pkg->exports[index];
		long tmp;

		/* Read some data into buffer. */
		if (!fops->eof(pkg->f)) {
			rc = fops->read(buf+nbuf, sizeof buf-nbuf, f);
			if (rc == 0 && nbuf == 0)
				goto err;
			nbuf += rc;
		}

		len = 0;
		rc = upkg_decode_index(&export->class, buf+len, nbuf-len);
		if (rc == 0) goto err;
		len += rc;

		rc = upkg_decode_index(&export->super, buf+len, nbuf-len);
		if (rc == 0) goto err;
		len += rc;

		if (nbuf-len < 4) goto err;
		export->pub.package = unpack_s32_le(buf+len);
		len += 4;

		rc = upkg_decode_index(&tmp, buf+len, nbuf-len);
		if (rc == 0 || tmp < 0 || tmp >= pkg->pub.name_count) goto err;
		export->pub.name = pkg->names[tmp].name;
		len += rc;

		if (nbuf-len < 4) goto err;
		export->pub.flags = unpack_32_le(buf+len);
		len += 4;

		rc = upkg_decode_index(&tmp, buf+len, nbuf-len);
		if (rc == 0 || tmp < 0) goto err;
		export->size = tmp;
		len += rc;

		if (export->size) {
			rc = upkg_decode_index(&tmp, buf+len, nbuf-len);
			if (rc == 0 || tmp < 0) goto err;
			export->offset = tmp;
			len += rc;
		}

		nbuf -= len;
		memmove(buf, buf+len, nbuf);
		index++;
	}

	return 0;
err:
	free(pkg->exports);
	return -1;
}

static int pkg_init_imports(struct upkg_priv *pkg)
{
	const struct upkg_file_ops *fops = pkg->fops;
	void *f = pkg->f;

	size_t rc, len, nbuf = 0;
	unsigned long index = 0;
	unsigned char buf[512];

	if (fops->seek(f, pkg->import_offset, SEEK_SET) != 0)
		return -1;

	pkg->imports = malloc(pkg->pub.import_count * sizeof *pkg->imports);
	if (!pkg->imports)
		return -1;

	while (index < pkg->pub.import_count) {
		struct upkg_import *import = &pkg->imports[index];
		long tmp;

		/* Read some data into buffer. */
		if (!fops->eof(pkg->f)) {
			rc = fops->read(buf+nbuf, sizeof buf-nbuf, f);
			if (rc == 0 && nbuf == 0)
				goto err;
			nbuf += rc;
		}

		len = 0;
		rc = upkg_decode_index(&tmp, buf+len, nbuf-len);
		if (rc == 0 || len >= pkg->pub.name_count) goto err;
		import->class_package = pkg->names[tmp].name;
		len += rc;

		rc = upkg_decode_index(&tmp, buf+len, nbuf-len);
		if (rc == 0 || len >= pkg->pub.name_count) goto err;
		import->class_name = pkg->names[tmp].name;
		len += rc;

		if (nbuf-len < 4) goto err;
		import->package = unpack_s32_le(buf+len);
		len += 4;

		rc = upkg_decode_index(&tmp, buf+len, nbuf-len);
		if (rc == 0 || len >= pkg->pub.name_count) goto err;
		import->object_name = pkg->names[tmp].name;
		len += rc;

		nbuf -= len;
		memmove(buf, buf+len, nbuf);
		index++;
	}

	return 0;
err:
	free(pkg->imports);
	return -1;
}

struct upkg *upkg_open(void *f, const struct upkg_file_ops *fops,
                       int (*destructor)(void *handle))
{
	unsigned char hdr_buf[UPKG_HDR_SIZE];
	struct upkg_priv *pkg;

	if (fops->read(hdr_buf, sizeof hdr_buf, f) != sizeof hdr_buf) {
		return NULL;
	}
	if (unpack_32_le(hdr_buf) != UPKG_HDR_MAGIC) {
		return NULL;
	}

	/* Initialize package structure. */
	pkg = init_upkg(hdr_buf);
	if (!pkg) {
		return NULL;
	}
	pkg->fops = fops;
	pkg->dtor = destructor;
	pkg->f    = f;

	if (pkg_init_guid(pkg) != 0) {
		goto err1;
	}

	if (pkg_init_names(pkg) != 0) {
		goto err1;
	}

	if (pkg_init_exports(pkg) != 0) {
		goto err2;
	}

	if (pkg_init_imports(pkg) != 0) {
		goto err3;
	}

	return &pkg->pub;
err3:
	free(pkg->exports);
err2:
	for (unsigned i = 0; i < pkg->pub.name_count; i++)
		free(pkg->names[i].name);
	free(pkg->names);
err1:
	free(pkg);
	return NULL;
}

struct upkg *upkg_fopen(const char *path)
{
	struct upkg *pkg;
	FILE *f;

	f = fopen(path, "rb");
	if (!f) {
		return NULL;
	}

	pkg = upkg_open(f, &upkg_default_fops, file_close);
	if (!pkg) {
		fclose(f);
	}

	return pkg;
}

int upkg_close(struct upkg *pub)
{
	struct upkg_priv *pkg = (struct upkg_priv *)pub;
	int rc = 0;

	if (pkg->dtor) {
		rc = pkg->dtor(pkg->f);
	}

	for (unsigned i = 0; i < pkg->pub.name_count; i++) {
		free(pkg->names[i].name);
	}

	free(pkg->imports);
	free(pkg->exports);
	free(pkg->names);
	free(pkg);

	return rc;
}

const char *upkg_get_name(struct upkg *pub, unsigned long idx)
{
	struct upkg_priv *pkg = (struct upkg_priv *)pub;

	if (idx >= pkg->pub.name_count)
		return 0;
	return pkg->names[idx].name;
}

long upkg_export_find(struct upkg *pub, long parent, const char *name)
{
	struct upkg_priv *pkg = (struct upkg_priv *)pub;

	/* This only makes sense if the assertion below is not violated. */
	long package = parent < 0 ? 0 : parent + 1;

	for (unsigned long i = 0; i < pkg->pub.export_count; i++) {
		struct upkg_export_priv *e = &pkg->exports[i];

		/* Assertion: an object's package is an export. */
		format_assert(e->pub.package >= 0, continue);
		if (e->pub.package == package
		    && strcmp(e->pub.name, name) == 0) {
			return i;
		}
	}

	return -1;
}

const struct upkg_export *upkg_get_export(struct upkg *pub, unsigned long idx)
{
	struct upkg_priv *pkg = (struct upkg_priv *)pub;

	if (idx < pkg->pub.export_count)
		return &pkg->exports[idx].pub;
	return NULL;
}

const char *upkg_export_class(struct upkg *pub, unsigned long idx,
                              const char **package)
{
	struct upkg_priv *pkg = (struct upkg_priv *)pub;
	struct upkg_export_priv *export;
	struct upkg_import *iclass, *ipackage;
	unsigned long pkg_idx;

	if (idx >= pkg->pub.export_count)
		return NULL;

	export = &pkg->exports[idx];

	/* Assumption: class references are always imports. */
	format_assert(export->class <= 0, return NULL);

	/* Get the class. */
	if (export->class == 0) {
		if (package) *package = "Core";
		return "Class";
	}

	pkg_idx = -(export->class + 1);
	if (pkg_idx >= pkg->pub.import_count)
		return NULL;
	iclass = &pkg->imports[pkg_idx];

	/* Assumption: class references are always Core.Class. */
	format_assert(!strcmp(iclass->class_package, "Core"), return NULL);
	format_assert(!strcmp(iclass->class_name, "Class"), return NULL);

	/* Assumption: package references are always imports. */
	format_assert(iclass->package <= 0, return NULL);

	/* Get the package. */
	pkg_idx = -(iclass->package + 1);
	if (pkg_idx >= pkg->pub.import_count)
		return NULL;
	ipackage = &pkg->imports[pkg_idx];

	/* Assumption: package references are always Core.Package. */
	format_assert(!strcmp(ipackage->class_package, "Core"), return NULL);
	format_assert(!strcmp(ipackage->class_name, "Package"), return NULL);

	if (package) *package = ipackage->object_name;
	return iclass->object_name;
}

struct upkg_file *upkg_export_open(struct upkg *pub, unsigned long idx)
{
	struct upkg_priv *pkg = (struct upkg_priv *)pub;
	struct upkg_file *f;

	if (idx >= pkg->pub.export_count)
		return NULL;

	f = malloc(sizeof *f);
	if (f == NULL)
		return NULL;

	*f = (struct upkg_file) {
		.pkg  = pkg,
		.base = pkg->exports[idx].offset,
		.len  = pkg->exports[idx].size,
		.name = pkg->exports[idx].pub.name,
	};

	return f;
}

void upkg_export_close(struct upkg_file *f)
{
	if (f->pkg->last_file == f)
		f->pkg->last_file = NULL;
	free(f);
}

long upkg_export_tell(struct upkg_file *f)
{
	return f->offset;
}

int upkg_export_seek(struct upkg_file *f, long offset, int whence)
{
	const struct upkg_file_ops *fops = f->pkg->fops;
	int rc = EOF;

	switch (whence) {
	case SEEK_CUR:
		offset = f->offset + offset;
	case SEEK_SET:
		if (offset < 0 || offset > f->len)
			return EOF;
		rc = fops->seek(f->pkg->f, f->base + offset, SEEK_SET);
		break;
	case SEEK_END:
		offset = -offset;
		if (offset < 0 || offset > f->len)
			return EOF;
		offset = f->len - offset;
		rc = fops->seek(f->pkg->f, f->base + offset, SEEK_SET);
		break;
	}

	if (rc == 0) {
		f->pkg->last_file = f;
		f->offset = offset;
		f->eof = 0;
	} else if (f->pkg->last_file == f) {
		f->pkg->last_file = NULL;
	}

	return rc;
}

size_t upkg_export_read(struct upkg_file *f, void *buf, size_t n)
{
	const struct upkg_file_ops *fops = f->pkg->fops;
	size_t want = MIN(n, f->len - f->offset);
	size_t rc;

	if (want == 0) {
		return 0;
	}

	if (f != f->pkg->last_file) {
		if (fops->seek(f->pkg->f, f->base + f->offset, SEEK_SET))
			return 0;
	}

	rc = fops->read(buf, want, f->pkg->f);
	f->offset += rc;

	if (want < n || (rc < want && fops->eof(f->pkg->f)))
		f->eof = 1;
	return rc;
}