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

#include "link.h"

/**
 * Read the phdr for this segment
 */
static int load_phdr(struct object *obj, struct segment *seg, size_t index)
{
	if (index >= obj->phdr_len) {
		ERROR("phdr index [%d] out of bounds", index);
		return M_ERROR;
	}

	seg->phdr = &obj->phdr[index];
	seg->phdr_idx = index;
	return M_SUCCESS;
}

/**
 * Read the shdr for this segment
 */
static int load_shdr(struct object *obj, struct segment *seg, size_t index)
{
	uint32_t shndx = obj->phdr_to_shdr_mapping[index];
	Elf32_Shdr *hdr = &obj->shdr[shndx];

	// get name
	uint32_t name = B32(hdr->sh_name);
	if (name >= obj->shstrtab->len) {
		ERROR("section name index [%d] out of bounds", name);
		return M_ERROR;
	}
	seg->name = &obj->shstrtab->data[name];

	// map bytes
	uint32_t len = B32(hdr->sh_size);
	uint32_t off = B32(hdr->sh_offset);
	seg->bytes = (unsigned char *) (obj->mapped + off);

	if (BOUND_CHK(obj, len, off)) {
		ERROR("cannot map seg %s:%d", seg->name, index);
		return M_ERROR;
	}

	seg->shdr = hdr;
	seg->shdr_idx = shndx;

	return M_SUCCESS;
}

/**
 * Read the relocation table for this segment (if one exists)
 */
static int load_reltab(struct object *obj, struct segment *seg)
{
	// default to none
	seg->reltab.len = 0;

	for (size_t i = 0; i < obj->shdr_len; i++) {
		Elf32_Shdr *hdr = &obj->shdr[i];
		uint32_t type = B32(hdr->sh_type);

		if (type != SHT_REL && type != SHT_RELA)
			continue;

		if ((type == SHT_REL && B32(hdr->sh_entsize) !=
			sizeof(Elf32_Rel)) || (type == SHT_RELA &&
			B32(hdr->sh_entsize) != sizeof(Elf32_Rela))) {
			ERROR("reltab [%d] has invalid entry size", i);
			return M_ERROR;
		}

		if (B32(hdr->sh_info) != seg->shdr_idx)
			continue;

		uint32_t symtab_idx = B32(hdr->sh_link);
		if (symtab_idx >= obj->shdr_len) {
			ERROR("reltab [%d] symtab index [%d] out of bounds",
				i, symtab_idx);
			return M_ERROR;
		}

		struct symbol_table *symtab = &obj->symtabs[symtab_idx];
		if (symtab->len < 1) {
			ERROR("reltab [%d] symtab is empty or invalid", i);
			return M_ERROR;
		}

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

		seg->reltab.symtab = symtab;
		seg->reltab.len = len / B32(hdr->sh_entsize);
		seg->reltab.raw = obj->mapped + off;
		seg->reltab.type = type;

		if (BOUND_CHK(obj, len, off)) {
			ERROR("cannot map reltab [%d] for %s",
				i, seg->name);
			return M_ERROR;
		}

		break;
	}

	return M_SUCCESS;
}

int segment_load(struct object *obj, struct segment *seg, size_t index)
{

	if (load_phdr(obj, seg, index))
		return M_ERROR;

	if (load_shdr(obj, seg, index))
		return M_ERROR;

	if (load_reltab(obj, seg))
		return M_ERROR;

	seg->off = B32(seg->phdr->p_offset);
	seg->vaddr = B32(seg->phdr->p_vaddr);
	seg->size = B32(seg->phdr->p_filesz);
	seg->new_off = 0;
	seg->new_vaddr = 0;

	seg->read = B32(seg->phdr->p_flags) & PF_R;
	seg->write = B32(seg->phdr->p_flags) & PF_W;
	seg->execute = B32(seg->phdr->p_flags) & PF_X;
	seg->align = B32(seg->phdr->p_align);

	seg->obj = obj;
	seg->ent = NULL;

	return M_SUCCESS;
}