#include <paging.h>
#include <printf.h>
#include <libc.h>
#include <io.h>

/** a disclaimer- 
 * I've never made an OS, so I don't know what kind of structures I'll have.
 * Because of this, I'm not YET using a slab allocator.
 * This allocator allocates sizes of 2^n, starting at 32, which can waste lots of space.
 * In other places, you'll see me reallocating when it's not nessesary, 
 * as if my allocator doesn't over-allocate.
 *
 * Later in development, I will create a slab allocator, making it so that this isn't an issue.
**/



#define MIN_CHUNK_ORDER 5 //equal to 32 bytes for heap_chunk struct. Keeping away from 24 for easy division.
#define MAX_CHUNK_I 6
#define CHUNK_LIST_LEN MAX_CHUNK_I + 1
#define LOG_SUB 63 - MIN_CHUNK_ORDER
#define PAGE_SIZE 0x1000 //TODO [minor] maybe move this to some header somewhere

#define CHUNK_SIZE_FROM_INDEX(i) ((1 << ((i) + MIN_CHUNK_ORDER)))
#define CHUNK_INDEX_FROM_SIZE(s) ((63 - MIN_CHUNK_ORDER - __builtin_clzl((s))))


typedef struct __attribute__((packed)) heap_chunk {
  unsigned int size:4; //will use with flags later if needed
  unsigned int lsize:4;
  unsigned int free:1;
  unsigned int mutex:1;
  unsigned long reserved:54;
  struct heap_chunk *fd;
  struct heap_chunk *bk;
} chunk;
 
struct heap_chunk_roots {
  void *fd;
};


//TODO this wastes some memory
chunk arena[CHUNK_LIST_LEN];


void debug_heap() {
  chunk *on_chunk, *last_chunk;
  size_t chunk_expected_size;
  size_t chunk_actual_size;
  for(unsigned int l = 0; l < CHUNK_LIST_LEN; l++) {
    chunk_expected_size = CHUNK_SIZE_FROM_INDEX(l);
    printf("\n%i sized chunks:\n", chunk_expected_size);
    for(on_chunk = arena[l].fd; on_chunk != 0; on_chunk = on_chunk->fd) {
      chunk_actual_size = CHUNK_SIZE_FROM_INDEX(on_chunk->size);
      printf("\t%p", on_chunk);
      if(chunk_expected_size != chunk_actual_size) {
        printf("\nFound chunk of size %i in %i sized freebin!!!!\n", chunk_actual_size, chunk_expected_size);
        hlt();
      }
      if(!on_chunk->free) {
        printf("\nChunk found in freebin isn't free!\n");
        hlt();
      }
      //may take a while
      if((uintptr_t)on_chunk & (0x1000 - 1)) {
        last_chunk = (void *)on_chunk - CHUNK_SIZE_FROM_INDEX(on_chunk->lsize);
        if(last_chunk->size != on_chunk->lsize) {
          printf("\nChunk before this one is a different size then lsize!");
          hlt();
        }
      }
    }
  }
  printf("\n\n");
}



void malloc_init() {
  bzero(&arena, sizeof(arena));
}

void free(void *addr) {
  chunk *on_chunk, *buddy_chunk, *next_chunk, *chunk_iter;
  on_chunk = addr - 24;


  for(;;) {
    if(on_chunk->size == (MAX_CHUNK_I + 1)) {
      pfree(on_chunk, 0x1000);
      printf("freeing page %p\n", on_chunk);
      return;
    }
    unsigned int chunk_size = CHUNK_SIZE_FROM_INDEX(on_chunk->size);
    if((uintptr_t)on_chunk & ((chunk_size * 2) - 1)) {

      buddy_chunk = (void *)on_chunk - CHUNK_SIZE_FROM_INDEX(on_chunk->lsize);
      if((buddy_chunk->size != on_chunk->size) || !(buddy_chunk->free)) {
        next_chunk = (void *)on_chunk + chunk_size;
        break;
      }
      chunk_iter = buddy_chunk;

    }
    else {

      buddy_chunk = (void *)on_chunk + chunk_size;
      if((buddy_chunk->size != on_chunk->size) || !(buddy_chunk->free)) {
        next_chunk = buddy_chunk;
        break;
      };
      chunk_iter = on_chunk;

    }

    if(buddy_chunk->fd) (buddy_chunk->fd)->bk = buddy_chunk->bk;
    (buddy_chunk->bk)->fd = buddy_chunk->fd;
    on_chunk = chunk_iter;
    on_chunk->size++;

  }


  next_chunk->lsize = on_chunk->size; 
  on_chunk->free = 1;

  on_chunk->bk = (chunk *)&arena[on_chunk->size];
  on_chunk->fd = arena[on_chunk->size].fd;
  if(arena[on_chunk->size].fd) (arena[on_chunk->size].fd)->bk = on_chunk;
  arena[on_chunk->size].fd = on_chunk;



}


static void split_chunk(chunk *full_chunk, unsigned int full_size, unsigned int req_size) {
  chunk *on_chunk;

  full_chunk->size = req_size;
  full_chunk->free = 0;

  //chunk isn't newly allocated page
  //our free funciton won't check lsize if addr is divisible by 0x1000
  if(full_chunk->bk) {
    if(full_chunk->fd) (full_chunk->fd)->bk = (chunk *)&arena[full_size];
    arena[full_size].fd = full_chunk->fd;
    on_chunk = (void *)full_chunk + CHUNK_SIZE_FROM_INDEX(full_size);
    if((uintptr_t)on_chunk & (0x1000 - 1)) on_chunk->lsize = full_size - 1;
  }
  
  on_chunk = (void *)full_chunk + CHUNK_SIZE_FROM_INDEX(req_size);

  for(unsigned int sz = req_size; sz < full_size; sz++) {
    on_chunk->lsize = (sz == req_size) ? req_size : sz - 1;
    on_chunk->fd = arena[sz].fd;
    on_chunk->bk = (chunk *)&arena[sz];
    on_chunk->size = sz;
    on_chunk->free = 1;

    arena[sz].fd = on_chunk;
    on_chunk = (void *)on_chunk + CHUNK_SIZE_FROM_INDEX(sz);
  }
}


void *malloc(size_t size) {
  size += sizeof(chunk);
  chunk *on_chunk;
  unsigned int chunk_sz_i; //desired chunk size index
  unsigned int chunk_ck_i; //backup chunk size index in case we need to count up
  //rounding size to nearest 2^n to find which size it fits into
  
  if(size < 32) size = 32;

  chunk_sz_i = CHUNK_INDEX_FROM_SIZE(size);
  if(size & (size - 1)) chunk_sz_i++;

  if(chunk_sz_i > MAX_CHUNK_I) return(void *)0;
  
  //search for direct hits
  if(arena[chunk_sz_i].fd) {
    on_chunk = arena[chunk_sz_i].fd;
    if(on_chunk->fd) (on_chunk->fd)->bk = (chunk *)&arena[chunk_sz_i];
    arena[chunk_sz_i].fd = on_chunk->fd;
    on_chunk->free = 0;
    return (void *)on_chunk + sizeof(chunk);
  }

  //search for chunks above our size
  for(chunk_ck_i = chunk_sz_i + 1; chunk_ck_i < CHUNK_LIST_LEN; chunk_ck_i++) {
    if(arena[chunk_ck_i].fd) {
      on_chunk = arena[chunk_ck_i].fd;
      split_chunk(on_chunk, chunk_ck_i, chunk_sz_i);
      return (void *)on_chunk + sizeof(chunk);
    }
  }

  on_chunk = (void *)palloc(0x1000);
  printf("allocating page %p\n", on_chunk);

  split_chunk(on_chunk, 7, chunk_sz_i);
  return (void *)on_chunk + sizeof(chunk);
   
}

void *realloc(void *old_chunk, size_t size) {
  void *new_chunk = malloc(size);
  memcpy(new_chunk, old_chunk, CHUNK_SIZE_FROM_INDEX(((chunk *)(old_chunk-24))->size)); 
  free(old_chunk);
  return new_chunk;
}