#include "arch/x86_64/object.h"
#include "object.h"
#include "bootargs.h"
#include "device/initrd.h"
#include "object.h"
#include "panic.h"
#include "string.h"
#include "device/processor.h"
#include "arch/x86_64/page.h"

static tcb_t s_init_tcb;

//Dynamic allocation hell!!!
static uintptr_t
s_new_mapping(objdir_t *untyped_dir)
{
    objdir_entry_t *untyped_entry = NULL;
    for(int i = 1; i < OBJECT_DIRECTORY_MAX_ENTRIES; i++) {
        untyped_entry = &untyped_dir->entries[i];
        if(untyped_entry->type != KO_MEMORY_UNTYPED) continue;
        uintptr_t mapping = 0;
        int err = untyped_retype_page(untyped_entry, (void**)&mapping);
        if(err != 0) continue;
        return mapping;
    }
    return 0;
}

static pmle_t*
s_ensure_mapping_layer(pmle_t *pml, objdir_t *untyped_dir, uintptr_t addr, uint8_t layer)
{
    pmli_t pmli = PML_I_FOR_LAYER(addr, layer);
    if(pml[pmli].p)
        return pmle_get_page(pml[pmli]);

    uintptr_t table_phys = s_new_mapping(untyped_dir);
    pmle_t *table = vmem_phys_tovirt(table_phys);
    memset(table, 0, 0x1000);

    pml[pmli] = (pmle_t) {
        .p = 1,
        .rw = 1,
        .us = 1,
        .osflg = (layer - 1) & 3,
        .paddr = table_phys >> 12
    };
    __asm__ volatile("invlpg (%0)":: "r"(table_phys));
    return vmem_phys_tovirt(table_phys);
}

static uintptr_t
s_map_page(pmle_t *pml4, objdir_t *untyped_dir, uintptr_t addr)
{
    pmle_t *pml3 = s_ensure_mapping_layer(pml4, untyped_dir, addr, 4);
    pmle_t *pml2 = s_ensure_mapping_layer(pml3, untyped_dir, addr, 3);
    pmle_t *pml1 = s_ensure_mapping_layer(pml2, untyped_dir, addr, 2);

    pmli_t pml1i = PML_I_FOR_LAYER(addr, 1);
    pmle_t *pmle = &pml1[pml1i];

    uintptr_t pptr = s_new_mapping(untyped_dir);

    *pmle = (pmle_t) {
        .p = 1,
        .rw = 1,
        .us = 1,
        .osflg = 0,
        .paddr = pptr >> 12
    };
    __asm__ volatile("invlpg (%0)":: "r"(addr));
    return pptr;
}

__attribute__((noreturn))
static void
s_enter_usermode(void *ip, void *sp)
{
    __asm__ volatile("mov %0, %%rsp; \
                      movq %1, %%rcx; \
                      movq $0x202, %%r11; \
                      cli; \
                      swapgs; \
                      sysretq"::
                      "r"(sp), "r"(ip): "memory");
    for(;;);
}

void
init_load(void)
{
    const char *init_filename = bootargs_getarg("init");
    if(init_filename == NULL) {
        kpanic("Missing boot argument \"init\"");
    }
    tar_header_t *init_header = initrd_find_file(init_filename);
    if(init_header == NULL) {
        kpanic("Init file not found in initrd. (expected \"%s\")", init_filename);
    }

    size_t tcb_diri = _initDirectory.self.data++;
    size_t message_diri = _initDirectory.self.data++;

    _initData.tcb_object = tcb_diri;
    _initData.message_object = message_diri;

    _initDirectory.entries[tcb_diri] = (objdir_entry_t) {
        .type = KO_TCB,
        .data = (uintptr_t)&s_init_tcb
    };

    size_t untyped_diri = _initData.untyped_data_dir;
    objdir_t *untyped_dir = (objdir_t*)_initDirectory.entries[untyped_diri].data;

    size_t pml4_diri = _initData.pm_object;
    pmle_t *pml4 = vmem_phys_tovirt(_initDirectory.entries[pml4_diri].data & ~3ULL);

    //Reserve and map pages for init binary
    size_t init_size = initrd_file_size(init_header);

    size_t init_pages = (init_size >> 12);
    uintptr_t init_base = 0x1000;

    for(size_t i = 0; i < init_pages + 1; i++) {
        s_map_page(pml4, untyped_dir, init_base + (i * 0x1000));
    }
    //Copy over init data
    memcpy((void*)init_base, (&((tar_block_t*)init_header)[1])->data, init_size);

    //Create a user stack
    uintptr_t stack_base = 0x00007FFFFFFFF000;
    s_map_page(pml4, untyped_dir, stack_base);
    uintptr_t sp = stack_base + 0xFF0;

    //Create a kernel stack for the init TCB
    uintptr_t ksp_base = (uintptr_t)s_init_tcb.kstack;
    s_init_tcb.ksp = ksp_base + 0xFF0;

    processor_t *proc = (processor_t*)processor_current();
    proc->tss.rsp[0] = ksp_base;

    //Create a message object for init
    uintptr_t message_base = init_base + (init_pages << 12);
    message_base += ((~(message_base & 0xFFF)) & 0xFFF) + 1;

    uintptr_t message_phys = s_map_page(pml4, untyped_dir, message_base);
    _initDirectory.entries[message_diri] = (objdir_entry_t) {
        .type = KO_MESSAGE,
        .data = (uintptr_t)vmem_phys_tovirt(message_phys)
    };
    _initData.message_object_address = message_base;

    //Write init data to user stack.
    sp -= sizeof(init_data_t);
    memcpy((void*)sp, &_initData, sizeof(init_data_t));
    sp -= sizeof(uintptr_t);
    *((uintptr_t*)sp) = sp + sizeof(uintptr_t);

    //Setup usermode and jump
    proc->tcb = &s_init_tcb;
    proc->tss.rsp[0] = s_init_tcb.ksp;
    s_enter_usermode((void*)init_base, (void*)sp);
}