#include "api/error.h"
#include "sys/lock.h"
#include "hal/obj.h"
#include "lib/printf.h"
#include "lib/string.h"
#include "panic/panic.h"
#include "memory.h"

jove_region_t *region_list_head;

jove_region_t*
region_from_addr(physaddr_t addr)
{
    jove_region_t *bestfit = NULL;
    for(jove_region_t *region = region_list_head; region; region = region->next)
    {
        virtaddr_t region_virt = phys_fromvirt((virtaddr_t)addr);
        if(region_virt < addr) {
            if(bestfit) return bestfit;
            continue;
        }

        virtaddr_t region_limit = region_virt + (region->bitmap.total_bits * 0x1000);
        if(region_limit < addr) {
            if(bestfit) return bestfit;
            continue;
        }

        if(addr < region_limit) {
            if(bestfit) return bestfit;
            continue;
        }

        if(!bestfit)
            bestfit = region;
        else {
            if(bestfit->bitmap.total_bits > region->bitmap.total_bits) {
                bestfit = region;
            }
        }
    }
    klogf("Failed to get region from phys address %p\n", addr);
    return NULL;
}

jove_region_t*
region_data(jove_owned_obj_t *obj)
{
    return obj_deref(obj);
}

physaddr_t 
region_base(jove_owned_obj_t *region)
{
    return obj_getvalue(region);
}

unsigned int 
region_pages(jove_owned_obj_t *region_obj)
{
    jove_region_t *region = region_data(region_obj);
    return region->bitmap.total_bits;
}

unsigned int
region_freepages(jove_owned_obj_t *region_obj)
{
    jove_region_t *region = region_data(region_obj);
    return region->bitmap.free_bits;
}

void 
region_new(jove_owned_obj_t *dest_obj, physaddr_t base, size_t pages)
{
    //Round to next region granularity.
    const size_t blocks = (pages + BITMAP_BLOCK_BITS - 1) / BITMAP_BLOCK_BITS;

    obj_settype(dest_obj, JOVE_OBJ_REGION);
    obj_setvalue(dest_obj, base);

    jove_region_t *dest = phys_tovirt(base);
    memset(dest, 0, sizeof(jove_region_t) + blocks);
    
    //Mark overhead pages as taken.
    const size_t overhead_bytes = sizeof(jove_region_t) + blocks;
    const size_t overhead_pages = (overhead_bytes + 0xFFF) / 0x1000;
    for(size_t i = 0; i < overhead_pages; i++) {
        bitmap_set(&dest->bitmap, i);
    }

    //Setup metadata for top-level pages.
    dest->toplevel = true;
    dest->bitmap.total_bits = pages;
    dest->bitmap.free_bits = pages - overhead_pages;

    //Place new region at the front of the region list.
    dest->next = region_list_head;
    if(dest->next) {
        dest->prev = dest->next->prev;
        dest->next->prev = dest;
    }
    release_mutex(&dest->lock);
}

int 
region_split(jove_owned_obj_t *source_obj, 
        jove_owned_obj_t *dest_obj, size_t pages)
{
    if(pages == 0) return -EBADSIZE;
    if(pages % JOVE_REGION_GRANULARITY != 0) return -EBADSIZE;

    /*Ensure that the source and destination memory regions keep a minimum size.*/
    if(pages > region_pages(source_obj)) return -ETOOBIG;
    if(region_pages(source_obj) - pages == 0) return -ETOOBIG;

    physaddr_t source_base = region_base(source_obj);
    jove_region_t *source = region_data(source_obj);

    //Ensure that the source region can't change as we're splitting.
    if(try_take_mutex(&source->lock)) {
        return -EBUSY;
    }

    //Check that the destination wouldn't clobber some allocated pages.
    const size_t dest_overhead = JOVE_REGION_OVERHEAD(pages);
    for(size_t i = 0; i < dest_overhead; i++) {
        if(bitmap_check(&source->bitmap, source->bitmap.total_bits - pages + i)) {
            release_mutex(&source->lock);
            return -EFULL;
        }
    }

    //Create and zero-fill new region.
    const physaddr_t dest_base = source_base + (pages * 0x1000);
    jove_region_t *dest = phys_tovirt(dest_base);
    memset(dest, 0, sizeof(jove_region_t));
    take_mutex(&dest->lock);

    //Link region to after source region.
    dest->prev = source;
    dest->next = source->next;
    if(dest->next) dest->next->prev = dest;

    //Populate destination bitmap.
    dest->bitmap.total_bits = pages;
    const size_t dest_blocks = pages / BITMAP_BLOCK_BITS;
    const size_t source_blocks = source->bitmap.total_bits / BITMAP_BLOCK_BITS;
    memcpy(dest->bitmap.data, &source->bitmap.data[source_blocks - dest_blocks], dest_blocks);
    memset(&source->bitmap.data[source_blocks - dest_blocks], 0, dest_blocks);

    //Count free pages per block.
    for(size_t i = 0; i < dest_blocks; i++) {
        const bitmap_block_t dest_block = dest->bitmap.data[i];
        for(bitmap_block_t bit = 0; bit < BITMAP_BLOCK_BITS; bit++) {
            source->bitmap.free_bits--;
            if(dest_block & (1U << bit)) continue;
            dest->bitmap.free_bits++;
        }
    }

    //Mark the overhead pages in the dest as taken.
    for(size_t i = 0; i < dest_overhead; i++) {
        bitmap_set(&dest->bitmap, i);
    }

    obj_settype(dest_obj, JOVE_OBJ_REGION);
    obj_setvalue(dest_obj, dest_base);

    release_mutex(&dest->lock);
    release_mutex(&source->lock);

    return 0;
}

int region_release(jove_owned_obj_t *region_obj)
{
    jove_region_t *region = region_data(region_obj);
    if(region->toplevel) return -EPERM;

    //Ensure this region doesn't have any children.
    if(region->next && !region->next->toplevel)
        return -EBUSY;

    if(try_take_mutex(&region->lock))
        return -EBUSY;

    int ret = 0;

    //Ensure that we can change the source mutex.
    jove_region_t *source_region = region->prev;
    if(try_take_mutex(&source_region->lock)) {
        ret = -EBUSY;
        goto release_region_return;
    }

    //Mark overhead pages as available.
    const size_t region_overhead = JOVE_REGION_OVERHEAD(region->bitmap.total_bits);
    for(size_t i = 0; i < region_overhead; i++) {
        bitmap_set(&region->bitmap, i);
    }

    //Copy child bitmap onto source bitmap.
    const size_t source_blocks = source_region->bitmap.total_bits / BITMAP_BLOCK_BITS;
    const size_t child_blocks = region->bitmap.total_bits / BITMAP_BLOCK_BITS;

    memcpy(&source_region->bitmap.data[source_blocks - child_blocks], region->bitmap.data, child_blocks);
    source_region->bitmap.free_bits += region->bitmap.free_bits;

    //Unlink the child region
    if(region->next) region->next->prev = region->prev;
    region->prev->next = region->next;
    
    //Set the region object to JOVE_OBJ_NONE
    obj_settype(region_obj, JOVE_OBJ_NONE);

    //Release and return.
    release_mutex(&source_region->lock);
release_region_return:
    release_mutex(&region->lock);
    return ret;
}

int 
region_alloc(jove_owned_obj_t *region_obj, jove_owned_obj_t *dest_obj)
{
    //Ensure exclusive access to the region.
    jove_region_t *region = region_data(region_obj);
    if(try_take_mutex(&region->lock)) {
        return -EBUSY;
    }

    int ret = 0;
    //Ensure that the region has free pages.
    if(region->bitmap.free_bits == 0) {
        ret = -EFULL;
        goto release_return;
    }

    //Ensure that the bitmap can find the free pages.
    int foundfree = bitmap_findfree(&region->bitmap);
    if(foundfree < 0) {
        ret = foundfree;
        goto release_return;
    }
    bitmap_set(&region->bitmap, foundfree);

    //Create new untyped object for the new page frame.
    obj_settype(dest_obj, JOVE_OBJ_UNTYPED);
    obj_setvalue(dest_obj, region_base(region_obj) + ((uintmax_t)foundfree * 0x1000));

    jove_untyped_t *untyped = obj_deref(dest_obj);
    untyped->limit = 0x1000;

release_return:
    release_mutex(&region->lock);
    return ret;
}

int
region_alloc_contig(jove_owned_obj_t *region_obj, jove_owned_obj_t *dest_obj, size_t pages)
{
    kpanic("Region alloc contig stub!\n");
}