#include "mem.h"
#include "klibc.h"
#include "list.h"
#include "types.h"

static struct mem_desc *page_desc = (struct mem_desc *)&__ld_kernel_end;
static struct mem_desc *free_page;
static struct mem_desc *used_page;
static unsigned long bottom_mem;
static unsigned long top_mem;

static unsigned long allocatedPage = 0;

int memSetup(paddr_t upperMem, paddr_t *lastUsedOut)
{
    // Align upper mem (in kB) on page size even if it does loose a page
    upperMem = ALIGN_DOWN(upperMem, PAGE_SIZE / 1024);

    printf("Available Mem from 0x%x to 0x%x: %dMB \n", &__ld_kernel_end, upperMem * 1024,
           (upperMem * 1024 - (uint32_t)&__ld_kernel_end) / (1024 * 1024));
    // Memory description is stored after the kernel. We need some place to store it
    unsigned long memdesc_end =
        (unsigned long)page_desc + ((upperMem) / (PAGE_SIZE / 1024)) * sizeof(struct mem_desc);
    uint lastUsed = (memdesc_end >> PAGE_SHIFT) + 1;
    list_init(free_page);
    list_init(used_page);
    bottom_mem   = lastUsed;
    *lastUsedOut = memdesc_end;
    top_mem      = upperMem * 1024;
    for (uint i = 0; i < (upperMem / (PAGE_SIZE / 1024)); i++) {
        struct mem_desc *mem = &page_desc[i];
        if (i < lastUsed) {
            mem->ref = 1;
            list_add_tail(used_page, mem);
        } else {
            mem->ref = 0;
            list_add_tail(free_page, mem);
        }
        mem->phy_addr = i * PAGE_SIZE;
    }
    return 0;
}

struct mem_desc *addr2memDesc(paddr_t addr)
{
    if (addr > top_mem || addr < bottom_mem)
        return NULL;

    int idx = addr >> PAGE_SHIFT;
    return page_desc + idx;
}

paddr_t allocPhyPage(void)
{
    if (list_is_empty(free_page)) {
        return (unsigned long)NULL;
    }
    struct mem_desc *mem = list_pop_head(free_page);
    mem->ref             = 1;
    list_add_tail(used_page, mem);
    allocatedPage++;
    return mem->phy_addr;
}

int unrefPhyPage(paddr_t addr)
{
    struct mem_desc *mem = addr2memDesc(addr);
    if (!mem) {
        return -1;
    }
    mem->ref--;
    if (mem->ref == 0) {
        allocatedPage--;
        list_delete(used_page, mem);
        list_add_tail(free_page, mem);
    }

    return mem->ref;
}

int refPhyPage(paddr_t addr)
{
    struct mem_desc *mem = addr2memDesc(addr);
    if (!mem) {
        return -1;
    }
    mem->ref++;
    if (mem->ref == 1) {
        allocatedPage++;
        list_add_tail(used_page, mem);
        list_delete(free_page, mem);
    }

    return 0;
}

unsigned long getNbAllocatedPage(void)
{
    return allocatedPage;
}