#include "errno.h"
#include "globals.h"
#include "mm/mm.h"
#include "util/debug.h"

#include "mm/mman.h"

/*
 * This function implements the brk(2) system call.
 *
 * This routine manages the calling process's "break" -- the ending address
 * of the process's dynamic region (heap)
 *
 * Some important details on the range of values 'p_brk' can take:
 * 1) 'p_brk' should not be set to a value lower than 'p_start_brk', since this
 *    could overrite data in another memory region. But, 'p_brk' can be equal to
 *    'p_start_brk', which would mean that there is no heap yet/is empty.
 * 2) Growth of the 'p_brk' cannot overlap with/expand into an existing
 *    mapping. Use vmmap_is_range_empty() to help with this.
 * 3) 'p_brk' cannot go beyond the region of the address space allocated for use by
 *    userland (USER_MEM_HIGH)
 *
 * Before setting 'p_brk' to 'addr', you must account for all scenarios by comparing
 * the page numbers of addr, 'p_brk' and 'p_start_brk' as the vmarea that represents the heap
 * has page granularity. Think about the following sub-cases (note that the heap 
 * should always be represented by at most one vmarea):
 * 1) The heap needs to be created. What permissions and attributes does a process
 *    expect the heap to have?
 * 2) The heap already exists, so you need to modify its end appropriately.
 * 3) The heap needs to shrink.
 *
 * Beware of page alignment!:
 * 1) The starting break is not necessarily page aligned. Since the loader sets
 *    'p_start_brk' to be the end of the bss section, 'p_start_brk' should always be
 *    aligned up to start the dynamic region at the first page after bss_end.
 * 2) vmareas only have page granularity, so you will need to take this
 *    into account when deciding how to set the mappings if p_brk or p_start_brk
 *    is not page aligned. The caller of do_brk() would be very disappointed if
 *    you give them less than they asked for!
 *
 * Some additional details:
 * 1) You are guaranteed that the process data/bss region is non-empty.
 *    That is, if the starting brk is not page-aligned, its page has
 *    read/write permissions.
 * 2) If 'addr' is NULL, you should return the current break. We use this to
 *    implement sbrk(0) without writing a separate syscall. Look in
 *    user/libc/syscall.c if you're curious.
 * 3) Return 0 on success, -errno on failure. The 'ret' argument should be used to 
 *    return the updated 'p_brk' on success.
 *
 * Error cases do_brk is responsible for generating:
 *  - ENOMEM: attempting to set p_brk beyond its valid range
 */
long do_brk(void *addr, void **ret)
{
    // If addr is NULL, return the current break
    if (addr == NULL) 
    {
        *ret = curproc->p_brk;
        return 0;
    }

    // Check if the address is within the valid range
    if ((uintptr_t)addr > USER_MEM_HIGH || (uintptr_t)addr < USER_MEM_LOW)
    {
        return -ENOMEM;
    }

    // Check if the address is within the valid range
    if (addr < curproc->p_start_brk)
    {
        return -ENOMEM;
    }

    // Check if the address is the same as the current break
    // if (addr == curproc->p_brk)
    // {
    //     *ret = curproc->p_brk;
    //     return 0;
    // }

    // Check if the address is page aligned
    uintptr_t addr_page_aligned = ADDR_TO_PN(PAGE_ALIGN_UP(addr));
    uintptr_t p_brk_page_aligned = ADDR_TO_PN(PAGE_ALIGN_UP(curproc->p_brk));
    uintptr_t p_start_brk_page_aligned = ADDR_TO_PN(PAGE_ALIGN_UP(curproc->p_start_brk));

    // Lookup the vmarea that represents the heap
    vmarea_t *heap_vmarea = vmmap_lookup(curproc->p_vmmap, p_start_brk_page_aligned);

    // Check if the address is the same as the current break
    // If so, set rets and end here
    if (addr_page_aligned == p_brk_page_aligned)
    {
        curproc->p_brk = addr;
        *ret = addr;
        return 0;
    }

    // Check the three cases, whether the heap needs to be created, modified or shrinked
    if (heap_vmarea == NULL)
    {
        // Create the heap
        long ret = vmmap_is_range_empty(curproc->p_vmmap, p_start_brk_page_aligned, addr_page_aligned - p_brk_page_aligned);
        if (!ret)
        {
            // On fail, return -ENOMEM
            return -ENOMEM;
        }

        // Map the heap
        int flags = MAP_PRIVATE | MAP_ANON;
        int prot = PROT_READ | PROT_WRITE;
        ret = vmmap_map(
            curproc->p_vmmap, NULL, 
            p_start_brk_page_aligned, 
            addr_page_aligned - p_start_brk_page_aligned, 
            prot, flags, 
            0, 
            VMMAP_DIR_LOHI, &heap_vmarea
        );
        if (ret < 0)
        {
            // On fail, return ret
            return ret;
        }
    }
    else if (addr_page_aligned < p_brk_page_aligned)
    {
        // Shrink the heap
        long ret = vmmap_remove(curproc->p_vmmap, addr_page_aligned, p_brk_page_aligned - addr_page_aligned);
        if (ret < 0)
        {
            // On fail, return ret
            return ret;
        }
    }
    else
    {
        // Modify the heap
        long ret = vmmap_is_range_empty(curproc->p_vmmap, p_brk_page_aligned, addr_page_aligned - p_brk_page_aligned);
        if (!ret)
        {
            // On fail, return -ENOMEM
            return -ENOMEM;
        }
        // Update the heap
        heap_vmarea->vma_end = addr_page_aligned;
    }

    
    // Update rets & return 0 on success
    curproc->p_brk = addr;
    *ret = addr;
    return 0;
}