#include "kernel.h"

// ================= Placement malloc:
// Assume here that nothing gets ever kfree'd. This makes it simple!

uint64_t placement_address = (uint64_t)&kernel_end;

// ================= Other malloc
// TODO: move this to user space?

#define HEAP_MAGIC 0xc0ffee
#define HEAP_INUSE 1
#define HEAP_FREE 0

/*
  malloc and free divide the chunks of memory present at the heap
  of the kernel into smaller parts.
  The heap is located at: 0x
*/

static heap_t* kernel_heap = 0; 

void expand_heap(heap_t *heap, uint64_t newsize);

/* Allocates 'size' bytes and returns the pointer if succesfull.
   Before the kernel is initialized, placement malloc is used.
   Kernelpanic in case of failure..
*/
static void* kmalloc_int(uint64_t size, uint64_t aligned) 
{
   // Check if there exists a kernel heap. If not, use placement malloc for now.
   if (kernel_heap == 0)
   {
      // Before that the heap is initialized, use placement malloc
      if (aligned == 1)
      {
         // Align placement address
         if ( (placement_address | 0xFFF) != 0 )
         {
            placement_address &= ~(0xFFF);
            placement_address += 0x1000;
         }
      }
      uint64_t tmp = placement_address;
      placement_address += size;
      return (void*)tmp;
   }
   else
   {
      // We have a kernel heap, search the heap for a block of suitable size.

      // Start at the beginning of our heap and search a free block:
      heap_header_t *current = kernel_heap->first_block;
      while (current->magic == HEAP_MAGIC)
      {
         if ((current->state == HEAP_FREE) && (current->size >= size))
         {
            // We found a hole with suitable size. Determine if split is of any use.

            // Insert a heap header if required:
            if (current->size > size + sizeof(heap_t) + 1)
            {
               // Calculate location of the inserted header:
               heap_header_t *newheader = (heap_header_t*) (((uint64_t)current)+size+sizeof(heap_header_t));

               // Set the new header fields:
               newheader->size = current->size - size - sizeof(heap_header_t);
               newheader->state = HEAP_FREE;
               newheader->magic = HEAP_MAGIC;

               // Set the size of this block
               current->size = size; 
            }
            else
            {
               // We allocate this whole block
               // Mark block as used:
               current->state = HEAP_INUSE;
            }
            // Calculate the size of the block:
            char *address = ((char*)current)+sizeof(heap_header_t);
            return address;
            
         }
         // Goto next heap block:
         current = (heap_header_t*)(((uint64_t) current) + current->size + sizeof(heap_header_t));
      }

      // We did not find a block large enough. Expand the heap and try again.
      expand_heap(kernel_heap, 0x10000000);
      return kmalloc_int(size, aligned); 
   }
}

// Convenient wrappers:
void* kmalloc(uint64_t size)
{
   return kmalloc_int(size, 0);
}

void* kmalloc_a(uint64_t size)
{
   return kmalloc_int(size, 1);
}

void kfree(void* ptr) 
{
   printf("Free address %x\n", ptr);
   if (kernel_heap == 0)
   {
      return;
   }
   else
   {
      // TODO
   }
   // TODO: free blocks:
}

heap_t* create_heap(uint64_t location, uint64_t size)
{
   // Make sure that the new heap location is mapped into the address space:
   uint64_t i = 0;
   while (i < size)
   {
      alloc_frame( get_page(location + i, kernel_map));
      i += 0x1000;
   }

   // create new heap structure:
   heap_t* h = (heap_t*)kmalloc(sizeof(heap_t));
   
   // Create one big hole:
   h->first_block = (heap_header_t*)location; // Place header at beginning of heap
   h->first_block->magic = HEAP_MAGIC;
   h->first_block->state = HEAP_FREE;
   h->first_block->size = size - sizeof(heap_header_t);
   return h;
}

// Make the total heap bigger:
void expand_heap(heap_t *heap, uint64_t newsize)
{
   uint64_t oldsize = heap->end_address - heap->start_address;
   if (newsize < oldsize)
   {
      panic("Cannot expand heap to a smaller size\n");
   }

   uint64_t i = oldsize;
   while (i < newsize)
   {
      alloc_frame( get_page(heap->start_address + i, kernel_map));
      i += 0x1000;
   }
}