#include "assert.h"
#include "errno.h"
#include "kernel.h"
#include "klibc.h"
#include "list.h"
#include "mem.h"
#include "types.h"

static struct phyMemDesc *pageDesc = (struct phyMemDesc *)&__ld_kernel_end;
static struct phyMemDesc *phyFreePage;
static struct phyMemDesc *phyUsedPage;
paddr_t pageDescEnd;
paddr_t upperMem;

static unsigned long allocatedPage = 0;

int memSetup(paddr_t upperMemKB, paddr_t *firstUsed, paddr_t *lastMemUsedOut)
{
    list_init(phyFreePage);
    list_init(phyUsedPage);

    // Align upper mem (in kB) on page size even if it does loose a page
    upperMemKB           = ALIGN_DOWN(upperMemKB, PAGE_SIZE / 1024);
    unsigned long nbPage = ((upperMemKB) / (PAGE_SIZE / 1024));

    printf("Available Mem from 0x%x to 0x%x: %dMB in %d Pages of %dB\n", &__ld_kernel_end,
           upperMemKB * 1024, (upperMemKB * 1024 - (uint32_t)&__ld_kernel_end) / (1024 * 1024),
           nbPage, PAGE_SIZE);

    // Memory description is stored after the kernel. We need some place to store it
    pageDescEnd     = (unsigned long)pageDesc + nbPage * sizeof(struct phyMemDesc);
    *lastMemUsedOut = ALIGN(pageDescEnd, PAGE_SIZE);
    upperMem        = upperMemKB * 1024;
    *firstUsed      = ALIGN_DOWN((paddr_t)&__ld_kernel_begin, PAGE_SIZE);

    // Remove addr from 0 to PAGE_SIZE so we can return 0 for no page available
    memAddBank(0, PAGE_SIZE, 0);

    memAddBank(*lastMemUsedOut, *lastMemUsedOut, 0);

    return 0;
}

int memAddBank(paddr_t bottomMem, paddr_t topMem, int isFree)
{
    topMem = min(topMem, upperMem);
    for (uint i = (bottomMem >> PAGE_SHIFT); i < (topMem >> PAGE_SHIFT); i++) {
        struct phyMemDesc *mem = &pageDesc[i];
        if (isFree) {
            mem->ref = 0;
            list_add_tail(phyFreePage, mem);
        } else {
            mem->ref = 1;
            list_add_tail(phyUsedPage, mem);
        }
        mem->phy_addr = i * PAGE_SIZE;
    }

    return 0;
}

void memGetStat(uint *free, uint *used)
{
    uint idx;
    struct phyMemDesc *mem;
    list_foreach(phyFreePage, mem, idx)
    {
        continue;
    }
    *free = idx;
    list_foreach(phyUsedPage, mem, idx)
    {
        continue;
    }
    *used = idx;
}

struct phyMemDesc *addr2memDesc(paddr_t addr)
{
    int idx = addr >> PAGE_SHIFT;
    return pageDesc + idx;
}

struct phyMemDesc *memFindConsecutiveFreePage(uint nbPage)
{
    struct phyMemDesc *mem, *head;
    uint memIdx, count;

    if (list_is_empty(phyFreePage)) {
        return NULL;
    }
    count  = 1;
    memIdx = 0;
    head   = phyFreePage;
    mem    = phyFreePage;

    while (count < nbPage && (!memIdx || mem != phyFreePage)) {
        memIdx++;
        mem = mem->next;
        if (mem->phy_addr == head->phy_addr + count * PAGE_SIZE) {
            count++;
        } else {
            count = 1;
            head  = mem;
        }
    }

    if (count < nbPage) {
        return NULL;
    }
    return head;
}

paddr_t allocPhyPage(uint nbPage)
{
    struct phyMemDesc *mem, *head, *next;
    uint count;

    head = memFindConsecutiveFreePage(nbPage);

    if (head == NULL) {
        pr_devel("Cannot find %d consecutive page\n", nbPage);
        return (unsigned long)NULL;
    }

    mem  = head;
    next = head->next;
    for (count = 0; count < nbPage; count++) {
        list_delete(phyFreePage, mem);
        mem->ref = 1;
        list_add_tail(phyUsedPage, mem);
        mem  = next;
        next = mem->next;
    }
    allocatedPage += nbPage;

    return head->phy_addr;
}

int unrefPhyPage(paddr_t addr)
{
    struct phyMemDesc *mem = addr2memDesc(addr);
    if (!mem) {
        return -EINVAL;
    }
    if(mem->ref == 0)
        return -EINVAL;
    mem->ref--;
    if (mem->ref == 0) {
        allocatedPage--;
        list_delete(phyUsedPage, mem);
        list_add_tail(phyFreePage, mem); // TODO find the right place to keep free_page sorted;
    }

    return mem->ref;
}

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

    return 0;
}

unsigned long getNbAllocatedPage(void)
{
    return allocatedPage;
}