/*
 * This file deals with loading programs into memory and executing them.  It
 * parses ELF files for details and loads them into memory, sets up the stack
 * and finally jumps into user-mode.
 */

#include <sys/errno.h>
#include <sys/file.h>
#include <sys/stat.h>
#include <nucleus/lib.h>
#include <nucleus/memory.h>
#include <nucleus/task.h>
#include <nucleus/vfs.h>

/* ELF Types */
enum ELFType {
	ELF_NONE,
	ELF_RELOCATABLE,
	ELF_EXECUTABLE,
	ELF_SHARED,
	ELF_CORE,
};

/* ELF Program Header Types */
enum ProgramType {
	PH_NONE,
	PH_LOAD,
	PH_DYNAMIC,
	PH_INTERPRET,
	PH_NOTE,
	PH_TLS = 7,
};

/* ELF Program Header Flags */
enum ProgramFlag {
	PH_EXECUTABLE = 1,
	PH_WRITABLE = 2,
	PH_READABLE = 4,
};

/* ELF Header */
struct ELFHeader {
	char magic[4];
	uint8_t arch;
	uint8_t endian;
	uint8_t headerVersion;
	uint8_t osAbi;
	uint8_t reserved[8];
	uint16_t type;
	uint16_t isa;
	uint32_t version;
	uint32_t entry;
	uint32_t programHeader;
	uint32_t sectionHeader;
	uint32_t flags;
	uint16_t headerSize;
	uint16_t programEntrySize;
	uint16_t numProgramEntries;
	uint16_t sectionEntrySize;
	uint16_t numSectionEntries;
	uint16_t sectionNames;
};

/* ELF Program Header */
struct ProgramHeader {
	uint32_t type;
	uint32_t offset;
	uint32_t address;
	uint32_t reserved;
	uint32_t filesz;
	uint32_t memsz;
	uint32_t flags;
	uint32_t align;
};

/* Execute a program */
int
execve(const char *file, char *argv[], char *envp[])
{
	if (!file || !verify_access(file, strnlen(file, PATH_MAX), PROT_READ))
		return -EFAULT;

	/* Count argv and envp */
	int argc, argi;
	if (argv == NULL) argc = 0;
	else for (argc = 0; argv[argc]; argc++);
	int envc, envi;
	if (envp == NULL) envc = 0;
	else for (envc = 0; envp[envc]; envc++);

	/* Find size of argv and envp strings */
	size_t ssz = sizeof(int)
	           + (sizeof(uintptr_t) * (argc + 1))
	           + (sizeof(uintptr_t) * (envc + 1));
	for (argi = 0; argi < argc; argi++) {
		if (!verify_access(argv[argi], strlen(argv[argi]), PROT_READ))
			return -EFAULT;
		ssz += strlen(argv[argi]) + 1;
	}
	for (envi = 0; envi < envc; envi++) {
		if (!verify_access(envp[envi], strlen(envp[envi]), PROT_READ))
			return -EFAULT;
		ssz += strlen(envp[envi]) + 1;
	}

	/* Open file */
	int fd = open(file, O_RDONLY);
	if (fd < 0)
		return fd;

	/* Only execute regular files */
	File *executable = get_file_by_fd(fd);
	if (!S_ISREG(executable->inode->mode)) {
		close(fd);
		return -EACCES;
	}

	/* Read ELF header */
	struct ELFHeader header;
	file_read(executable, (char *) &header, sizeof(struct ELFHeader));
	if (memcmp(header.magic, "\x7F""ELF", 4) || header.isa != 3) {
		close(fd);
		return -ENOEXEC;
	}
	if (header.type != ELF_EXECUTABLE) {
		close(fd);
		return -ENOEXEC;
	}

	/*
	 * POINT OF NO RETURN
	 */

	exit_syscall_context();

	/* Store everything (pointers adjusted) in temporary buffer */
	uintptr_t esp = 0xE0000000 - ssz;
	char *istack = kmalloc(ssz);
	char *isp = istack + ssz;
	for (envi = envc - 1; envi >= 0 && envi < envc; envi--) {
		isp -= strlen(envp[envi]) + 1;
		memcpy(isp, envp[envi], strlen(envp[envi]) + 1);
		envp[envi] = (char *) (isp - istack) + esp;
	}
	if (envp)
		envp[envc] = NULL;
	for (argi = argc - 1; argi >= 0 && argi < argc; argi--) {
		isp -= strlen(argv[argi]) + 1;
		memcpy(isp, argv[argi], strlen(argv[argi]) + 1);
		argv[argi] = (char *) (isp - istack) + esp;
	}
	if (argv)
		argv[argc] = NULL;
	isp -= sizeof(uintptr_t);
	*((uintptr_t *) isp) = (uintptr_t) NULL;
	isp -= sizeof(uintptr_t) * envc;
	memcpy(isp, envp, sizeof(uintptr_t) * envc);
	isp -= sizeof(uintptr_t);
	*((uintptr_t *) isp) = (uintptr_t) NULL;
	isp -= sizeof(uintptr_t) * argc;
	memcpy(isp, argv, sizeof(uintptr_t) * argc);
	isp -= sizeof(int);
	*((int *) isp) = argc;

	/* Destroy previous executable */
	VMRegion *region;
	put(current->vm);
	current->vm = new(&virtualMemoryType);

	/* Program headers */
	size_t p;
	off_t off;
	uintptr_t pgbrk, heapEnd;
	struct ProgramHeader ph, tlsph;
	memset(&tlsph, 0, sizeof(struct ProgramHeader));
	for (p = 0; p < header.numProgramEntries; p++) {
		off = header.programHeader + (p * header.programEntrySize);
		executable->pos = off;
		file_read(executable, (char *) &ph,
		          sizeof(struct ProgramHeader));
		if (ph.type != PH_LOAD) {
//			if (ph.type == PH_TLS)
			continue;
		}

		/* Map data into region */
		mmap((void *) ph.address, ph.filesz, ph.flags, MAP_PRIVATE, fd,
		     PAGE_ADDR(ph.offset));
		/* Space left before */
		if (ph.address % PAGE_SIZE) {
			memset((void *) PAGE_ADDR(ph.address), 0,
			       ph.address - PAGE_ADDR(ph.address));
		}
		/* Unset memory */
		if (ph.memsz > ph.filesz) {
			pgbrk = PAGE_ADDR(ph.address + ph.filesz);
			if ((ph.address + ph.filesz) % PAGE_SIZE)
				pgbrk += PAGE_SIZE;
			memset((void *) ph.address + ph.filesz, 0,
			       pgbrk - ph.address - ph.filesz);
			if (ph.memsz > pgbrk - ph.address) {
				mmap((void *) pgbrk,
				     ph.memsz - (pgbrk - ph.address),
				     ph.flags, MAP_PRIVATE | MAP_ANONYMOUS,
				     -1, 0);
			}
		}
		if (ph.address + ph.memsz > heapEnd)
			heapEnd = ph.address + ph.memsz;
	}

	/* Store executable */
	if (current->executable)
		put(current->executable);
	current->executable = get(executable);
	close(fd);

	/* Stack area */
	memcpy((void *) esp, istack, ssz);
	kfree(istack);

	/* Switch to user-mode */
	asm volatile(
		"cli;"
		"mov $0x23, %%ax;"
		"mov %%ax, %%ds;"
		"mov %%ax, %%es;"
		"mov %%ax, %%fs;"
		"mov %%ax, %%gs;"
		"mov %%esi, %%eax;"
		"pushl $0x23;"
		"pushl %%eax;"
		"pushf;"
		"pop %%eax;"
		"or $0x200, %%eax;"
		"push %%eax;"
		"pushl $0x1B;"
		"pushl %%ebx;"
		"iret;"
		: : "b" (header.entry), "S" (esp)
	);
	__builtin_unreachable();
}