#include "allocArea.h"
#include "alloc.h"
#include "assert.h"
#include "irq.h"
#include "kernel.h"
#include "list.h"
#include "mem.h"
#include "stdarg.h"

struct memArea {
    vaddr_t startAddr;
    uint nbPages;

    struct memArea *next;
    struct memArea *prev;
};

static struct memArea *freeArea;
static struct memArea *usedArea;

static int areaMergeFreeArea(struct memArea *prev, struct memArea **next);

void areaInit(vaddr_t firstMemUsed, vaddr_t lastUsed, vaddr_t stackBottom, vaddr_t stackTop)
{
    list_init(freeArea);
    list_init(usedArea);

    vaddr_t areaAddr, descAddr, entryAddr;
    allocSetup(sizeof(struct memArea), &areaAddr, &descAddr, &entryAddr);
    areaAdd(descAddr, descAddr + PAGE_SIZE, FALSE);
    if (entryAddr != descAddr)
        areaAdd(entryAddr, entryAddr + PAGE_SIZE, FALSE);
    if (areaAddr != descAddr && areaAddr != entryAddr)
        areaAdd(areaAddr, areaAddr + PAGE_SIZE, FALSE);

    // kernel bootstrap part
    areaAdd(ALIGN_DOWN(firstMemUsed, PAGE_SIZE), stackBottom, FALSE);

    // Initial kernel stack
    areaAdd(stackBottom, stackTop, FALSE);

    // Rest of kernel code
    areaAdd(stackTop, lastUsed, FALSE);

    // Rest of virtual mem is free
    areaAdd(areaAddr + PAGE_SIZE, AREA_MEM_TOP, TRUE);

    // Create allocBank for the rest of the system
    allocPopulate();
}

struct memArea *areaFindFit(unsigned int nbPages)
{
    struct memArea *area;
    int count;

    list_foreach(freeArea, area, count)
    {
        if (area->nbPages >= nbPages)
            return area;
    }
    return NULL;
}

static void insertSorted(struct memArea **list, struct memArea *item)
{
    struct memArea *prevArea, *cur;
    int count;
    prevArea = NULL;

    list_foreach(*list, cur, count)
    {
        if (cur->startAddr > item->startAddr)
            break;
        prevArea = cur;
    }
    if (prevArea)
        list_insert_after(*list, prevArea, item);
    else
        list_add_head(*list, item);
}

vaddr_t areaBook(unsigned int nbPages, uint32_t flags)
{

    struct memArea *area;
    vaddr_t allocated;
    uint32_t irqFlags;

    disable_IRQs(irqFlags);

    area = areaFindFit(nbPages);
    if (!area) {
        printf("NULL<\n");
        restore_IRQs(irqFlags);

        return (vaddr_t)NULL;
    }

    if (area->nbPages == nbPages) {
        list_delete(freeArea, area);
        allocated = area->startAddr;
    } else {

        allocated = area->startAddr;

        area->nbPages -= nbPages;
        area->startAddr += nbPages * PAGE_SIZE;
    }

    if (flags & AREA_PHY_MAP) {
        for (uint i = 0; i < nbPages; i++) {
            paddr_t page = allocPhyPage(1);
            if (page) {
                if (pageMap(allocated + i * PAGE_SIZE, page, PAGING_MEM_WRITE)) {
                    unrefPhyPage(page);
                    page = (paddr_t)NULL;
                } else {
                    unrefPhyPage(page);
                }
            }
            if (page == (paddr_t)NULL) {
                areaFree(allocated);
                restore_IRQs(irqFlags);

                return (vaddr_t)NULL;
            }
        }
    }

    restore_IRQs(irqFlags);
    return allocated;
}

vaddr_t areaAlloc(unsigned int nbPages, uint32_t flags)
{
    struct memArea *area, *allocated;
    uint32_t irqFlags;

    disable_IRQs(irqFlags);

    area = areaFindFit(nbPages);
    if (!area) {
        restore_IRQs(irqFlags);

        return (vaddr_t)NULL;
    }

    if (area->nbPages == nbPages) {
        list_delete(freeArea, area);
        insertSorted(&usedArea, area);
        allocated = area;
    } else {
        struct memArea *newArea = (struct memArea *)malloc(sizeof(struct memArea));
        if (!newArea) {
            pr_devel("Failed to allocated area of %d pages\n", nbPages);
            restore_IRQs(irqFlags);
            return (vaddr_t)NULL;
        }

        // Avoid insertSorted(freeArea, newArea) call by modifying area

        newArea->nbPages   = nbPages;
        newArea->startAddr = area->startAddr;

        area->nbPages -= nbPages;
        area->startAddr += nbPages * PAGE_SIZE;

        insertSorted(&usedArea, newArea);

        allocated = newArea;
    }

    if (flags & AREA_PHY_MAP) {
        for (uint i = 0; i < nbPages; i++) {
            paddr_t page = allocPhyPage(1);
            if (page) {
                if (pageMap(allocated->startAddr + i * PAGE_SIZE, page, PAGING_MEM_WRITE)) {
                    page = (paddr_t)NULL;
                } else {
                    unrefPhyPage(page);
                }
            }
            if (page == (paddr_t)NULL) {
                areaFree(allocated->startAddr);
                restore_IRQs(irqFlags);

                return (vaddr_t)NULL;
            }
        }
    }

    restore_IRQs(irqFlags);
    return allocated->startAddr;
}

int areaAdd(vaddr_t start, vaddr_t end, int isFree)
{
    struct memArea **area;

    int nbPages = (end - start) / PAGE_SIZE;
    pr_devel("Add %s area 0x%p->0x%p (%d)\n", isFree ? "free" : "used", (void *)start, (void *)end, nbPages);

    assert(nbPages > 0);
    assert(IS_ALIGNED(start, PAGE_SIZE));
    assert(IS_ALIGNED(end, PAGE_SIZE));

    struct memArea *newArea = (struct memArea *)malloc(sizeof(struct memArea));
    if (!newArea)
        return (vaddr_t)NULL;

    newArea->nbPages   = nbPages;
    newArea->startAddr = start;

    if (isFree) {
        area = &freeArea;
    } else {
        area = &usedArea;
    }

    insertSorted(area, newArea);

    if (isFree) {
        if (newArea->prev)
            areaMergeFreeArea(newArea->prev, area);

        if (newArea && newArea->next)
            areaMergeFreeArea(*area, &(newArea->next));
    }

    return 0;
}

static int areaMergeFreeArea(struct memArea *prevArea, struct memArea **nextArea)
{
    if (prevArea->startAddr + prevArea->nbPages * PAGE_SIZE != (*nextArea)->startAddr) {
        return -1;
    }
    struct memArea *toDelete = *nextArea;

    prevArea->nbPages += (*nextArea)->nbPages;
    list_delete(freeArea, toDelete);
    free(*nextArea);

    *nextArea = prevArea;

    return 0;
}

static struct memArea *areaFindMemArea(struct memArea *list, vaddr_t addr)
{
    struct memArea *area;
    int count;
    list_foreach(list, area, count)
    {
        if (area->startAddr <= addr && addr < area->startAddr + area->nbPages * PAGE_SIZE)
            return area;
    }

    return NULL;
}

int areaFree(vaddr_t addr)
{
    struct memArea *area;

    area = areaFindMemArea(usedArea, addr);
    if (!area) {
        pr_info("Cannot find memArea associated to %p\n", (void *)addr);
        return -1;
    }

    for (uint i = 0; i < area->nbPages; i++) {
        pageUnmap(area->startAddr + i * PAGE_SIZE);
    }

    list_delete(usedArea, area);
    insertSorted(&freeArea, area);

    if (area->prev)
        areaMergeFreeArea(area->prev, &area);

    if (area && area->next)
        areaMergeFreeArea(area, &(area->next));

    return 0;
}