#include "uaddrspace.h"
#include "alloc.h"
#include "kernel.h"
#include "klibc.h"
#include "list.h"
#include "mem.h"
#include "mmuContext.h"
#include "paging.h"
#include "process.h"
#include "stdarg.h"
#include "thread.h"
#include "types.h"
#include <assert.h>
#include <stddef.h>
#include <stdint.h>

struct uAddrSpace {
    struct process *process; // The process that is represented by this AS
    struct mmu_context *ctx; // The corresponding MMU configuration

    struct uAddrVirtualReg *listVirtualReg; // List of Virtual Region used by this process

    uaddr_t heapStart; // Start of the Head
    size_t heapSize;   // Heap size -> modified by brk()
};

static int hasOverlap(uaddr_t addr1, size_t size1, uaddr_t addr2, size_t size2)
{
    return max(addr1, addr2) < min(addr1 + size1, addr2 + size2);
}

static struct uAddrVirtualReg *findVirtualRegionFromAddr(struct uAddrSpace *as, uaddr_t uaddr, size_t size)
{
    struct uAddrVirtualReg *reg;
    int idx;

    list_foreach_named(as->listVirtualReg, reg, idx, prevInAddrSpace, nextInAddrSpace)
    {
        if (hasOverlap(reg->addr, reg->size, uaddr, size))
            return reg;
    }
    return NULL;
}

static struct uAddrVirtualReg *findVirtualRegionBeforeAddr(struct uAddrSpace *as,
                                                           uaddr_t uaddr)
{
    struct uAddrVirtualReg *reg, *prev = NULL;
    int idx;

    list_foreach_named(as->listVirtualReg, reg, idx, prevInAddrSpace, nextInAddrSpace)
    {
        if (uaddr > reg->addr)
            break;
        prev = reg;
    }
    return prev;
}

/**
 * Find a address not alreay used in the virtual region of the AS.
 * This address will be equals to uaddr if possible and could take size
 */
static uaddr_t findFreeAddrInVirtualRegion(struct uAddrSpace *as, uaddr_t uaddr, size_t size)
{
    struct uAddrVirtualReg *reg, *regNext, *regOver;

    if (uaddr < PAGING_BASE_USER_ADDRESS)
        uaddr = PAGING_BASE_USER_ADDRESS;
    if(uaddr  > PAGING_TOP_USER_ADDRESS - size)
        uaddr = PAGING_TOP_USER_ADDRESS - size;

    reg = findVirtualRegionFromAddr(as, uaddr, size);

    if (!reg)
        return uaddr;

    int idx;

    regOver = reg;

    //Find last region that overlap
    list_foreach_named(reg->nextInAddrSpace, regNext, idx, prevInAddrSpace, nextInAddrSpace){
        if(!hasOverlap(uaddr, size, regNext->addr, regNext->size))
            break;
        regOver = regNext;
    }

    uaddr = regOver->addr + regOver->size;
    list_foreach_named(regOver->nextInAddrSpace, regNext, idx, prevInAddrSpace,
                       nextInAddrSpace)
    {
        if (!hasOverlap(uaddr, size, regNext->addr, regNext->size) &&
            uaddr <= (PAGING_TOP_USER_ADDRESS - size)) {
            return uaddr;
        }
        if (reg == regNext) // Already checked region
            break;
        uaddr = regNext->addr + regNext->size;
    }
    return (uaddr_t)NULL;
}

struct uAddrSpace *uAddrSpaceCreate(struct process *proc)
{
    struct uAddrSpace *addr = (struct uAddrSpace *)zalloc(sizeof(struct uAddrSpace));

    if (addr == NULL)
        return NULL;

    addr->ctx = mmuContextCreate();

    if (addr->ctx == NULL) {
        free(addr);

        return NULL;
    }

    addr->process = proc;
    list_init(addr->listVirtualReg);

    return addr;
}

int uAddrSpaceDelete(struct uAddrSpace *addr)
{
    struct uAddrVirtualReg *reg;

    list_collapse_named(addr->listVirtualReg, reg, nextInAddrSpace, prevInAddrSpace) {
        // TODO Implement me with real ressources
        assertmsg(reg->res == NULL, "Unsupported mapper ressource");
        // This is memory allocated for the heap just unmap it to free it
        pr_devel("Freeing heap 0x%lx for process %s\n", reg->addr, processGetName(addr->process));
        pageUnmap(reg->addr);
        free(reg);
    }
    return mmuContextUnref(addr->ctx);
}

/**
 * Find the associated uAddrVirtualReg associated to the unmap space and free them;
 */
int uAddrSpaceUnmap(struct uAddrSpace *as, uaddr_t uaddr, size_t size)
{
    if (uaddr < PAGING_BASE_USER_ADDRESS || uaddr > PAGING_TOP_USER_ADDRESS - size)
        return -EINVAL;
    if (!ALIGN(uaddr, PAGE_SIZE) || size <= 0)
        return -EINVAL;

    size                            = ALIGN(size, PAGE_SIZE);
    struct uAddrVirtualReg *reg     = as->listVirtualReg;
    struct uAddrVirtualReg *lastReg = reg ? reg->prevInAddrSpace : NULL;

    while (reg != NULL) {
        if (reg->addr > uaddr + size)
            break;
        struct uAddrVirtualReg *next = reg->nextInAddrSpace;
        // The Virtual Region is completly inside the unmaped space
        if (reg->addr >= uaddr && (reg->addr + reg->size <= uaddr + size)) {
            list_delete_named(as->listVirtualReg, reg, prevInAddrSpace, nextInAddrSpace);
            list_delete_named(reg->res->listVirtualReg, reg, prevInMappedRes, nextInMappedRes);
            if (reg->res->ops && reg->res->ops->close)
                reg->res->ops->close(reg);

            free(reg);
            // Unmaped space is inside and smaller than the VR
            // VR should be splitted
        } else if (reg->addr > uaddr && (reg->addr + reg->size <= uaddr + size)) {
            struct uAddrVirtualReg *new =
                (struct uAddrVirtualReg *)zalloc(sizeof(struct uAddrSpace));
            if (!new)
                return -ENOMEM;

            new->addr   = uaddr + size;
            new->size   = reg->addr + reg->size - (uaddr + size);
            new->right  = reg->right;
            new->offset = uaddr + size - reg->addr;

            reg->size = uaddr - reg->addr;

            list_insert_after_named(as->listVirtualReg, reg, new, prevInAddrSpace,
                                    nextInAddrSpace);
            list_insert_after_named(reg->res->listVirtualReg, reg, new, prevInMappedRes,
                                    nextInMappedRes);
            if (reg->res->ops && reg->res->ops->unmap)
                reg->res->ops->unmap(reg, uaddr, size);
            if (new->res->ops &&new->res->ops->open)
                new->res->ops->open(new);
            break;
            // Only affect the beginning
        } else if (uaddr <= reg->addr && uaddr + size > reg->addr) {
            size_t offset = uaddr + size - reg->addr;
            reg->size -= offset;
            reg->offset += offset;
            reg->addr += offset;
            if (reg->res->ops && reg->res->ops->unmap)
                reg->res->ops->unmap(reg, uaddr, size);
            break;
            // Only affect the end
        } else if (uaddr > reg->addr && uaddr + size > reg->addr + reg->size) {
            size_t unmapSize = reg->addr + reg->size - uaddr;
            reg->size        = uaddr - reg->addr;

            if (reg->res->ops && reg->res->ops->unmap)
                reg->res->ops->unmap(reg, uaddr, unmapSize);
        }
        reg = next;
        if (reg == lastReg)
            break;
    }

    int needMMUSetup = as->ctx != getCurrentThread()->squattedContext;

    if (needMMUSetup)
        threadChangeCurrentContext(as->ctx);
    for (vaddr_t addr = uaddr; addr < uaddr + size; addr += PAGE_SIZE) {
        pageUnmap(addr);
    }
    if (needMMUSetup)
        threadChangeCurrentContext(NULL);

    return 0;
}

struct mmu_context *uAddrSpaceGetMMUContext(struct uAddrSpace *addr)
{
    return addr->ctx;
}

int uAddrSpaceSetHeap(struct uAddrSpace *as, uaddr_t addr, size_t size)
{
    as->heapStart = addr;
    as->heapSize  = size;

    return 0;
}

uaddr_t sysBrk(struct uAddrSpace *as, uaddr_t newHeapTop)
{
    int incSize;
    assert(as->heapStart);
    if (!newHeapTop || newHeapTop < as->heapStart)
        return as->heapStart + as->heapSize;

    newHeapTop = ALIGN(newHeapTop, PAGE_SIZE);

    if (newHeapTop == as->heapStart + as->heapSize)
        return newHeapTop;

    incSize = ALIGN(newHeapTop - (as->heapStart + as->heapSize), PAGE_SIZE);

    if (incSize < 0){
        //TODO how to free allocated page by uAddrSpaceHeapCheckNAlloc
        return as->heapStart + as->heapSize;
    }

    as->heapSize += incSize;

    return as->heapStart + as->heapSize;
}

int uAddrSpaceHeapCheckNAlloc(struct uAddrSpace *as, vaddr_t addr)
{
    struct uAddrVirtualReg *newReg;
    int right = PAGING_MEM_USER | PAGING_MEM_WRITE | PAGING_MEM_READ;

    pr_devel("Checking 0x%lx inside 0x%lx and 0x%lx\n", addr, as->heapStart,
             as->heapStart + as->heapSize);

    if (addr < as->heapStart || addr >= as->heapStart + as->heapSize) {
        return -1;
    }

    pr_devel("Alloc heap for process %s\n", processGetName(as->process));

    vaddr_t addrAlign = ALIGN_DOWN(addr, PAGE_SIZE);
    paddr_t ppage     = allocPhyPage(1);

    if (0 != pageMap(addrAlign, ppage, right))
        goto free_ppage;

    newReg = zalloc(sizeof(struct uAddrVirtualReg));

    if (newReg == NULL)
        goto free_ppage;

    newReg->addr = addrAlign;
    newReg->size = PAGE_SIZE;
    newReg->right = right;

    list_add_tail_named(as->listVirtualReg, newReg, nextInAddrSpace, prevInAddrSpace);

    unrefPhyPage(ppage);

    return 0;
free_ppage:
    unrefPhyPage(ppage);
    return -1;
}

int uAddrSpaceMmap(struct uAddrSpace *as, uaddr_t *uaddr, size_t size, uint32_t rights,
                   uint32_t flags, struct mappedRessource *res, uint32_t offset )
{
    int ret            = 0;
    uaddr_t hint_uaddr = *uaddr;

    if (res == NULL || res->ops == NULL || res->ops->nopage == NULL)
        return -ENOENT;
    if (!IS_ALIGNED(hint_uaddr, PAGE_SIZE) || size <= 0)
        return -EINVAL;
    if (flags & UA_MAP_SHARED) {
        if (((rights & PAGING_MEM_READ) && !(res->allowedRight & PAGING_MEM_READ)) ||
            ((rights & PAGING_MEM_WRITE) && !(res->allowedRight & PAGING_MEM_WRITE)) ||
            ((rights & PAGING_MEM_EXEC) && !(res->allowedRight & PAGING_MEM_EXEC)))
            return -EPERM;
    }

    struct uAddrVirtualReg *reg =
        (struct uAddrVirtualReg *)malloc(sizeof(struct uAddrVirtualReg));
    if (!reg)
        return -ENOMEM;

    if (flags & UA_MAP_FIXED) {
        if (hint_uaddr < PAGING_BASE_USER_ADDRESS ||
            hint_uaddr > PAGING_TOP_USER_ADDRESS - size) {
            ret = -EINVAL;
            goto free_reg;
        }
        ret = uAddrSpaceUnmap(as, hint_uaddr, size);
        if (ret)
            goto free_reg;

    } else {
        hint_uaddr = findFreeAddrInVirtualRegion(as, hint_uaddr, size);
        if (!hint_uaddr) {
            ret = -ENOMEM;
            goto free_reg;
        }
    }

    reg->addr   = hint_uaddr;
    reg->size   = size;
    reg->right  = rights;
    reg->res    = res;
    reg->offset = offset;
    // TODO merge it with prev/next one

    // keep the AddrSpace list sorted
    struct uAddrVirtualReg *prev = findVirtualRegionBeforeAddr(as, hint_uaddr);
    if (prev)
        list_insert_after_named(as->listVirtualReg, prev, reg, prevInAddrSpace,
                                nextInAddrSpace);
    else
        list_add_tail_named(as->listVirtualReg, reg, prevInAddrSpace, nextInAddrSpace);

    list_add_tail_named(reg->res->listVirtualReg, reg, prevInMappedRes, nextInMappedRes);
    if (res->onResMapped) {
        int cbret = res->onResMapped(reg);
        if (cbret) {
            pr_devel("Call back failed on ressource mmaped\n");
            ret = uAddrSpaceUnmap(as, reg->addr, reg->size);
        }
    }

    if (res->ops->open)
        res->ops->open(reg);

    *uaddr = hint_uaddr;

    return ret;
free_reg:
    free(reg);
    return ret;
}

/**
 * Check if @param faultAddr could be managed by the uAddrVirtualReg from @param as
 **/
int uAddrSpaceSolvePageFault(struct uAddrSpace *as, vaddr_t faultAddr, int isWriteAccess)
{
    struct uAddrVirtualReg *reg;
    int rights = PAGING_MEM_READ;

    reg = findVirtualRegionFromAddr(as, faultAddr, 1);

    if (reg == NULL)
        return -EFAULT;

    if (isWriteAccess && !(reg->right & PAGING_MEM_WRITE))
        return -EFAULT;

    if (isWriteAccess)
        rights |= PAGING_MEM_WRITE;

    if (reg->res->ops->nopage(reg, faultAddr, rights))
        return -EFAULT;

    return 0;
}