#include <elf.h>
#include <merror.h>
#include <melf.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>

#include "link.h"

extern char *current_file;

static int obj_assert_ehdr(const void *given, const void *assert, size_t size)
{
	if (memcmp(given, assert, size) == 0)
		return M_SUCCESS;
	ERROR("invalid elf file");
	return M_ERROR;
}

/**
 * Map the ehdr
 */
static int load_ehdr(struct object *object)
{
	uint32_t off = 0;
	object->ehdr = (Elf32_Ehdr *) (object->mapped + off);

	if (BOUND_CHK(object, sizeof(Elf32_Ehdr), off)) {
		ERROR("cannot read ehdr");
		return M_ERROR;
	}

	/**
	 * Compare each "static" value in the ehdr and make sure it
	 * is what it should be. If not throw and error and eventually
	 * return.
	 */

	#define EHDR_ASSERT(name, size)                                   \
			if (res == M_SUCCESS)                             \
			  res |= obj_assert_ehdr(&MIPS_ELF_EHDR.e_##name, \
			  &object->ehdr->e_##name, size)                  \

	int res = 0;
	EHDR_ASSERT(ident,     EI_NIDENT);
	EHDR_ASSERT(type,      sizeof(Elf32_Half));
	EHDR_ASSERT(machine,   sizeof(Elf32_Half));
	EHDR_ASSERT(version,   sizeof(Elf32_Word));
	EHDR_ASSERT(flags,     sizeof(Elf32_Word));
	EHDR_ASSERT(ehsize,    sizeof(Elf32_Half));
	EHDR_ASSERT(phentsize, sizeof(Elf32_Half));
	EHDR_ASSERT(shentsize, sizeof(Elf32_Half));

	#undef EHDR_ASSERT

	if (res)
		return M_ERROR;

	return M_SUCCESS;
}

/**
 * Map the shdr
 */
static int load_shdr(struct object *object)
{
	size_t shdr_len = B16(object->ehdr->e_shentsize) *
		B16(object->ehdr->e_shnum);
	size_t shdr_off = B32(object->ehdr->e_shoff);
	object->shdr = (Elf32_Shdr *) (object->mapped + shdr_off);
	object->shdr_len = B16(object->ehdr->e_shnum);

	if (BOUND_CHK(object, shdr_len, shdr_off)) {
		ERROR("cannot read shdr");
		return M_ERROR;
	}

	return M_SUCCESS;
}

/**
 * Map the phdr
 */
static int load_phdr(struct object *object)
{
	size_t phdr_len = B16(object->ehdr->e_phentsize) *
		B16(object->ehdr->e_phnum);
	size_t phdr_off = B32(object->ehdr->e_phoff);
	object->phdr = (Elf32_Phdr *) (object->mapped + phdr_off);
	object->phdr_len = B16(object->ehdr->e_phnum);

	if (BOUND_CHK(object, phdr_len, phdr_off)) {
		ERROR("cannot read phdr");
		return M_ERROR;
	}

	return M_SUCCESS;
}

/**
 * Load the strtabs
 */
static int load_strtabs(struct object *object)
{
	uint32_t max_entries = object->shdr_len;
	struct string_table *strtabs = malloc(max_entries *
					sizeof(struct string_table));

	if (strtabs == NULL) {
		PERROR("cannot alloc");
		return M_ERROR;
	}

	for (size_t i = 0; i < object->shdr_len; i++) {
		Elf32_Shdr *hdr = &object->shdr[i];
		struct string_table *strtab = &strtabs[i];

		if (B32(hdr->sh_type) != SHT_STRTAB) {
			strtab->len = 0;
			continue;
		}

		uint32_t off = B32(hdr->sh_offset);
		uint32_t len = B32(hdr->sh_size);

		strtab->len = len;
		strtab->data = object->mapped + off;

		if (len < 1) {
			ERROR("invalid or empty strtab [%d]", i);
			return M_ERROR;
		}

		if (BOUND_CHK(object, len, off)) {
			ERROR("cannot map strtab [%d]", i);
			return M_ERROR;
		}

		if (strtab->data[0] != '\0' ||
			strtab->data[len - 1] != '\0') {
			ERROR("strtab [%d] doesn't not have bounding null "
				"values", i);
			return M_ERROR;
		}
	}

	object->strtabs = strtabs;
	return M_SUCCESS;
}

/**
 * Load the symtabs
 */
static int load_symtabs(struct object *object)
{
	uint32_t max_entries = object->shdr_len;
	struct symbol_table *symtabs = malloc(max_entries *
					sizeof(struct symbol_table));

	if (symtabs == NULL) {
		PERROR("cannot alloc");
		return M_ERROR;
	}

	for (size_t i = 0; i < object->shdr_len; i++) {
		Elf32_Shdr *hdr = &object->shdr[i];
		struct symbol_table *symtab = &symtabs[i];

		if (B32(hdr->sh_type) != SHT_SYMTAB) {
			symtab->len = 0;
			continue;
		}

		uint32_t off = B32(hdr->sh_offset);
		uint32_t len = B32(hdr->sh_size);
		uint32_t stridx = B32(hdr->sh_link);

		if (stridx >= max_entries) {
			ERROR("strtab index [%d] out of bounds", stridx);
			return M_ERROR;
		}

		struct string_table *strtab = &object->strtabs[stridx];
		if (strtab->len < 1) {
			ERROR("strtab index [%d] empty or invalid", stridx);
			return M_ERROR;
		}

		symtab->strtab = strtab;
		symtab->len = len;
		symtab->syms = (Elf32_Sym *) (object->mapped + off);

		if (BOUND_CHK(object, len, off)) {
			ERROR("cannot map symtab index [%d]", i);
			return M_ERROR;
		}
	}

	object->symtabs = symtabs;
	return M_SUCCESS;
}

/**
 * Load the shstrtab
 */
static int load_shstrtab(struct object *object)
{
	uint16_t idx = B16(object->ehdr->e_shstrndx);
	if (idx >= object->shdr_len) {
		ERROR("shstrndx [%d] out of bounds", idx);
		return M_ERROR;
	}

	struct string_table *strtab = &object->strtabs[idx];
	if (strtab->len < 1) {
		ERROR("shstrndx is invalid or empty");
		return M_ERROR;
	}

	object->shstrtab = strtab;
	return M_SUCCESS;
}

/**
 * Load the program segments
 */
static int load_segments(struct object *object)
{
	object->segment_len = B16(object->ehdr->e_phnum);
	object->segments = malloc(sizeof(struct segment) *
				object->segment_len);

	if (object->segments == NULL) {
		PERROR("cannot alloc");
		return M_ERROR;
	}

	for (size_t i = 0; i < object->segment_len; i++) {
		struct segment *seg = &object->segments[i];
		if (segment_load(object, seg, i)) {
			return M_ERROR;
		}
	}

	return M_SUCCESS;
}

static int map_file(struct object *obj, char *path)
{
	obj->fd = open(path, O_RDONLY);
	if (obj->fd == -1) {
		PERROR("cannot read");
		return M_ERROR;
	}

	struct stat st;
	if (stat(path, &st)) {
		PERROR("cannot stat");
		return M_ERROR;
	}

	obj->mapped_size = st.st_size;
	obj->mapped = mmap(NULL, st.st_size, PROT_READ, MAP_SHARED,
				obj->fd, 0);

	if (obj->mapped == MAP_FAILED) {
		PERROR("cannot map");
		return M_ERROR;
	}

	return M_SUCCESS;
}

int object_load(struct object *object, char *path)
{
	current_file = path;

	object->fd = 0;
	object->segments = NULL;
	object->symtabs = NULL;
	object->strtabs = NULL;
	object->mapped = NULL;
	object->name = path;

	/** load the file */
	if (map_file(object, path))
		return M_ERROR;

	/* ehdr */
	if (load_ehdr(object))
		return M_ERROR;

	/* shdr */
	if (load_shdr(object))
		return M_ERROR;

	/* phdr */
	if (load_phdr(object))
		return M_ERROR;

	/* strtabs */
	if (load_strtabs(object))
		return M_ERROR;

	/* symtabs */
	if (load_symtabs(object))
		return M_ERROR;

	/* shstrtab */
	if (load_shstrtab(object))
		return M_ERROR;

	/* segments */
	if (load_segments(object))
		return M_ERROR;

	return M_SUCCESS;
}

void object_free(struct object *obj)
{
	if (obj->symtabs != NULL)
		free(obj->symtabs);
	if (obj->strtabs != NULL)
		free(obj->strtabs);
	if (obj->segments != NULL)
		free(obj->segments);
	if (obj->fd > 0)
		close(obj->fd);
	if (obj->mapped != NULL && obj->mapped != MAP_FAILED)
		munmap(obj->mapped, obj->mapped_size);
}