// loader.c: boot mechanism agnostic, bootboot-compliant loader

#include <stdint.h>
#include <stddef.h>

// minimal set of string routines, copied from KarlOS

void *memset(void *ptr, int value, size_t num)
{
    __asm__ ("rep stosb"
        : "+D"(ptr), "+c"(num)
        : "a"(value)
        : "memory");
    return ptr;
}

void *memcpy(void *restrict dest, const void *restrict src, size_t num)
{
    __asm__ ("rep movsb"
        : "+D"(dest), "+S"(src), "+c"(num)
        :
        : "memory");
    return dest;
}

int memcmp(const void *ptr, const void *ptr2, size_t num)
{
    unsigned char *d = (unsigned char *)ptr;
    unsigned char *b = (unsigned char *)ptr2;
    for (size_t i = 0; i < num; i++) {
        if (d[i] != b[i]) {
            return (d[i] - b[i]);
        }
    }
    return 0;
}

size_t strlen(const char *s)
{
    size_t len = 0;
    while (*s != '\0') {
        s++;
        len++;
    }
    return len;
}

// Helpers from bztsrc's bootboot codebase

/**
 * convert ascii octal number to binary number
 */
int octbin(unsigned char *str,int size)
{
    int s=0;
    unsigned char *c=str;
    while(size-->0){
        s*=8;
        s+=*c-'0';
        c++;
    }
    return s;
}

/**
 * convert ascii hex number to binary number
 */
int hexbin(unsigned char *str, int size)
{
    int v=0;
    while(size-->0){
        v <<= 4;
        if(*str>='0' && *str<='9')
            v += (int)((unsigned char)(*str)-'0');
        else if(*str >= 'A' && *str <= 'F')
            v += (int)((unsigned char)(*str)-'A'+10);
        str++;
    }
    return v;
}

typedef struct {
    uint8_t *ptr;
    uint64_t size;
} file_t;

#include "bootboot.h"
#include "fs.h"

#define KILO(x)  ((intptr_t)(x) * 1024)
#define MEG(x)   (KILO(x) * 1024)
#define PAGES(x) ((intptr_t)(x) * 0x1000)

#define ADDR_FRAMEBUFFER MEG(-64)
#define ADDR_MMIO        MEG(-128)
#define ADDR_BOOTBOOT    MEG(-2)
#define ADDR_ENVIRONMENT (MEG(-2) + PAGES(1))
#define ADDR_LOAD        (MEG(-2) + PAGES(2))

// ELF64

typedef struct {
	unsigned char e_ident[16];
	uint16_t      e_type;
	uint16_t      e_machine;
	uint32_t      e_version;
	uint64_t      e_entry;
	uint64_t      e_phoff;
	uint64_t      e_shoff;
	uint32_t      e_flags;
	uint16_t      e_ehsize;
	uint16_t      e_phentsize;
	uint16_t      e_phnum;
	uint16_t      e_shentsize;
	uint16_t      e_shnum;
	uint16_t      e_shstrndx;
} Elf64_Ehdr;

typedef struct {
	uint32_t p_type;
	uint32_t p_flags;
	uint64_t p_offset;
	uint64_t p_vaddr;
	uint64_t p_paddr;
	uint64_t p_filesz;
	uint64_t p_memsz;
	uint64_t p_align;
} Elf64_Phdr;

static void
panic(const char *message)
{
	for (const char *c = message; *c; c++) {
		uint8_t status = 0;
		do {
			__asm__ ("inb\t%%dx" : "=a"(status) : "d"(0x3F8 + 5));
		} while (!(status & 0x20));
		__asm__ __volatile__ ("outb\t%%dx" : : "a"(*c), "d"(0x3F8));
	}

	for (;;) {
		__asm__ __volatile__ ("hlt");
	}
}

static int
load_elf(file_t file, uintptr_t *entry)
{
	if (file.size < sizeof (Elf64_Ehdr)) {
		return -1;
	}

	Elf64_Ehdr *ehdr = (Elf64_Ehdr *)file.ptr;
	if (memcmp(ehdr->e_ident, "\177ELF\2\1\1", 7) != 0) {
		return -1;
	}

	if (ehdr->e_type != 2) {
		return -1;
	}

	if (ehdr->e_phentsize < sizeof (Elf64_Phdr)) {
		return -1;
	}
	if (file.size < ehdr->e_phoff + ehdr->e_phentsize * ehdr->e_phnum) {
		return -1;
	}

	for (uint16_t p = 0; p < ehdr->e_phnum; p++) {
		Elf64_Phdr *phdr = (Elf64_Phdr *)(file.ptr + ehdr->e_phoff + ehdr->e_phentsize * p);
		if (phdr->p_type == 1) { // PT_LOAD
			if (file.size < phdr->p_offset + phdr->p_filesz) {
				return -1;
			}
			panic("Loading segment");
			memcpy((void *)ADDR_LOAD, file.ptr + phdr->p_offset, phdr->p_filesz);
			if (phdr->p_filesz < phdr->p_memsz) {
				memset((void *)(ADDR_LOAD + phdr->p_filesz), 0, phdr->p_memsz - phdr->p_filesz);
			}
		}
	}

	*entry = ehdr->e_entry;

	return 0;
}

static void *
alloc(BOOTBOOT *bootboot_ptr, size_t size, size_t align)
{
	MMapEnt *mmap_end = (MMapEnt *)((char *)bootboot_ptr + bootboot_ptr->size);
	for (MMapEnt *ent = &bootboot_ptr->mmap; ent < mmap_end; ent++) {
		if (!MMapEnt_IsFree(ent)) continue;
		int64_t mstart = MMapEnt_Ptr(ent);
		int64_t msize  = MMapEnt_Size(ent);
		int64_t mend   = mstart + msize;
		mend -= size;
		mend &= -align;
		if (mend < mstart) continue;
		ent->size = ((mend - mstart) & INT64_C(-16)) | MMAP_FREE;
		return (void *)mend;
	}
	panic("panic: Could not allocate from memory map.");
	return NULL;
}

void
loader_main(BOOTBOOT *bootboot_ptr)
{
	uint64_t entry = 0;
	file_t elf = { (void *)bootboot_ptr->initrd_ptr, MEG(2) };
	if (load_elf(elf, &entry) < 0) {
		panic("panic: Malformed ELF64 executable");
	}

	alloc(bootboot_ptr, PAGES(1), PAGES(1));

	uint64_t stack = 0;

	__asm__ __volatile__ (
		"mov\t%1, %%rbp\n\t"
		"mov\t%%rbp, %%rsp\n\t"
		"call\t*%0"
		: : "r"(entry), "r"(stack) : "memory");

	panic("panic: Ran past kernel invocation");
}