#include "kernel/vm.h"
#include "kernel/kalloc.h"
#include "kernel/lib.h"
#include "kernel/riscv.h"

enum
{
    VM_VALID    = 1,
    VM_USER     = 16,
    VM_GLOBAL   = 32,
    VM_ACCESSED = 64,
    VM_DIRTY    = 128,
};

struct PageTable
{
    uint32_t entries[1024];
};
typedef struct PageTable PageTable;
static_assert(sizeof(PageTable) == PAGE_SIZE);

static PageTable* kroot = nullptr;

static inline bool is_aligned(uint32_t p)
{
    return (p & (PAGE_SIZE - 1)) == 0;
}

static inline void* offset_page(void* p)
{
    return (void*)((char*)p + PAGE_SIZE);
}

void kvm_map_one(PageTable* root, void* va, void* pa, uint32_t mode)
{
    uint32_t vpn[] = {
        ((uint32_t)va >> 12) & 0x3ff,
        ((uint32_t)va >> 22) & 0x3ff,
    };

    uint32_t fl = root->entries[vpn[1]];
    if (!(fl & VM_VALID))
    {
        PageTable* child = (PageTable*)kzalloc();
        fl = (((uint32_t)child >> 2) & 0xffff'fc00) | VM_VALID;
        root->entries[vpn[1]] = fl;
    }

    PageTable* child = (PageTable*)((fl & 0xffff'fc00) << 2);
        
    uint32_t ppn[] = {
        ((uint32_t)pa >> 12) & 0x3ff,
        ((uint32_t)pa >> 22) & 0x3ff,
    };

    child->entries[vpn[0]] = (ppn[1] << 20) | (ppn[0] << 10) | VM_VALID | mode;
}

void kvm_map(void* va, void* pa, uint32_t size, uint32_t mode)
{
    if (!is_aligned((uint32_t)va)) panic("kvm_map: virtual address not aligned");
    if (!is_aligned((uint32_t)pa)) panic("kvm_map: physical address not aligned");
    if (!is_aligned(size)) panic("kvm_map: size not aligned");

    for (; size > 0; size -= PAGE_SIZE)
    {
        kvm_map_one(kroot, va, pa, mode);
        va = offset_page(va);
        pa = offset_page(pa);
    }
}

extern char _text_start;
extern char _text_end;

extern char _rodata_start;
extern char _rodata_end;

extern char _data_start;
extern char _data_end;

extern char _bss_start;
extern char _bss_end;

extern char _stack_start;
extern char _stack_end;

extern char _heap_start;
extern char _heap_end;

void kvm_init()
{
    kroot = (PageTable*)kzalloc();

    kvm_map(&_text_start, &_text_start, &_text_end - &_text_start, VM_R | VM_X);
    kvm_map(&_rodata_start, &_rodata_start, &_rodata_end - &_rodata_start, VM_R);
    kvm_map(&_data_start, &_data_start, &_data_end - &_data_start, VM_RW);
    kvm_map(&_bss_start, &_bss_start, &_bss_end - &_bss_start, VM_RW);
    kvm_map(&_heap_start, &_heap_start, &_heap_end - &_heap_start, VM_RW);
    kvm_map(&_stack_start, &_stack_start, &_stack_end - &_stack_start, VM_RW);
    
    kvm_map((void*)0x1000'0000, (void*)0x1000'0000, PAGE_SIZE, VM_RW); // UART
    kvm_map((void*)0x0c00'0000, (void*)0x0c00'0000, 0x8000, VM_RW); // APLIC M
    kvm_map((void*)0x0d00'0000, (void*)0x0d00'0000, 0x8000, VM_RW); // APLIC S
    kvm_map((void*)0x2400'0000, (void*)0x2400'0000, 0x1000, VM_RW); // IMSIC M
    kvm_map((void*)0x2800'0000, (void*)0x2800'0000, 0x1000, VM_RW); // IMSIC S

    w_satp(SATP_MODE_SV32 | ((uint32_t)kroot >> 12));
    sfence_vma();
}