#pragma once
/**
 * Kernel - Memory Map
 *
 * @version 0.1.0
 * @license MIT License
 */
#ifndef _KERNEL_MEMORY_H
#define _KERNEL_MEMORY_H

#include <utility.h>

/**
 * Memory block size
 * NOTE: No clue if 4096 is the right optimal size for memory blocks,
 * ... but the first google search says so.
 */
#define MEMORY_BLOCK_SIZE (4096)

/**
 * Memory alignment functions
 */
// Align an memory item/block to the left
static uint32_t kmemory_align_left(const uint32_t item)
{
  return (item / MEMORY_BLOCK_SIZE) * MEMORY_BLOCK_SIZE;
}
// Align an memory item/block to the right
static uint32_t kmemory_align_right(const uint32_t item)
{
  return (item / MEMORY_BLOCK_SIZE + (item % MEMORY_BLOCK_SIZE != 0)) * MEMORY_BLOCK_SIZE;
}

/**
 * Representation of a region of memory
 */
typedef struct memory_region
{
  /**
   * The memory region base address pointer.
   */
  address_t *base;

  /**
   * The memory region size.
   */
  size_t size;
} kernel_memory_region_t;

/**
 * Representation of a memory map
 */
typedef struct memory_map
{
  /**
   * The memory bitmap base address
   */
  uint8_t *bitmap_address;

  /**
   * The amount of bitmap blocks
   */
  size_t bitmap_count;

  /**
   * The memory base address
   */
  uint8_t *data_address;

  /**
   * The amount of data blocks
   */
  size_t data_count;

  /**
   * The amount of blocks used
   */
  size_t used;
} kernel_memory_map_t;

/**
 * Initialize the memory map
 * @param region. The kernel memory region
 */
void kmemory_init(kernel_memory_map_t *mm, const kernel_memory_region_t &region)
{
  // Get the first, and last memory block
  uint32_t first_block = kmemory_align_right((uint32_t)region.base);
  uint32_t last_block  = kmemory_align_left((uint32_t)region.base + region.size);

  // The total blocks of the memory region
  uint32_t size = last_block - first_block;
  // The block count of the memory region
  uint32_t block_count = size / MEMORY_BLOCK_SIZE;
  // The amount of blocks required for the memory map
  uint32_t blocks_required = (block_count / MEMORY_BLOCK_SIZE + (block_count % MEMORY_BLOCK_SIZE != 0)) / MEMORY_BLOCK_SIZE + ((block_count / MEMORY_BLOCK_SIZE + (block_count % MEMORY_BLOCK_SIZE != 0)) % MEMORY_BLOCK_SIZE != 0);

  // The memory map bitmap
  mm->bitmap_address = reinterpret_cast<uint8_t *>(first_block);
  mm->bitmap_count   = blocks_required;

  // The memory map data
  mm->data_address = reinterpret_cast<uint8_t *>(first_block + blocks_required * MEMORY_BLOCK_SIZE);
  mm->data_count   = block_count - blocks_required;

  // Allocate the memory map into memory
  for(size_t i = 0; i < blocks_required; i++)
  {
    *(mm->bitmap_address + i) = 0;
  }

  // mark 0 so whe can compare it to 0x0 or NULL, as
  kmemory_mark_block_used(mm, 0);
}

/**
 * Mark a memory block as free
 * @param mm. The kernel memory map
 * @param block. The block inside the memory map to set
 */
void kmemory_mark_block_free(kernel_memory_map_t *mm, const uint32_t block)
{
  if (block < mm->data_count)
  {
    reinterpret_cast<uint32_t *>(mm->data_address)[block/32] |= ~(1 << (block % 32));
  }
}

/**
 * Mark a memory block as used
 * @param mm. The kernel memory map
 * @param block. The block inside the memory map to set
 */
void kmemory_mark_block_used(kernel_memory_map_t *mm, const uint32_t block)
{
  if (block < mm->data_count)
  {
    reinterpret_cast<uint32_t *>(mm->data_address)[block/32] &= ~(1 << (block % 32));
  }
}

/**
 * Get the total block count of the memory map
 */
size_t kmemory_block_count(const kernel_memory_map_t *mm)
{
  return mm->data_count;
}

/**
 * Get the amount of used blocks in the memory map
 * @param mm. The kernel memory map
 * @return The amount of used blocks in the memory map
 */
size_t kmemory_block_used_count(const kernel_memory_map_t* mm)
{
  return mm->used;
}

/**
 * Get the amount of free blocks in the memory map
 * @param mm. The kernel memory map
 * @return The amount of free blocks in the memory map
 */
size_t kmemory_block_free_count(const kernel_memory_map_t* mm)
{
  return mm->data_count - mm->used;
}

/**
 * Allocate blocks on the memory map
 * @param mm. The kernel memory map
 * @param count. The amount of blocks to allocate
 * @return The base address of the allocated blocks, or null if allocation failed.
 */
address_t* kmemory_allocate_blocks(kernel_memory_map_t* mm, const uint32_t count)
{
  // don't allocate 0 blocks
  if (count == 0)
  {
    return nullptr;
  }

  uint32_t current_block_count = 0;
  for(size_t i = 0; i < mm->data_count; i++)
  {
    // try to find enough blocks to allocate
    if (!(i < mm->block_count) && !(reinterpret_cast<uint32_t *>(mm->data_address)[i/32] & (1 << (i % 32))))
    {
      current_block_count++;
    }
    else {
      current_block_count = 0;
    }

    if (current_block_count == count)
    {
      size_t first_block = i + 1 - count;

      for (size_t block = 0; block < count; block++)
      {
        kmemory_mark_block_used(mm, first_block + block);
      }
      mm->used += count;
      return mm->data_address + first_block + MEMORY_BLOCK_SIZE;
    }
  }

  return nullptr;
}

/**
 * Free blocks on the memory map
 * @param mm. The kernel memory map
 * @param address. The base address to free
 */
void kmemory_free_blocks(kernel_memory_map_t* mm, const address_t* address)
{
  uint32_t offset = reinterpret_cast<uint32_t *>(reinterpret_cast<uint8_t*>(address) - reinterpret_cast<uint32_t>(mm->data_address));

  if (offset % MEMORY_BLOCK_SIZE == 0)
  {
    uint32_t block = offset / MEMORY_BLOCK_SIZE;

    if (!(block < mm->block_count) && !(reinterpret_cast<uint32_t *>(mm->data_address)[block/32] & (1 << (block % 32))))
    {
      mm->used--;
    }

    kmemory_mark_block_free(mm, block);
  }
}

#endif /* end of kernel.h */