#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
#define PAGE_SIZE 4096

#define BITFILED_LEN (UINT32_MAX / PAGE_SIZE / 8)

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)                  \

	// ignore abi ver
	ehdr.e_ident[EI_ABIVERSION] = 0x00;

	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 void set_page(uint8_t *bitfield, uint32_t addr)
{
	int idx = (addr / PAGE_SIZE) / 8;
	int off = (addr / PAGE_SIZE) % 8;

	bitfield[idx] |= 1 << off;
}

static int get_page(const uint8_t *const bitfield, uint32_t addr)
{
	int idx = (addr / PAGE_SIZE) / 8;
	int off = (addr / PAGE_SIZE) % 8;

	return ((bitfield[idx] >> off) & 1);
}

static int load_segment(struct simulator *sim, struct load_state *state,
			Elf32_Phdr *hdr, uint8_t* bitfield)
{
	uint32_t addr = B32(hdr->p_vaddr),
	         off = B32(hdr->p_offset),
	         len = B32(hdr->p_filesz),
	         flags = B32(hdr->p_flags);

	bool exec = flags & PF_X;

	// ignore if empty
	if (len < 1)
		return M_SUCCESS;

	// make sure segment is acutally inside
	// the file
	if (BOUND_CHK(state->file_sz, len, off)) {
		ERROR("segment location invalid");
		return M_ERROR;
	}

	// make sure the vitural address is also
	// valid
	if (BOUND_CHK(UINT32_MAX, len, addr)) {
		ERROR("segment vitural addr invalid");
		return M_ERROR;
	}

	// update text seg bounds
	if (exec) {
		if (addr < sim->text_min)
			sim->text_min = addr;
		if (addr + len > sim->text_max)
			sim->text_max = addr + len;
	}

	// align the mapping ptr to
	// the page size
	uintptr_t ptr = (addr / PAGE_SIZE) * PAGE_SIZE;

	// map each page that the segment
	// requires
	for (; ptr < addr + len; ptr += PAGE_SIZE) {
		// dont remap if address is
		// already mapped
		if (get_page(bitfield, ptr))
			continue;

		// set page as mapped
		set_page(bitfield, ptr);

		void *res = mmap(
			(void *) ptr, PAGE_SIZE,
			PROT_READ | PROT_WRITE,
			MAP_PRIVATE | MAP_FIXED | MAP_ANONYMOUS, -1, 0);

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

	}

	// load the segment into the mapped memory
	fseek(state->file, off, SEEK_SET);
	fread((void *) (uintptr_t) addr, 1, len, state->file);

	return M_SUCCESS;
}

static int load_memory(struct simulator *sim, struct load_state *state)
{
	// map each page in a 32bit address space to a single bit
	// in the bitfield
	uint8_t *bitfield = malloc(BITFILED_LEN);
	if (bitfield == NULL) {
		PERROR("cannot alloc");
		return M_ERROR;
	}
	memset(bitfield, 0, BITFILED_LEN);

	for (uint32_t i = 0; i < state->phdr_len; i++) {
		Elf32_Phdr *hdr = &state->phdr[i];

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

	free(bitfield);

	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;
}