#include <kernel/mem.h>

extern kernel_ctx_t ctx;

// TODO: Fix this, it wastes memory when switching regions because the previous offset isnt being subtracted
// Should work fine for now?
void *mmap_block_to_physical(int block) {
    uint32_t offset = block * BLOCK_SIZE;

    for (uint32_t i = 0; i < ctx.multi_mmap_size; i++) {
        struct multiboot_mmap_entry entry = ctx.multi_mmap[i];
        if (entry.type != MULTIBOOT_MEMORY_AVAILABLE || entry.len > UINT32_MAX || entry.addr > UINT32_MAX) continue;

        uint32_t len  = (uint32_t)(entry.len);
        uint32_t addr = (uint32_t)(entry.addr);
	if (len > offset) {
	    uint32_t phys = (addr + offset);
	    if (phys < 0x00400000) {
		return (void *)(phys + 0xC0000000);
	    }
	    return (void *)phys;
	}
    }

    return 0;
}

void memset(void *s, int c, uint32_t len) {
    unsigned char *dst = s;
    while (len > 0) {
        *dst = (unsigned char) c;
        dst++;
        len--;
    }
}

void memcpy(void *dest, void *src, int n) {  
    for (int i=0; i<n; i++)  
        ((char *)dest)[i] = ((char *)src)[i];  
}


uint8_t mmap_init(void) {
    // Calculate the number of bytes available to us
    uint64_t available_bytes = 0;
    for (uint32_t i = 0; i < ctx.multi_mmap_size; i++) {
        struct multiboot_mmap_entry entry = ctx.multi_mmap[i];
        if (entry.type != MULTIBOOT_MEMORY_AVAILABLE || entry.len > UINT32_MAX || entry.addr > UINT32_MAX) continue;

        //uint32_t len = (uint32_t)(entry.len_low | entry.len_high);
        available_bytes += entry.len;
        LOG("Entry: Addr: %016llX -- Len: %016llX\n", (long long)entry.addr, (long long)entry.len);
    }

    ctx.mmap_size = available_bytes / BLOCK_SIZE / 8;

    // Loop again to find the first region with enough space to hold the memory map
    int index = -1;
    for (uint32_t i = 0; i < ctx.multi_mmap_size; i++) {
        struct multiboot_mmap_entry entry = ctx.multi_mmap[i];
        if (entry.type != MULTIBOOT_MEMORY_AVAILABLE || entry.len > UINT32_MAX || entry.addr > UINT32_MAX) continue;

        //uint32_t len = (uint32_t)(entry.len_low | entry.len_high);
        if (entry.len > ctx.mmap_size) {
            index = i;
            break;
        }
    }

    if (index == -1) return 0; // Failed

    // Create map
    struct multiboot_mmap_entry entry = ctx.multi_mmap[index];
    //uint32_t addr = (uint32_t)(entry.addr_low | entry.addr_high);
    if ((uint32_t)(entry.addr & UINT32_MAX) < 0x00400000)
	ctx.mmap = (uint32_t *)((uint32_t)(entry.addr & UINT32_MAX) + 0xC0000000);
    else
	ctx.mmap = (uint32_t *)(uint32_t)(entry.addr & 0xffffffff);
    // Zero the map
    memset(ctx.mmap, 0, ctx.mmap_size);
    // Reserve the blocks that hold the memory map + 1
    uint32_t blocks_to_set = ctx.mmap_size / BLOCK_SIZE + 1;
    for (uint32_t i = 0; i < blocks_to_set; i++)
        mmap_set_block(i);
    // Reserve the kernel space
    int bits = ctx.mmap_size * 8;
    for (int i = 0; i < bits; i++) {
        uint32_t phys = (uint32_t)mmap_block_to_physical(i);
        if (phys > KERNEL_START && phys < KERNEL_START + KERNEL_SIZE)
            mmap_set_block(i);
    }

    return 1;
}

void mmap_set_block(int block) {
    int index = block / BITS;
    int bit = block;
    if (block >= BITS)
        bit -= BITS;

    uint32_t mask = 1 << bit;
    if (!(ctx.mmap[index] & mask))
        ctx.mmap[index] |= mask;
}

void mmap_free_block(int block) {
    int index = block / BITS;
    int bit = block;
    if (block >= BITS)
        bit -= BITS;

    uint32_t mask = 1 << bit;
    if (ctx.mmap[index] & mask)
        ctx.mmap[index] ^= mask;
}

void *mmap_find_first_free_block(void) {
    int bits = ctx.mmap_size * 8;
    for (int i = 0; i < bits; i++) {
        int index = i / BITS;
        int bit = i;
        if (i >= BITS)
            bit -= BITS;

        uint32_t mask = 1 << bit;

        uint32_t dword = ctx.mmap[index];
        if (!(dword & mask)) {
            mmap_set_block(i);
            return mmap_block_to_physical(i);
        }
    }
    
    return 0;
}