#include "elf.h"

#include <mango/config.h>
#include <mango/handle.h>
#include <mango/log.h>
#include <mango/vm.h>
#include <stdbool.h>
#include <stdio.h>
#include <string.h>

#define MAX(x, y)          ((x) > (y) ? (x) : (y))
#define MIN(x, y)          ((x) < (y) ? (x) : (y))

#define NEEDS_NOTHING      0
#define NEEDS_VDSO         1
#define NEEDS_MORE         2

#define ACL                (PF_R | PF_W | PF_X)
#define ACCESS(x)          ((x) & ACL)

/* TODO in case we ever support ELF32 images */
#define elf_class_bits(x)  (64)

#define PAGE_SIZE          (image->e_page_size)
#define PAGE_MASK          (image->e_page_size - 1)
#define PAGE_OFFSET(v)     ((v) & (PAGE_SIZE - 1))
#define PAGE_ALIGN_DOWN(v) (v) &= ~(PAGE_SIZE - 1)
#define PAGE_ALIGN_UP(v)                                                       \
	do {                                                                   \
		if ((v) & (PAGE_SIZE - 1)) {                                   \
			v &= ~(PAGE_SIZE - 1);                                 \
			v += PAGE_SIZE;                                        \
		}                                                              \
	} while (0)

#undef DEBUG_LOG

static enum launch_status elf_validate_ehdr(elf_ehdr_t *hdr)
{
	if (hdr->e_ident[EI_MAG0] != ELF_MAG0) {
		return LAUNCH_ERR_INVALID_EXECUTABLE;
	}

	if (hdr->e_ident[EI_MAG1] != ELF_MAG1) {
		return LAUNCH_ERR_INVALID_EXECUTABLE;
	}

	if (hdr->e_ident[EI_MAG2] != ELF_MAG2) {
		return LAUNCH_ERR_INVALID_EXECUTABLE;
	}

	if (hdr->e_ident[EI_MAG3] != ELF_MAG3) {
		return LAUNCH_ERR_INVALID_EXECUTABLE;
	}

	if (hdr->e_ident[EI_CLASS] != ELFCLASS64) {
		return LAUNCH_ERR_UNSUPPORTED_EXECUTABLE;
	}

	if (hdr->e_machine != EM_X86_64) {
		return LAUNCH_ERR_UNSUPPORTED_EXECUTABLE;
	}

	if (hdr->e_ident[EI_DATA] != ELFDATA2LSB) {
		return LAUNCH_ERR_UNSUPPORTED_EXECUTABLE;
	}

	if (hdr->e_ident[EI_VERSION] != EV_CURRENT) {
		return LAUNCH_ERR_UNSUPPORTED_EXECUTABLE;
	}

	return LAUNCH_OK;
}

static enum launch_status read_header(struct elf_image *image)
{
	size_t nr_read = 0;
	vm_object_read(
		image->e_image,
		&image->e_hdr,
		0,
		sizeof image->e_hdr,
		&nr_read);
	if (nr_read != sizeof image->e_hdr) {
		return LAUNCH_ERR_INVALID_EXECUTABLE;
	}

	return elf_validate_ehdr(&image->e_hdr);
}

static enum launch_status parse_phdr(struct elf_image *image)
{
	elf_phdr_t phdr;
	size_t r = 0;
	image->e_total_size = 0;
	image->e_data_size = 0;
	off_t vaddr, vlimit;

	for (size_t i = 0; i < image->e_hdr.e_phnum; i++) {
		off_t offset
			= image->e_hdr.e_phoff + (i * image->e_hdr.e_phentsize);
		kern_status_t status = vm_object_read(
			image->e_image,
			&phdr,
			offset,
			sizeof phdr,
			&r);

		if (status != KERN_OK || r != sizeof phdr) {
			return LAUNCH_ERR_INVALID_EXECUTABLE;
		}

		vaddr = phdr.p_vaddr;
		vlimit = phdr.p_vaddr + phdr.p_memsz;
		if (vaddr & (PAGE_SIZE - 1)) {
			vaddr &= ~(PAGE_SIZE - 1);
		}

		if (vlimit & (PAGE_SIZE - 1)) {
			vlimit &= ~(PAGE_SIZE - 1);
			vlimit += PAGE_SIZE;
		}

		switch (phdr.p_type) {
		case PT_DYNAMIC:
			image->e_dynamic = phdr;
			break;
		case PT_LOAD:

			image->e_total_size = MAX(image->e_total_size, vlimit);
			break;
		case PT_INTERP: {
			size_t r = 0;
			vm_object_read(
				image->e_image,
				image->e_interp,
				phdr.p_offset,
				MIN(sizeof image->e_interp - 1, phdr.p_filesz),
				&r);
			image->e_interp[r] = 0;
			break;
		}
		default:
			break;
		}

		if (phdr.p_flags & PF_W) {
			image->e_data_size = MAX(image->e_data_size, vlimit);
		}
	}

	return LAUNCH_OK;
}

static kern_status_t create_exec_regions(struct elf_image *image)
{
	kern_tracef(
		"launch: executable region size=%zu %zx",
		image->e_total_size,
		image->e_total_size);
	kern_status_t status = KERN_OK;
	if (image->e_local_space != KERN_HANDLE_INVALID) {
		status = vm_region_create(
			image->e_local_space,
			NULL,
			0,
			VM_REGION_ANY_OFFSET,
			image->e_total_size,
			VM_PROT_READ | VM_PROT_WRITE | VM_PROT_EXEC
				| VM_PROT_USER,
			&image->e_local_exec,
			&image->e_local_base);
	}

	if (status != KERN_OK) {
		return status;
	}

	if (image->e_remote_space != KERN_HANDLE_INVALID) {
		status = vm_region_create(
			image->e_remote_space,
			NULL,
			0,
			VM_REGION_ANY_OFFSET,
			image->e_total_size,
			VM_PROT_READ | VM_PROT_WRITE | VM_PROT_EXEC
				| VM_PROT_USER,
			&image->e_remote_exec,
			&image->e_remote_base);
	}

	if (status != KERN_OK) {
		vm_region_kill(image->e_local_exec);
		kern_handle_close(image->e_local_exec);
		image->e_local_exec = KERN_HANDLE_INVALID;

		return status;
	}

	return KERN_OK;
}

static enum launch_status map_executable(struct elf_image *image)
{
	elf_phdr_t phdr;
	size_t r = 0;
	image->e_total_size = 0;
	image->e_data_size = 0;

	size_t data_offset = 0;

	for (size_t i = 0; i < image->e_hdr.e_phnum; i++) {
		off_t phdr_offset
			= image->e_hdr.e_phoff + (i * image->e_hdr.e_phentsize);
		kern_status_t status = vm_object_read(
			image->e_image,
			&phdr,
			phdr_offset,
			sizeof phdr,
			&r);

		if (status != KERN_OK || r != sizeof phdr) {
			return LAUNCH_ERR_INVALID_EXECUTABLE;
		}

		if (phdr.p_type != PT_LOAD) {
			continue;
		}

		kern_handle_t vmo = image->e_image;
		vm_prot_t prot = VM_PROT_USER;
		size_t offset = phdr.p_offset;

		phdr.p_flags &PF_R && (prot |= VM_PROT_READ);
		phdr.p_flags &PF_W && (prot |= VM_PROT_WRITE);
		phdr.p_flags &PF_X && (prot |= VM_PROT_EXEC);
		if (phdr.p_flags & PF_W) {
			vmo = image->e_data;
			offset = data_offset;
			size_t tmp = 0;

			status = vm_object_copy(
				image->e_data,
				data_offset + (phdr.p_offset & PAGE_MASK),
				image->e_image,
				phdr.p_offset,
				phdr.p_filesz,
				&tmp);

			if (tmp != phdr.p_filesz) {
				return LAUNCH_ERR_IMAGE_DATA_LOAD_FAILED;
			}
		}

		if (status != KERN_OK) {
			return LAUNCH_ERR_IMAGE_DATA_LOAD_FAILED;
		}

		if (image->e_local_exec != KERN_HANDLE_INVALID) {
			status = vm_region_map_relative(
				image->e_local_exec,
				phdr.p_vaddr,
				vmo,
				offset,
				phdr.p_memsz,
				prot,
				NULL);
		}

		if (status != KERN_OK) {
			return LAUNCH_ERR_MEMORY_MAP_FAILED;
		}

		if (image->e_remote_exec != KERN_HANDLE_INVALID) {
			status = vm_region_map_relative(
				image->e_remote_exec,
				phdr.p_vaddr,
				vmo,
				offset,
				phdr.p_memsz,
				prot,
				NULL);
		}

		if (status != KERN_OK) {
			return LAUNCH_ERR_MEMORY_MAP_FAILED;
		}

		if (phdr.p_flags & PF_W) {
			data_offset += phdr.p_memsz;
			if (data_offset & (PAGE_SIZE - 1)) {
				data_offset &= (PAGE_SIZE - 1);
				data_offset += PAGE_SIZE;
			}
		}
	}

	return LAUNCH_OK;
}

static elf_sym_t *get_dynsym(struct elf_image *image, size_t index)
{
	elf_sym_t *sym = (elf_sym_t *)(image->e_local_base + image->e_dynsym
	                               + (index * image->e_dynsym_entsize));

	if (!sym->st_value) {
		return NULL;
	}

	return sym;
}

static enum launch_status do_rela(struct elf_image *image, elf_rela_t *rela)
{
	int type = ELF64_R_TYPE(rela->r_info);
	elf_sym_t *sym = NULL;

	switch (type) {
	case R_X86_64_JUMP_SLOT:
		sym = get_dynsym(image, ELF64_R_SYM(rela->r_info));
		if (!sym) {
			return LAUNCH_ERR_MISSING_SYMBOL;
		}

		*(uint64_t *)(image->e_local_base + rela->r_offset)
			= image->e_remote_base + sym->st_value + rela->r_addend;
		kern_tracef(
			"JUMP_SLOT: offset=%zx, symbol=%zu, addend=%zx",
			rela->r_offset,
			ELF64_R_SYM(rela->r_info),
			rela->r_addend);
		break;
	default:
		kern_trace("Unknown relocation type");
		return LAUNCH_ERR_UNSUPPORTED_EXECUTABLE;
	}

	return LAUNCH_OK;
}

static enum launch_status do_rela_list(
	struct elf_image *image,
	off_t offset,
	size_t size,
	size_t entsize)
{
	kern_tracef(
		"do_rela_list(%p, %d, %d, %d)",
		image,
		offset,
		size,
		entsize);
	size_t entries = size / entsize;
	elf_rela_t *rela = (elf_rela_t *)(image->e_local_base + offset);
	enum launch_status status = LAUNCH_OK;

	for (size_t i = 0; i < entries; i++) {
		status = do_rela(image, rela);

		if (status != LAUNCH_OK) {
			break;
		}

		rela = (elf_rela_t *)((char *)rela + entsize);
	}

	return status;
}

static enum launch_status do_rel(
	struct elf_image *image,
	off_t offset,
	size_t size,
	size_t entsize)

{
	return LAUNCH_ERR_UNSUPPORTED_EXECUTABLE;
}

static enum launch_status relocate(struct elf_image *image)
{
	elf_dyn_t *dyn
		= (elf_dyn_t *)(image->e_local_base + image->e_dynamic.p_vaddr);

	enum {
		RT_REL,
		RT_RELA,
		RT_PLTREL,
		RT_COUNT,
	};

	int pltrel_type = DT_NULL;
	off_t offsets[RT_COUNT] = {0};
	size_t sizes[RT_COUNT] = {0}, entsizes[RT_COUNT] = {0};

	size_t nr_dyn = image->e_dynamic.p_filesz / sizeof *dyn;
	for (size_t i = 0; i < nr_dyn; i++) {
		kern_tracef("DYN:%zx", dyn[i].d_tag);
		switch (dyn[i].d_tag) {
		case DT_SYMTAB:
			image->e_dynsym = dyn[i].d_un.d_ptr;
			break;
		case DT_SYMENT:
			image->e_dynsym_entsize = dyn[i].d_un.d_val;
			break;
		case DT_REL:
			offsets[RT_REL] = dyn[i].d_un.d_ptr;
			break;
		case DT_RELSZ:
			sizes[RT_REL] = dyn[i].d_un.d_val;
			break;
		case DT_RELENT:
			entsizes[RT_REL] = dyn[i].d_un.d_val;
			break;
		case DT_RELA:
			offsets[RT_RELA] = dyn[i].d_un.d_ptr;
			break;
		case DT_RELASZ:
			sizes[RT_RELA] = dyn[i].d_un.d_val;
			break;
		case DT_RELAENT:
			entsizes[RT_RELA] = dyn[i].d_un.d_val;
			break;
		case DT_PLTREL:
			pltrel_type = dyn[i].d_un.d_val;
			switch (pltrel_type) {
			case DT_REL:
				entsizes[RT_PLTREL] = 0;
				break;
			case DT_RELA:
				entsizes[RT_PLTREL] = sizeof(elf_rela_t);
				break;
			default:
				break;
			}
			break;
		case DT_JMPREL:
			offsets[RT_PLTREL] = dyn[i].d_un.d_ptr;
			break;
		case DT_PLTRELSZ:
			sizes[RT_PLTREL] = dyn[i].d_un.d_val;
			break;
		default:
			break;
		}

		if (dyn[i].d_tag == DT_NULL) {
			break;
		}
	}

	enum launch_status status = LAUNCH_OK;
	if (offsets[RT_RELA] && sizes[RT_RELA] && entsizes[RT_RELA]) {
		kern_trace("RELA");
		status = do_rela_list(
			image,
			offsets[RT_RELA],
			sizes[RT_RELA],
			entsizes[RT_RELA]);

		if (status != LAUNCH_OK) {
			return status;
		}
	}

	if (offsets[RT_PLTREL] && entsizes[RT_PLTREL]) {
		kern_trace("PLTREL");
		if (pltrel_type == DT_REL) {
			status = do_rel(
				image,
				offsets[RT_PLTREL],
				sizes[RT_PLTREL],
				entsizes[RT_PLTREL]);
		} else {
			status = do_rela_list(
				image,
				offsets[RT_PLTREL],
				sizes[RT_PLTREL],
				entsizes[RT_PLTREL]);
		}

		if (status != LAUNCH_OK) {
			return status;
		}
	}

	return LAUNCH_OK;
}

void elf_image_init(struct elf_image *out)
{
	memset(out, 0x0, sizeof(*out));

	kern_config_get(
		KERN_CFG_PAGE_SIZE,
		&out->e_page_size,
		sizeof out->e_page_size);
	out->e_image = KERN_HANDLE_INVALID;
	out->e_data = KERN_HANDLE_INVALID;
	out->e_local_space = KERN_HANDLE_INVALID;
	out->e_remote_space = KERN_HANDLE_INVALID;
	out->e_local_exec = KERN_HANDLE_INVALID;
	out->e_remote_exec = KERN_HANDLE_INVALID;
}

enum launch_status elf_image_load(
	struct elf_image *image,
	kern_handle_t exec_object,
	kern_handle_t local_space,
	kern_handle_t remote_space)
{
	image->e_image = exec_object;
	image->e_local_space = local_space;
	image->e_remote_space = remote_space;

	enum launch_status status = read_header(image);
	if (status != LAUNCH_OK) {
		return status;
	}

	status = parse_phdr(image);
	if (status != LAUNCH_OK) {
		return status;
	}

	if (image->e_interp[0] != 0) {
		return LAUNCH_ERR_INTERPRETER_REQUIRED;
	}

	kern_status_t kstatus = vm_object_create(
		".data",
		5,
		image->e_data_size,
		VM_PROT_READ | VM_PROT_WRITE | VM_PROT_USER,
		&image->e_data);
	if (kstatus != KERN_OK) {
		return LAUNCH_ERR_NO_MEMORY;
	}

	status = create_exec_regions(image);
	if (status != KERN_OK) {
		return LAUNCH_ERR_NO_MEMORY;
	}

	status = map_executable(image);
	if (status != LAUNCH_OK) {
		return status;
	}

	status = relocate(image);
	if (status != LAUNCH_OK) {
		return status;
	}

	return LAUNCH_OK;
}

void elf_image_cleanup(struct elf_image *image)
{
	vm_region_unmap_relative(image->e_local_exec, 0, image->e_total_size);
	kern_handle_close(image->e_data);
	vm_region_kill(image->e_local_exec);
	kern_handle_close(image->e_local_exec);
	kern_handle_close(image->e_remote_exec);
}