#include <mercury/kernel/paging.hpp>
#include <cstddef>

namespace mercury::kernel::allocator {

  struct free_entry {
    uint64_t start;
    uint64_t len;//0 for unused
  };

  struct free_page {
    free_page *next;
    free_entry entries[255];
  };

  free_page *first_page;

  static_assert(sizeof(free_page) == 4088);

  free_entry *get_entry(uint64_t start, uint64_t len) {
    for (free_page *fp = first_page; fp; fp = fp->next)
      for (int i = 0; i < 255; ++i)
        if (fp->entries[i].start == start && fp->entries[i].len == len)
          return fp->entries + i;
    return 0;
  }

  void add_entry(uint64_t start, uint64_t len) {
    for (free_page *fp = first_page; fp; fp = fp->next)
      for (int i = 0; i < 255; ++i)
        if (fp->entries[i].len == 0) {
          fp->entries[i].start = start;
          fp->entries[i].len = len;
          return;
        }
    free_page *new_page = (free_page *)paging::map_new_kernel_pages(1);
    new_page->next = first_page;
    first_page = new_page;
    for (int i = 2; i < 255; ++i)
      new_page->entries[i].len = 0;
    new_page->entries[0].start = (uint64_t)new_page + 4088;
    new_page->entries[0].len = 8;
    new_page->entries[1].start = start;
    new_page->entries[1].len = len;
  }

  //len is power of 2, start is len-aligned
  void free_block(uint64_t start, uint64_t len) {
    free_entry *buddy = get_entry(start ^ len, len);
    if (buddy) {
      buddy->start = start & ~len;
      buddy->len = len * 2;
    }
    else
      add_entry(start, len);
  }

  void free_region(uint64_t start, uint64_t len) {
    uint64_t block_size = 1;
    while (block_size <= len) {
      if (start & block_size) {
        free_block(start, block_size);
        start += block_size;
        len -= block_size;
      }
      block_size *= 2;
    }
    while (len) {
      block_size /= 2;
      if (block_size <= len) {
        free_block(start, block_size);
        start += block_size;
        len -= block_size;
      }
    }
    //testing
    if (len != 0)
      while (1)
        ;
  }

  uint64_t take_region(uint64_t len) {

    uint64_t min_size = 1;
    while (min_size < len)
      min_size *= 2;

    free_entry *entry = 0;

    for (free_page *fp = first_page; fp; fp = fp->next)
      for (int i = 0; i < 255; ++i)
        if (fp->entries[i].len >= min_size) {
          if (fp->entries[i].len == min_size) {
            entry = fp->entries + i;
            goto loop_done;
          }
          if (entry == 0 || fp->entries[i].len < entry->len)
            entry = fp->entries + i;
        }

  loop_done:
    if (entry != 0) {
      uint64_t start = entry->start;
      uint64_t block_len = entry->len;
      entry->len = 0;
      if (block_len != len)
        free_region(start + len, block_len - len);
      return start;
    }

    uint64_t pages = (len - 1) / 4096 + 1;
    uint64_t start = (uint64_t)paging::map_new_kernel_pages(pages);
    if (pages * 4096 != len)
      free_region(start + len, pages * 4096 - len);
    return start;

  }

}

using namespace mercury::kernel::allocator;

void *_new(size_t len) {
  if (len == 0)
    return 0;
  uint64_t vaddr = take_region(len + sizeof(size_t));
  *(size_t *)vaddr = len;
  return (void *)(vaddr + sizeof(size_t));
}

void _delete(void *ptr) {
  if ((uint64_t)ptr == 0)
    return;
  uint64_t vaddr = (uint64_t)ptr - sizeof(size_t);
  free_region(vaddr, *(size_t *)vaddr + sizeof(size_t));
}

void *operator new(size_t len) {
  return _new(len);
}

void operator delete(void *ptr, size_t) {
  return _delete(ptr);
}

void operator delete(void *ptr) {
  return _delete(ptr);
}

void *operator new[](size_t len) {
  return _new(len);
}

void operator delete[](void *ptr, size_t) {
  return _delete(ptr);
}

void operator delete[](void *ptr) {
  return _delete(ptr);
}