#include <mem.h>
#include <printf.h>

// 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(struct kernel_context *ctx, 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) continue;

        uint32_t len  = (uint32_t)(entry.len_low | entry.len_high);
        uint32_t addr = (uint32_t)(entry.addr_low | entry.addr_high);
        if (len > offset)
           return (void *)(addr + offset);
    }

    return 0;
}

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


uint8_t mmap_init(struct kernel_context *ctx) {
    // Calculate the number of bytes available to us
    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) continue;

        uint32_t len = (uint32_t)(entry.len_low | entry.len_high);
        ctx->available_bytes += len;
    }

    ctx->mmap_size = ctx->available_bytes / BLOCK_SIZE / 8;
    printf("Memory map size: %d bytes\n", ctx->mmap_size);

    // 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) continue;

        uint32_t len = (uint32_t)(entry.len_low | entry.len_high);
        if (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);
    ctx->mmap = (uint32_t *)addr;
    // 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(ctx, 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(ctx, i);
        if (phys > KERNEL_START && phys < KERNEL_START + KERNEL_SIZE)
            mmap_set_block(ctx, i);
    }

    return 1;
}

void mmap_set_block(struct kernel_context *ctx, 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(struct kernel_context *ctx, 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(struct kernel_context *ctx) {
    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)) {
            return mmap_block_to_physical(ctx, i);
        }
    }
    
    return 0;
}