#include <merror.h>
#include <elf.h>
#include <melf.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/mman.h>

#include "sim.h"

#define SEC_ALIGN 0x1000

struct load_state {
	FILE *file;
	int fd;
	uint32_t file_sz;

	Elf32_Phdr *phdr;
	uint32_t phdr_len;
	Elf32_Ehdr ehdr;
};

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

static int load_ehdr(struct simulator *sim, struct load_state *state)
{
	Elf32_Ehdr ehdr;
	fseek(state->file, 0, SEEK_SET);
	if (fread(&ehdr, sizeof(Elf32_Ehdr), 1, state->file) != 1) {
		ERROR("cannot load 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.
	 */

	Elf32_Ehdr baseline = MIPS_ELF_EHDR;
	baseline.e_type = B16(ET_EXEC);

	#define EHDR_ASSERT(name, size)                          \
			if (res == M_SUCCESS)                    \
			  res |= assert_ehdr(&baseline.e_##name, \
			  &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;

	state->ehdr = ehdr;
	sim->entry = B32(ehdr.e_entry);

	return M_SUCCESS;
}

static int load_phdr(struct load_state *state)
{
	Elf32_Phdr *phdr = NULL;

	uint32_t off = B32(state->ehdr.e_phoff);
	uint16_t len = B16(state->ehdr.e_phnum);

	if (BOUND_CHK(state->file_sz, len, off)) {
		ERROR("elf phdr location invalid");
		return M_ERROR;
	}

	phdr = malloc(sizeof(Elf32_Phdr) * len);
	if (phdr == NULL) {
		PERROR("cannot alloc");
		return M_ERROR;
	}

	fseek(state->file, off, SEEK_SET);
	if (fread(phdr, sizeof(Elf32_Phdr), len, state->file) != len) {
		free(phdr);
		PERROR("cannot read phdr");
		return M_ERROR;
	}

	state->phdr = phdr;
	state->phdr_len = len;
	return M_SUCCESS;
}

static int load_segment(struct simulator *sim,
			struct load_state *state, Elf32_Phdr *hdr)
{
	uint32_t off = B32(hdr->p_offset);
	uint32_t len = B32(hdr->p_filesz);
	uint32_t is_text = B32(hdr->p_flags) & PF_X;

	if (BOUND_CHK(state->file_sz, len, off)) {
		ERROR("segment location invalid");
		return M_ERROR;
	}

	uintptr_t addr = 0;
	uint32_t add = 0;

	if (len % SEC_ALIGN) {
		add = SEC_ALIGN - (len % SEC_ALIGN);
		len += add;
	}

	if (is_text) {
		addr = sim->text_max;
		sim->text_max += len;
	} else {
		addr = sim->data_max;
		sim->data_max += len;
	}

	bool read = B32(hdr->p_flags) & PF_R;
	bool write = B32(hdr->p_flags) & PF_W;

	uint32_t prot = 0;
	if (read)
		prot |= PROT_READ;
	if (write)
		prot |= PROT_WRITE;

	void *res = mmap((void*)addr, len, prot, MAP_PRIVATE | MAP_FIXED,
		  state->fd, off);

	if ((uintptr_t)res != addr) {
		PERROR("failed to map executable");
		return M_ERROR;
	}

	return M_SUCCESS;
}

static int load_memory(struct simulator *sim, struct load_state *state)
{
	uint32_t base = 0;

	for (uint32_t i = 0; i < state->phdr_len; i++) {
		Elf32_Phdr *hdr = NULL;
		uint32_t min = UINT32_MAX;

		if (B32(state->phdr[i].p_filesz) < 1)
			continue;

		// we need to load segments in order
		for (uint32_t j = 0; j < state->phdr_len; j++) {
			Elf32_Phdr *temp = &state->phdr[j];
			uint32_t off = B32(temp->p_offset);
			uint32_t len = B32(temp->p_filesz);

			if (len < 1)
				continue;

			if (off <= base)
				continue;

			if (off >= min)
				continue;

			min = off;
			hdr = temp;
		}

		base = min;

		if (hdr == NULL) {
			ERROR("invalid elf phdr");
			return M_ERROR;
		}

		if (load_segment(sim, state, hdr))
			return M_ERROR;
	}

	return M_SUCCESS;
}

int sim_load_file(struct simulator *sim) {
	struct load_state state;

	state.file = fopen(sim->args->executable, "r");
	if (state.file == NULL) {
		PERROR("cannot read '%s'", state.file);
		return M_ERROR;
	}
	state.fd = fileno(state.file);

	// get filesize
	fseek(state.file, 0, SEEK_END);
	state.file_sz = ftell(state.file);
	state.phdr = NULL;

	int res = M_SUCCESS;
	if (res == M_SUCCESS)
		res = load_ehdr(sim, &state);
	if (res == M_SUCCESS)
		res = load_phdr(&state);
	if (res == M_SUCCESS)
		res = load_memory(sim, &state);

	if (state.phdr)
		free(state.phdr);
	fclose(state.file);

	return res;
}