#include "elf.h"

#include "resolve.h"

#include <errno.h>
#include <fcntl.h>
#include <mango/config.h>
#include <mango/handle.h>
#include <mango/log.h>
#include <mango/vm.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <unistd.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

const char *elf_image_status_to_string(enum elf_image_status status)
{
#define ENUM_STR(s)                                                            \
	case s:                                                                \
		return #s
	switch (status) {
		ENUM_STR(ELF_IMAGE_NONE);
		ENUM_STR(ELF_IMAGE_OPEN);
		ENUM_STR(ELF_IMAGE_PARSED);
		ENUM_STR(ELF_IMAGE_LOADED);
		ENUM_STR(ELF_IMAGE_LINKED);
	default:
		return "UNKNOWN";
	}
#undef ENUM_STR
}

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

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

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

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

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

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

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

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

	return true;
}

static int map_image(struct elf_image *image)
{
	elf_phdr_t phdr;
	size_t r = 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);
		lseek(image->e_fd, phdr_offset, SEEK_SET);
		int r = read(image->e_fd, &phdr, sizeof phdr);

		if (r < 0) {
			return -r;
		}

		if (r != sizeof phdr) {
			return ENOEXEC;
		}

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

		int prot = 0;
		size_t offset = phdr.p_offset & ~PAGE_MASK;

		phdr.p_flags &PF_R && (prot |= PROT_READ);
		phdr.p_flags &PF_W && (prot |= PROT_WRITE);
		phdr.p_flags &PF_X && (prot |= PROT_EXEC);

		virt_addr_t vaddr = phdr.p_vaddr;
		virt_addr_t vlimit = phdr.p_vaddr + phdr.p_memsz;
		if (vaddr & PAGE_MASK) {
			vaddr &= ~PAGE_MASK;
		}

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

		if (image->e_hdr.e_type == ET_DYN) {
			vaddr += image->e_base;
			vlimit += image->e_base;
		}

		int fd = image->e_fd;
		int flags = MAP_SHARED | MAP_EXECUTABLE | MAP_FIXED;

		if (phdr.p_flags & PF_W) {
			fd = -1;
			flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED;
			offset = 0;
		}

		void *p
			= mmap((void *)vaddr,
		               vlimit - vaddr,
		               prot,
		               flags,
		               fd,
		               offset);
		if (p == MAP_FAILED) {
			return EIO;
		}

		kern_tracef(
			"mapped PHDR %u [%zx-%zx] at %p",
			i,
			phdr.p_vaddr,
			phdr.p_vaddr + phdr.p_memsz,
			p);

		if (phdr.p_flags & PF_W) {
			lseek(image->e_fd, phdr.p_offset, SEEK_SET);
			void *dst = (void *)image->e_base + phdr.p_vaddr;
			r = read(image->e_fd, dst, phdr.p_filesz);
			if (r < 0) {
				return -r;
			}
		}
	}

	return SUCCESS;
}

static int parse_phdr(struct elf_image *image)
{
	elf_phdr_t phdr;
	size_t r = 0;
	image->e_length = 0;
	image->e_data_length = 0;
	off_t vaddr, vlimit;

	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);
		lseek(image->e_fd, phdr_offset, SEEK_SET);
		int r = read(image->e_fd, &phdr, sizeof phdr);

		if (r < 0) {
			return -r;
		}

		if (r != sizeof phdr) {
			return ENOEXEC;
		}

		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_length = MAX(image->e_length, vlimit);
			break;
#if 0
		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;
		}
#endif
		default:
			break;
		}

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

	return SUCCESS;
}

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

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

	return sym;
}

static void resolve_symbol(unsigned int slot)
{
	kern_tracef("request for symbol %u", slot);
}

static int do_rela(struct elf_image *image, elf_rela_t *rela, bool lazy)
{
	kern_tracef(
		"do_rela(%p, %d, %d, %d)",
		image,
		rela->r_info,
		rela->r_addend,
		rela->r_offset);
	int type = ELF64_R_TYPE(rela->r_info);
	elf_sym_t *sym = NULL;

	switch (type) {
	case R_X86_64_JUMP_SLOT:
		*(uint64_t *)(image->e_base + rela->r_offset) += image->e_base;
		kern_tracef(
			"JUMP_SLOT: offset=%zx, symbol=%zu, addend=%zx",
			rela->r_offset,
			ELF64_R_SYM(rela->r_info),
			rela->r_addend);
		break;
	case R_X86_64_RELATIVE:
		*(uint64_t *)(image->e_base + rela->r_offset)
			= image->e_base + rela->r_addend;
		kern_tracef(
			"RELATIVE: offset=%zx, addend=%zx",
			rela->r_offset,
			rela->r_addend);
		break;
	default:
		kern_log("Unknown relocation type");
		return ENOEXEC;
	}

	return SUCCESS;
}

static int relocate_pltrel(
	struct elf_image *image,
	off_t offset,
	size_t size,
	size_t entsize)
{
	size_t entries = size / entsize;
	elf_rela_t *rela = (elf_rela_t *)(image->e_base + offset);
	int status = SUCCESS;

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

		if (status != SUCCESS) {
			break;
		}

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

	return status;
}

static int relocate_rela(
	struct elf_image *image,
	off_t offset,
	size_t size,
	size_t entsize)
{
	size_t entries = size / entsize;
	elf_rela_t *rela = (elf_rela_t *)(image->e_base + offset);
	int status = SUCCESS;

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

		if (status != SUCCESS) {
			break;
		}

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

	return status;
}

static int relocate_rel(
	struct elf_image *image,
	off_t offset,
	size_t size,
	size_t entsize)
{
	return ENOEXEC;
}

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

{
	kern_tracef("do_rel (unsupported)");
	return ENOEXEC;
}
#endif

static int load_dependency(struct elf_image *image, const char *name)
{
	kern_tracef("required library: %s", name);
	return ENOEXEC;
}

static int parse_dynamic(struct elf_image *image)
{
	if (image->e_dynamic.p_type != PT_DYNAMIC) {
		return SUCCESS;
	}

	image->e_dyn = (elf_dyn_t *)(image->e_base + image->e_dynamic.p_vaddr);

	int status = SUCCESS;

	size_t nr_dyn = image->e_dynamic.p_filesz / sizeof *image->e_dyn;
	for (size_t i = 0; i < nr_dyn; i++) {
		if (image->e_dyn[i].d_tag == DT_NULL) {
			break;
		}

		switch (image->e_dyn[i].d_tag) {
		case DT_NEEDED:
			image->e_nr_links++;
			break;
		case DT_STRTAB:
			image->e_strtab = image->e_dyn[i].d_un.d_ptr;
			break;
		case DT_SYMTAB:
			image->e_dynsym = image->e_dyn[i].d_un.d_ptr;
			break;
		case DT_SYMENT:
			image->e_dynsym_entsize = image->e_dyn[i].d_un.d_val;
			break;
		case DT_PLTGOT:
			image->e_got_plt = image->e_dyn[i].d_un.d_val;
			break;
		case DT_HASH:
			image->e_hash_type = ELF_HASH_STANDARD;
			image->e_hash_table = image->e_dyn[i].d_un.d_ptr;
			break;
		case DT_GNU_HASH:
			image->e_hash_type = ELF_HASH_GNU;
			image->e_hash_table = image->e_dyn[i].d_un.d_ptr;
			break;
		case DT_REL:
			image->e_rel_offset[ELF_RT_REL]
				= image->e_dyn[i].d_un.d_ptr;
			break;
		case DT_RELSZ:
			image->e_rel_size[ELF_RT_REL]
				= image->e_dyn[i].d_un.d_val;
			break;
		case DT_RELENT:
			image->e_rel_entsize[ELF_RT_REL]
				= image->e_dyn[i].d_un.d_val;
			break;
		case DT_RELA:
			image->e_rel_offset[ELF_RT_RELA]
				= image->e_dyn[i].d_un.d_ptr;
			break;
		case DT_RELASZ:
			image->e_rel_size[ELF_RT_RELA]
				= image->e_dyn[i].d_un.d_val;
			break;
		case DT_RELAENT:
			image->e_rel_entsize[ELF_RT_RELA]
				= image->e_dyn[i].d_un.d_val;
			break;
		case DT_PLTREL:
			image->e_pltrel_type = image->e_dyn[i].d_un.d_val;
			switch (image->e_pltrel_type) {
			case DT_REL:
				image->e_rel_entsize[ELF_RT_PLTREL] = 0;
				break;
			case DT_RELA:
				image->e_rel_entsize[ELF_RT_PLTREL]
					= sizeof(elf_rela_t);
				break;
			default:
				break;
			}
			break;
		case DT_JMPREL:
			image->e_rel_offset[ELF_RT_PLTREL]
				= image->e_dyn[i].d_un.d_ptr;
			break;
		case DT_PLTRELSZ:
			image->e_rel_size[ELF_RT_PLTREL]
				= image->e_dyn[i].d_un.d_val;
			break;
		default:
			break;
		}

		image->e_dyn_count++;
	}

	return SUCCESS;
}

static int reserve_exec_region(struct elf_image *image)
{
	void *base
		= mmap(NULL,
	               image->e_length,
	               PROT_NONE,
	               MAP_ANONYMOUS | MAP_PRIVATE,
	               -1,
	               0);
	if (base == MAP_FAILED) {
		return ENOMEM;
	}

	image->e_base = (virt_addr_t)base;
	return KERN_OK;
}

static int create_image_with_name(const char *name, struct elf_image **out)
{
	struct elf_image *elf = malloc(sizeof *elf);
	if (!elf) {
		return ENOMEM;
	}

	memset(elf, 0x0, sizeof *elf);

	snprintf(elf->e_leaf.l_name, sizeof elf->e_leaf.l_name, "%s", name);

	kern_config_get(
		KERN_CFG_PAGE_SIZE,
		&elf->e_page_size,
		sizeof elf->e_page_size);

	*out = elf;
	return SUCCESS;
}

int elf_image_open(const char *path, struct elf_image **out)
{
	struct elf_image *elf = malloc(sizeof *elf);
	if (!elf) {
		return ENOMEM;
	}

	memset(elf, 0x0, sizeof *elf);

	kern_config_get(
		KERN_CFG_PAGE_SIZE,
		&elf->e_page_size,
		sizeof elf->e_page_size);

	int fd = open(path, O_RDONLY);
	if (fd < 0) {
		elf_image_close(elf);
		return -fd;
	}

	elf->e_status = ELF_IMAGE_OPEN;
	elf->e_fd = fd;

	*out = elf;
	return SUCCESS;
}

int elf_image_parse(struct elf_image *img)
{
	if (img->e_status != ELF_IMAGE_OPEN) {
		return EINVAL;
	}

	int e = read(img->e_fd, &img->e_hdr, sizeof img->e_hdr);
	if (e < 0) {
		return -e;
	}

	if (e != sizeof img->e_hdr) {
		return ENOEXEC;
	}

	if (!elf_validate_ehdr(&img->e_hdr)) {
		return ENOEXEC;
	}

	e = parse_phdr(img);
	if (e != SUCCESS) {
		return e;
	}

	img->e_status = ELF_IMAGE_PARSED;
	return SUCCESS;
}

int elf_image_load(struct elf_image *img)
{
	if (img->e_status != ELF_IMAGE_PARSED) {
		return EINVAL;
	}

	int e = reserve_exec_region(img);
	if (e != SUCCESS) {
		return e;
	}

	e = map_image(img);
	if (e != SUCCESS) {
		return e;
	}

	e = parse_dynamic(img);
	if (e != SUCCESS) {
		return e;
	}

	img->e_status = ELF_IMAGE_LOADED;
	return SUCCESS;
}

int elf_image_link(struct elf_image *img)
{
	if (img->e_status != ELF_IMAGE_LOADED) {
		return EINVAL;
	}

	int status = SUCCESS;

	if (img->e_rel_offset[ELF_RT_REL]) {
		status = relocate_rel(
			img,
			img->e_rel_offset[ELF_RT_REL],
			img->e_rel_size[ELF_RT_REL],
			img->e_rel_entsize[ELF_RT_REL]);

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

	if (img->e_rel_offset[ELF_RT_RELA]) {
		status = relocate_rela(
			img,
			img->e_rel_offset[ELF_RT_RELA],
			img->e_rel_size[ELF_RT_RELA],
			img->e_rel_entsize[ELF_RT_RELA]);

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

#if 1
	if (img->e_rel_offset[ELF_RT_PLTREL]) {
		status = relocate_pltrel(
			img,
			img->e_rel_offset[ELF_RT_PLTREL],
			img->e_rel_size[ELF_RT_PLTREL],
			img->e_rel_entsize[ELF_RT_PLTREL]);

		if (status != SUCCESS) {
			return status;
		}
	}
#endif
	*(uintptr_t *)(img->e_base + img->e_got_plt + 16)
		= (uintptr_t)_dl_runtime_resolve;
	*(uintptr_t *)(img->e_base + img->e_got_plt + 8) = (uintptr_t)img;

	img->e_entry = (virt_addr_t)img->e_base + img->e_hdr.e_entry;
	img->e_status = ELF_IMAGE_LINKED;
	return SUCCESS;
}

extern int elf_image_collect_dependencies(
	struct elf_image *img,
	struct image_list *dest)
{
	if (!img->e_nr_links || img->e_links) {
		return SUCCESS;
	}

	int nr_added = 0;
	img->e_links = calloc(img->e_nr_links, sizeof(struct elf_image *));

	for (size_t i = 0; i < img->e_dyn_count; i++) {
		if (img->e_dyn[i].d_tag != DT_NEEDED) {
			continue;
		}

		const char *name = (const char *)img->e_base + img->e_strtab
		                 + img->e_dyn[i].d_un.d_val;

		if (image_list_get(dest, name)) {
			continue;
		}

		struct elf_image *dep = NULL;
		int status = create_image_with_name(name, &dep);
		if (status != SUCCESS) {
			return -status;
		}

		image_list_put(dest, &dep->e_leaf);
		img->e_links[nr_added] = dep;
		nr_added++;
	}

	return nr_added;
}

void elf_image_close(struct elf_image *image)
{
	if (image->e_fd) {
		close(image->e_fd);
	}

	free(image);
}

static uint32_t std_hash(const char *name)
{
	uint32_t h = 0, g;
	for (; *name; name++) {
		h = (h << 4) + *name;
		if ((g = h & 0xf0000000)) {
			h ^= g >> 24;
		}
		h &= ~g;
	}
	return h;
}

static uint32_t gnu_hash(const char *name)
{
	uint32_t h = 5381;

	for (; *name; name++) {
		h = (h << 5) + h + *name;
	}

	return h;
}

static virt_addr_t find_symbol_stdhash(
	struct elf_image *img,
	const char *name,
	uint32_t hash)
{
	const uint32_t *hashtab
		= (void *)((virt_addr_t)img->e_base + img->e_hash_table);

	const char *strtab = (void *)((virt_addr_t)img->e_base + img->e_strtab);

	const elf_sym_t *symtab
		= (void *)((virt_addr_t)img->e_base + img->e_dynsym);

	const uint32_t nbucket = hashtab[0];
	const uint32_t nchain = hashtab[1];
	const uint32_t *bucket = &hashtab[2];
	const uint32_t *chain = &bucket[nbucket];

	for (uint32_t i = bucket[hash % nbucket]; i; i = chain[i]) {
		if (strcmp(name, strtab + symtab[i].st_name) == 0) {
			return img->e_base + symtab[i].st_value;
		}
	}

	return 0;
}

static virt_addr_t find_symbol_gnuhash(
	struct elf_image *img,
	const char *name,
	uint32_t hash)
{
	return 0;
}

static virt_addr_t find_symbol_slow(struct elf_image *img, const char *name)
{
	return 0;
}

static virt_addr_t find_symbol(
	struct elf_image *img,
	const char *name,
	uint32_t std_hash,
	uint32_t gnu_hash)
{
	switch (img->e_hash_type) {
	case ELF_HASH_STANDARD:
		return find_symbol_stdhash(img, name, std_hash);
	case ELF_HASH_GNU:
		return find_symbol_gnuhash(img, name, gnu_hash);
	default:
		return find_symbol_slow(img, name);
	}
}

virt_addr_t elf_image_find_symbol(struct elf_image *img, const char *name)
{
	uint32_t std_hash_val = std_hash(name);
	uint32_t gnu_hash_val = gnu_hash(name);
	return find_symbol(img, name, std_hash_val, gnu_hash_val);
}

virt_addr_t elf_image_find_linked_symbol(
	struct elf_image *img,
	const char *name)
{
	uint32_t std_hash_val = std_hash(name);
	uint32_t gnu_hash_val = gnu_hash(name);

	virt_addr_t sym = 0;
	for (size_t i = 0; i < img->e_nr_links; i++) {
		sym = find_symbol(
			img->e_links[i],
			name,
			std_hash_val,
			gnu_hash_val);
		if (sym) {
			break;
		}
	}

	return sym;
}