#include "init.h"
#include "arch/x86_64/hal/msr.h"
#include "config.h"
#include "hal/process.h"
#include "jove.h"
#include "hal/obj.h"
#include "lib/string.h"
#include "lib/panic.h"
#include "lib/printf.h"
#include "arch/x86_64/lib/elf.h"
#include "arch/x86_64/hal/paging.h"
#include "boot/cmdline.h"
#include "boot/initrd.h"

PAGEALIGN static page_t s_init_program_pml3[512];
PAGEALIGN static page_t s_init_program_pml2[512];
PAGEALIGN static page_t s_init_program_pml1[512];
PAGEALIGN static uint8_t s_init_program_pages[INIT_PROGRAM_PAGES_MAX][0x1000];

PAGEALIGN static page_t s_init_stack_pml3[512];
PAGEALIGN static page_t s_init_stack_pml2[512];
PAGEALIGN static page_t s_init_stack_pml1[512];
PAGEALIGN static uint8_t s_init_stack_pages[INIT_STACK_PAGES][0x1000];

static int
elf_is_valid(struct elf_header *ehdr)
{
    if(ehdr->e_ident[0] != EI_MAG0 ||
       ehdr->e_ident[1] != EI_MAG1 ||
       ehdr->e_ident[2] != EI_MAG2 ||
       ehdr->e_ident[3] != EI_MAG3) {
        kprintf("ELF: Magic does not match!\n");
        return 0;
    }
    if(ehdr->e_ident[4] != EI_CLASS_64) {
        kprintf("ELF: Class is not 64-bit!\n");
        return 0;
    }
    if(ehdr->e_ident[5] != EI_DATA_LE) {
        kprintf("ELF: Data is not little-endian!\n");
        return 0;
    }
    if(ehdr->e_type != ET_EXEC) {
        kprintf("ELF: Program is not standard executable type!\n");
        return 0;
    }
    if(ehdr->e_machine != EM_x86_64) {
        kprintf("ELF: Executable is not x86_64!\n");
        return 0;
    }
    return 1;
}

void
init_setup(void)
{
    memset(s_init_program_pml3, 0, 0x1000);
    memset(s_init_program_pml2, 0, 0x1000);
    memset(s_init_program_pml1, 0, 0x1000);

    memset(s_init_stack_pml3, 0, 0x1000);
    memset(s_init_stack_pml2, 0, 0x1000);
    memset(s_init_stack_pml1, 0, 0x1000);

    uintptr_t init_program_pml3_pptr = KPTR_TOPHYS(&s_init_program_pml3);
    uintptr_t init_program_pml2_pptr = KPTR_TOPHYS(&s_init_program_pml2);
    uintptr_t init_program_pml1_pptr = KPTR_TOPHYS(&s_init_program_pml1);
    uintptr_t init_program_pml0_pptr = KPTR_TOPHYS(&s_init_program_pages);

    uintptr_t init_stack_pml3_pptr = KPTR_TOPHYS(&s_init_stack_pml3);
    uintptr_t init_stack_pml2_pptr = KPTR_TOPHYS(&s_init_stack_pml2);
    uintptr_t init_stack_pml1_pptr = KPTR_TOPHYS(&s_init_stack_pml1);
    uintptr_t init_stack_pml0_pptr = KPTR_TOPHYS(&s_init_stack_pages);

    //Get the init file data from the initrd.
    const char *init_path = cmdline_getarg("init");
    if(init_path == NULL) {
        kpanic("Kernel commandline is missing 'init' arg.\n");
    }

    jove_owned_obj_t *initrd_obj = objdir_getlocal(_rootdir_obj, ROOTDIR_INITRD_I);
    jove_owned_obj_t *init_obj = initrd_seek(initrd_obj, init_path);
    if(init_obj == NULL) {
        kpanic("init: Cannot find file '%s' in initrd.\n", init_path);
    }
    obj_release(initrd_obj);

    void *init_data = initrd_file_data(init_obj);
    obj_release(init_obj);

    struct elf_header *init_ehdr = init_data;
    if(init_ehdr == NULL) {
        kpanic("init: Cannot get file data for file '%s': %i\n", init_path, errno);
    }

    if(!elf_is_valid(init_ehdr)) {
        kpanic("init: Either passed file '%s' is not an ELF file or is not loadable.\n", init_path);
    }

    jove_owned_obj_t *init_pagedir = objdir_getlocal(_rootdir_obj, ROOTDIR_PAGEDIR_I);
    page_t *pml4 = obj_deref(init_pagedir);
    size_t last_base = 0;

    //Load the init ELF
    for(size_t i = 0; i < init_ehdr->e_phnum; i++) {
        struct elf_phdr *phdr = init_data + init_ehdr->e_phoff + i * init_ehdr->e_phentsize;
        if(phdr->p_type != PT_LOAD) continue;
        kprintf("init: phdr %i flags %b\n", i, phdr->p_flags);

        uintptr_t base = phdr->p_vaddr;
        size_t limit = phdr->p_memsz;

        uintptr_t page_base = base & ~0xFFFULL;
        size_t pml3i = (page_base >> PMLI_SHL(3)) & 0x1FF;
        size_t pml2i = (page_base >> PMLI_SHL(2)) & 0x1FF;
        size_t pml1i = (page_base >> PMLI_SHL(1)) & 0x1FF;
        size_t base_pml0i = (page_base >> PMLI_SHL(0)) & 0x1FF;

        pml4[pml3i].value = init_program_pml3_pptr | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_USER;
        s_init_program_pml3[pml2i].value = init_program_pml2_pptr | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_USER;
        s_init_program_pml2[pml1i].value = init_program_pml1_pptr | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_USER;

        size_t pages = (limit + 4095) / 4096;

        for(size_t pmlo = 0; pmlo < pages; pmlo++) {
            size_t pml0i = base_pml0i + pmlo;
            if(pml0i > INIT_PROGRAM_PAGES_MAX) {
                kpanic("init: Program exceeds maximum program pages in the same hugepage frame!\n");
            }
            if(last_base < base + limit) last_base = base + limit;

            uintptr_t pml0_pptr = init_program_pml0_pptr + (pml0i << 12);
            s_init_program_pml1[pml0i].value = pml0_pptr | PAGE_PRESENT | PAGE_USER | PAGE_WRITEABLE;
        }

        //Copy data from section if SHT_PROGBITS
        kprintf("init: load %p : %p\n", base, limit);
        memcpy((void*)phdr->p_vaddr, init_data + phdr->p_offset, phdr->p_memsz);
    }

    //Map the stack
    pml4[255].value = init_stack_pml3_pptr | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_USER | PAGE_XD;
    s_init_stack_pml3[511].value = init_stack_pml2_pptr | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_USER | PAGE_XD;
    s_init_stack_pml2[511].value = init_stack_pml1_pptr | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_USER | PAGE_XD;
    for(size_t pmlo = 0; pmlo < INIT_STACK_PAGES; pmlo++) {
        size_t pml0i = 511 - pmlo;
        uintptr_t page_pptr = init_stack_pml0_pptr + (pmlo << 12);

        s_init_stack_pml1[pml0i].value = page_pptr | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_USER | PAGE_XD;
    }

    //Map the syscall message.
    last_base = (last_base + 0xFFF) & ~0xFFF;
    size_t last_page = (last_base >> 12) & 0x1FF;
    uintptr_t message_pptr = init_program_pml0_pptr + (last_page << 12);
    s_init_program_pml1[last_page].value = message_pptr | PAGE_PRESENT | PAGE_WRITEABLE | PAGE_USER | PAGE_XD;

    jove_owned_obj_t *message_obj = objdir_getlocal(_rootdir_obj, ROOTDIR_MESSAGE_I);
    obj_settype(message_obj, JOVE_OBJ_MESSAGE);
    obj_setvalue(message_obj, (message_pptr & ~0xFFF) | 1);
    obj_release(message_obj);

    obj_release(init_pagedir);

    //Enable syscalls
    extern void _syscall_entry(void);
    msr_set_lstar(_syscall_entry);

    //Setup KSP for syscalls
    jove_owned_obj_t *pcb_obj = process_current();
    pcb_t *pcb = obj_deref(pcb_obj);

    obj_release(pcb_obj);

    uintptr_t usp = 0x00007FFFFFFFFFF8;
    uintptr_t uip = init_ehdr->e_entry;

    __asm__ volatile("movq %0, %%rsp; \
                      movq %1, %%rcx; \
                      movq %2, %%rdi; \
                      movq $0x202, %%r11; \
                      swapgs; \
                      sysretq"::
                     "r"(usp), "r"(uip), "m"(last_base): "memory");
}