#include "init.h"
#include "bootargs.h"
#include "device/initrd.h"
#include "arch/x86_64/page.h"
#include "arch/x86_64/processor.h"
#include "device/processor.h"
#include "print.h"
#include "elf.h"
#include "string.h"

uintptr_t 
init_alloc_pageframe()
{
    objdir_entry_t *untyped_dir_entry = &_initDirectory.entries[INIT_OBJECT_UNTYPED_DIR];
    objdir_t *untyped_dir = ko_entry_data(untyped_dir_entry);

    /*Seek to the last member of untyped directory.*/
    uint8_t lastmembi = 1;
    for(; untyped_dir->entries[lastmembi].type == KO_MEMORY_UNTYPED; lastmembi++);
    objdir_entry_t *lastmemb = &untyped_dir->entries[lastmembi-1];

    objdir_entry_t pageframe;
    ko_untyped_split(lastmemb, &pageframe, 0x1000);
    return pageframe.data;
}

void 
init_map_pageframe(uintptr_t pptr, uintptr_t vptr, uint64_t pflags)
{
    uint64_t pathval;
    uint16_t *path = (uint16_t*)&pathval;
    pml4_get_path(vptr, 4, path);

    processor_t *proc = processor_current();
    pmle_t *pml4 = pptr_tovirt_ident(proc->pdir);
    
    pmle_t *pmld = pml4;
    for(uint8_t d = 0; d < 3; d++) {
        uint16_t pmli = path[d];
        pmle_t *pmle = &pmld[pmli];
        if(!pmle->p) {
            pmle->value = init_alloc_pageframe() | PAGE_PRESENT | PAGE_RW | PAGE_US;
        }
        pmld = pptr_tovirt_ident(pmle->paddr << 12);
    }
    pmld[path[3]].value = pptr | pflags;
    __asm__ volatile("invlpg (%0)":: "r"(vptr): "memory");
}

void
init_ensure_page(uintptr_t vptr, uint8_t pflags)
{
    uint64_t pathval;
    uint16_t *path = (uint16_t*)&pathval;
    pml4_get_path(vptr, 4, path);

    processor_t *proc = processor_current();
    pmle_t *pml4 = pptr_tovirt_ident(proc->pdir);
    
    pmle_t *pmld = pml4, *pmle = NULL;
    for(uint8_t d = 0; d < 4; d++) {
        uint16_t pmli = path[d];
        pmle = &pmld[pmli];
        if(!pmle->p) {
            physptr_t frame = init_alloc_pageframe();
            pmle->value = frame | (d == 3 ? pflags : (PAGE_PRESENT | PAGE_RW | PAGE_US));
            if(d != 3) memset(pptr_tovirt_ident(frame), 0, 0x1000);
            pmle->p = 1;
        }
        pmld = pptr_tovirt_ident(pmle->paddr << 12);
    }
    __asm__ volatile("invlpg (%0)":: "r"(vptr): "memory");
}

void
init_reperm_page(uintptr_t vptr, uint64_t pflags)
{
    uint64_t pathval;
    uint16_t *path = (uint16_t*)&pathval;
    pml4_get_path(vptr, 4, path);

    processor_t *proc = processor_current();
    pmle_t *pml4 = pptr_tovirt_ident(proc->pdir);
    
    pmle_t *pmld = pml4, *pmle = NULL;
    for(uint8_t d = 0; d < 4; d++) {
        uint16_t pmli = path[d];
        pmle = &pmld[pmli];
        if(!pmle->p) {
            return;
        }
        pmld = pptr_tovirt_ident(pmle->paddr << 12);
    }
    pmle->rw = pflags & PAGE_RW ? 1 : 0;
    pmle->xd = pflags & PAGE_XD ? 1 : 0;
    __asm__ volatile("invlpg (%0)":: "r"(vptr): "memory");
}

void
init_ensure_range(uintptr_t base, size_t width, uint64_t pflags)
{
    if(width & 0xFFF) width = (width | 0xFFF) + 1;
    width >>= 12;
    for(size_t i = 0; i < width; i++) {
        init_ensure_page(base + (i << 12), pflags);
    }
}

void
init_reperm_range(uintptr_t base, size_t width, uint64_t pflags)
{
    if(width & 0xFFF) width = (width | 0xFFF) + 1;
    width >>= 12;
    for(size_t i = 0; i < width; i++) {
        init_reperm_page(base + (i << 12), pflags);
    }
}

void
init_load(void)
{
    char *init_arg = bootargs_getarg("init");
    if(init_arg == NULL) {
        klogf("init not specified in bootargs!\n");
        return;
    }
    tar_header_t *init_file = initrd_find_file(init_arg);
    if(init_file == NULL) {
        klogf("init not found in initrd! (filename %s)\n", init_arg);
        return;
    }
    uint8_t *init_data = (uint8_t*)&((tar_block_t*)init_file)[1];
    Elf64_Ehdr *init_ehdr = (Elf64_Ehdr*)init_data;
    if(!elf64_ehdr_valid(init_ehdr)) {
        klogf("init is not a proper ELF executable!\n");
        return;
    }

    uintptr_t program_end = 0;
    for(size_t phdri = 0; phdri < init_ehdr->e_phnum; phdri++) {
        Elf64_Phdr *phdr = &((Elf64_Phdr*)(init_data + init_ehdr->e_phoff))[phdri];
        if(phdr->p_type == PT_LOAD) {
            uintptr_t section_base = phdr->p_vaddr;
            size_t section_size = phdr->p_memsz;
            size_t section_cpysize = phdr->p_filesz;

            if(section_base + section_size > program_end) {
                program_end = section_base + section_size;
                if(program_end & 0xFFF) program_end = ((program_end >> 12) + 1) << 12;
            }
            klogf("ELF LOAD %p %p %p %x\n", section_base, section_size, section_cpysize, phdr->p_flags);

            init_ensure_range(section_base, section_size, PAGE_PRESENT | PAGE_US | PAGE_RW);

            memcpy((void*)section_base, init_data + phdr->p_offset, section_cpysize);
            memset((void*)(section_base + section_cpysize), 0, (section_size - section_cpysize));

            init_reperm_range(section_base, section_size, (phdr->p_flags & PF_W ? PAGE_RW : 0) |
                                                          (!(phdr->p_flags & PF_X) ? PAGE_XD : 0));
        }
    }

    //Allocate usermode stack.
    uintptr_t usermode_stack_base = 0x00007FFFFFFFE000;
    init_ensure_range(usermode_stack_base, 0x2000, PAGE_PRESENT | PAGE_RW | PAGE_US | PAGE_XD);

    uintptr_t *usermode_sp = (uintptr_t*)(usermode_stack_base + 0x1FF8);

    //Allocate syscall message.
    uintptr_t message_frame = init_alloc_pageframe();
    init_map_pageframe(message_frame, program_end, PAGE_PRESENT | PAGE_RW | PAGE_US | PAGE_XD);

    _initDirectory.entries[INIT_OBJECT_MESSAGE] = (objdir_entry_t) {
        .type = KO_MESSAGE,
        .extra = KODE_EX_MESSAGE_MAPPED,
        .data = program_end
    };
    *(--usermode_sp) = program_end;

    extern tcb_t *_init_tcb;
    _init_tcb->sp = (uintptr_t)usermode_sp;

    usermode((void*)init_ehdr->e_entry, usermode_sp);
}