#include <stdbool.h>
#include <libc/stdlib.h>
#include <lib/elf.h>
#include <posix/unistd.h>

static bool is_valid_ident(Elf_Byte const e_ident[EI_NIDENT])
{
	return (e_ident[EI_MAG1] == ELFMAG1
			&& e_ident[EI_MAG1] == ELFMAG1
			&& e_ident[EI_MAG2] == ELFMAG2 
			&& e_ident[EI_MAG3] == ELFMAG3
			&& e_ident[EI_CLASS] == ELFCLASS64
			&& e_ident[EI_DATA] == ELFDATA2LSB
			&& e_ident[EI_VERSION] == EV_CURRENT);
}

bool elf64_validate(Elf64_Ehdr const *ehdr, unsigned type, unsigned machine)
{
	return (is_valid_ident(ehdr->e_ident)
			&& ehdr->e_type == type
			&& ehdr->e_machine == machine);
}

static ssize_t read_ehdr(int fd, Elf64_Ehdr *ehdr)
{
	ssize_t pos = lseek(fd, 0, SEEK_CUR);

	ssize_t result;

	if ((result = lseek(fd, 0, SEEK_SET)) == 0)
	{
		result = read(fd, ehdr, sizeof(*ehdr));
	}

	lseek(fd, pos, SEEK_SET);

	return result;
}

static ssize_t read_phdr(int fd, Elf64_Ehdr *ehdr, int phdr_index, Elf64_Phdr *phdr)
{
	if (phdr_index > ehdr->e_phnum)
	{
		return -1;
	}

	ssize_t pos = lseek(fd, 0, SEEK_CUR);

	ssize_t offset = ehdr->e_phoff + ehdr->e_phentsize * phdr_index;

	ssize_t result;
	
	if ((result = lseek(fd, offset, SEEK_SET)) == offset)
	{
		result = read(fd, phdr, ehdr->e_phentsize);
	}

	lseek(fd, pos, SEEK_SET);

	return result;
}

static ssize_t read_phdrs(int fd, Elf64_Ehdr *ehdr, Elf64_Phdr phdrs[ehdr->e_phnum])
{
	ssize_t result = 0;

	for (int i = 0; result >= 0 && i < ehdr->e_phnum; i++)
	{
		result = read_phdr(fd, ehdr, i, &phdrs[i]);
	}

	return result;
}

bool elf64_validate_fd(int fd, unsigned type, unsigned machine)
{
	Elf64_Ehdr ehdr;
	read_ehdr(fd, &ehdr);
	return elf64_validate(&ehdr, type, machine);
}

static bool segment_is_loadable(Elf64_Phdr const *phdr)
{
	return phdr != nullptr && phdr->p_type == PT_LOAD;
}

static size_t aligned_size(Elf64_Phdr const *phdr)
{
	size_t size = phdr->p_offset + phdr->p_memsz;
	return (size + phdr->p_align - 1) & ~(phdr->p_align - 1);
}

static size_t segment_memsz(Elf64_Phdr const *phdr)
{
	return segment_is_loadable(phdr) ? aligned_size(phdr) : 0;
}

size_t elf64_size(Elf64_Ehdr const *ehdr, Elf64_Phdr *phdrs)
{
	size_t size = 0;

	for (int i = 0; i < ehdr->e_phnum; i++)
	{
		size_t size_next = segment_memsz(&phdrs[i]);
		size = size_next > size ? size_next : size;
	}

	return size;
}

size_t elf64_size_fd(int fd)
{
	size_t size = 0;

	Elf64_Ehdr ehdr;
	Elf64_Phdr phdr;

	ssize_t result = read_ehdr(fd, &ehdr);

	for (int i = 0; result >= 0 && i < ehdr.e_phnum; i++)
	{
		result = read_phdr(fd, &ehdr, i, &phdr);
		size_t size_next = segment_memsz(&phdr);
		size = size_next > size ? size_next : size;
	}

	return result > 0 ? size : 0;
}

static Elf64_Dyn *get_dynent(Elf64_Dyn *dyntab, unsigned long dt_type)
{
	while (dyntab && dyntab->d_tag != DT_NULL)
	{
		if (dyntab->d_tag == dt_type) 
		{
			return dyntab;
		}

		dyntab++;
	}

	return nullptr;
}

void *elf64_dt_ptr(char *base, Elf64_Dyn *dyntab, unsigned long dt_type)
{
	Elf64_Dyn *dynent = get_dynent(dyntab, dt_type);

	if (dynent == nullptr)
	{
		return nullptr;
	}

	return dynent->d_ptr + base;
}

uintptr_t elf64_dt_val(Elf64_Dyn *dyntab, unsigned long dt_type)
{
	Elf64_Dyn *dynent = get_dynent(dyntab, dt_type);

	if (dynent == nullptr)
	{
		return (uintptr_t)-1;
	}

	return dynent->d_val;
}

static ssize_t load_segment(int fd, Elf64_Phdr *phdr, size_t buffer_size, char buffer[buffer_size])
{
	if (!segment_is_loadable(phdr))
	{
		return 0;
	}
	
	ssize_t pos = lseek(fd, 0, SEEK_CUR);

	ssize_t result;

	if ((result = lseek(fd, phdr->p_offset, SEEK_SET)) >= 0)
	{
		result = read(fd, &buffer[phdr->p_offset], phdr->p_memsz); 
	}

	lseek(fd, pos, SEEK_SET);

	return result;
}

ssize_t elf64_load_segments(int fd, size_t buffer_size, char buffer[buffer_size])
{
	Elf64_Ehdr *ehdr = (Elf64_Ehdr *)&buffer[0];

	ssize_t result;

	if ((result = read_ehdr(fd, ehdr)) < 0)
	{
		return result;
	}

	Elf64_Phdr *phdrs = (Elf64_Phdr *)&buffer[ehdr->e_phoff];

	if ((result = read_phdrs(fd, ehdr, phdrs)) < 0)
	{
		return result;
	}

	for (int i = 0; result > 0 && i < ehdr->e_phnum; i++)
	{
		result = load_segment(fd, &phdrs[i], buffer_size, buffer);
	}

	return result;
}