#define KL_LOG KL_KMEM
#include <sys/libkern.h>
#include <sys/mimiker.h>
#include <sys/klog.h>
#include <sys/linker_set.h>
#include <sys/sched.h>
#include <sys/malloc.h>
#include <sys/pool.h>
#include <sys/kmem.h>
#include <sys/vm.h>
#include <machine/vm_param.h>

#define PI_ALIGNMENT sizeof(uint64_t)

#define POOL_DEBUG 0

#if defined(POOL_DEBUG) && POOL_DEBUG > 0
#define debug(...) klog(__VA_ARGS__)
#else
#define debug(...)
#endif

static TAILQ_HEAD(, pool) pool_list = TAILQ_HEAD_INITIALIZER(pool_list);
static MTX_DEFINE(pool_list_lock, 0);
static KMALLOC_DEFINE(M_POOL, "pool allocators");

static void *slab_item_at(slab_t *slab, unsigned i) {
  return slab->ph_items + i * slab->ph_itemsize;
}

static void add_slab(pool_t *pool, slab_t *slab, size_t slabsize) {
  assert(mtx_owned(&pool->pp_mtx));
  assert(is_aligned(slab, PAGESIZE));
  assert(slabsize >= pool->pp_slabsize);
  assert(is_aligned(slabsize, PAGESIZE));

  klog("add slab at %p to '%s' pool", slab, pool->pp_desc);

  slab->ph_pool = pool;
  slab->ph_size = slabsize;
  slab->ph_itemsize = pool->pp_itemsize;
#if KASAN
  slab->ph_itemsize += pool->pp_redzone;
#endif /* !KASAN */

  /*
   * Now we need to calculate maximum possible number of items of given `size`
   * in slab that occupies one page, taking into account space taken by:
   *  - items: ntotal * itemsize,
   *  - slab + bitmap: sizeof(slab_t) + bitstr_size(ntotal)
   * With:
   *  - usable = slabsize - sizeof(slab_t)
   * ... inequation looks as follow:
   * (1) ntotal * itemsize + (ntotal + 7) / 8 <= usable
   * (2) ntotal * 8 * itemsize + ntotal + 7 <= usable * 8
   * (3) ntotal * (8 * itemsize + 1) <= usable * 8 - 7
   * (4) ntotal <= (usable * 8 - 7) / (8 * itemsize + 1)
   */
  size_t usable = slabsize - sizeof(slab_t);
  slab->ph_ntotal = (usable * 8 - 7) / (8 * slab->ph_itemsize + 1);
  slab->ph_nused = 0;

  assert(slab->ph_ntotal > 0);

  size_t header = sizeof(slab_t) + bitstr_size(slab->ph_ntotal);
  slab->ph_items = (void *)slab + align(header, pool->pp_alignment);

  bzero(slab->ph_bitmap, bitstr_size(slab->ph_ntotal));

  LIST_INSERT_HEAD(&pool->pp_empty_slabs, slab, ph_link);

  pool->pp_ntotal += slab->ph_ntotal;
  pool->pp_npages += slabsize;

  for (size_t i = 0; i < slabsize; i += PAGESIZE) {
    vm_page_t *pg = kva_find_page((vaddr_t)slab + i);
    assert(pg != NULL);
    pg->slab = slab;
  }
}

void *pool_alloc(pool_t *pool, kmem_flags_t flags) {
  void *ptr;

  debug("pool_alloc: pool=%p", pool);

  WITH_MTX_LOCK (&pool->pp_mtx) {
    slab_t *slab;

    if (!(slab = LIST_FIRST(&pool->pp_part_slabs))) {
      if (!(slab = LIST_FIRST(&pool->pp_empty_slabs))) {
        size_t slabsize = pool->pp_slabsize;
        slab = kmem_alloc(slabsize, flags);
        assert(slab != NULL);
        add_slab(pool, slab, slabsize);
      }
      /* We're going to allocate from empty slab
       * -> move it to the list of non-empty slabs. */
      assert(slab->ph_nused == 0);
      LIST_REMOVE(slab, ph_link);
      LIST_INSERT_HEAD(&pool->pp_part_slabs, slab, ph_link);
    }

    assert(slab->ph_nused < slab->ph_ntotal);
    int i = 0;
    bit_ffc(slab->ph_bitmap, slab->ph_ntotal, &i);
    bit_set(slab->ph_bitmap, i);
    ptr = slab_item_at(slab, i);
    debug("slab_alloc: allocated item %p at slab %p, index %d", ptr, slab, i);

    if (++slab->ph_nused == slab->ph_ntotal) {
      /* We've allocated last item from non-empty slab
       * -> move it to the list of full slabs. */
      LIST_REMOVE(slab, ph_link);
      LIST_INSERT_HEAD(&pool->pp_full_slabs, slab, ph_link);
    }

    pool->pp_nused++;
    pool->pp_nmaxused = max(pool->pp_nmaxused, pool->pp_nused);
  }

  /* Create redzone after the item. */
  kasan_mark(ptr, pool->pp_itemsize, pool->pp_itemsize + pool->pp_redzone,
             KASAN_CODE_POOL_OVERFLOW);
  if (flags & M_ZERO)
    bzero(ptr, pool->pp_itemsize);

  return ptr;
}

/* TODO: destroy empty slabs when their number reaches a certain threshold
 * (maybe leave one). */
static void _pool_free(pool_t *pool, void *ptr) {
  assert(mtx_owned(&pool->pp_mtx));

  debug("pool_free: pool = %p, ptr = %p", pool, ptr);

  vm_page_t *pg = kva_find_page((vaddr_t)ptr);
  assert(pg != NULL);
  slab_t *slab = pg->slab;

  unsigned index =
    ((intptr_t)ptr - (intptr_t)slab->ph_items) / slab->ph_itemsize;
  bitstr_t *bitmap = slab->ph_bitmap;

  if (!bit_test(bitmap, index))
    panic("Double free detected in '%s' pool at %p!", pool->pp_desc, ptr);
  bit_clear(bitmap, index);

  if (slab->ph_nused == slab->ph_ntotal) {
    LIST_REMOVE(slab, ph_link);
    LIST_INSERT_HEAD(&pool->pp_part_slabs, slab, ph_link);
  }

  if (--slab->ph_nused == 0) {
    LIST_REMOVE(slab, ph_link);
    LIST_INSERT_HEAD(&pool->pp_empty_slabs, slab, ph_link);
  }

  pool->pp_nused--;

  debug("pool_free: freed item %p at slab %p, index %d", ptr, slab, index);
}

void pool_free(pool_t *pool, void *ptr) {
  SCOPED_MTX_LOCK(&pool->pp_mtx);

  kasan_mark_invalid(ptr, pool->pp_itemsize + pool->pp_redzone,
                     KASAN_CODE_POOL_FREED);
  kasan_quar_additem(&pool->pp_quarantine, pool, ptr);
#if !KASAN
  /* Without KASAN, call regular free method */
  _pool_free(pool, ptr);
#endif /* !KASAN */
}

static void pool_ctor(pool_t *pool) {
  LIST_INIT(&pool->pp_empty_slabs);
  LIST_INIT(&pool->pp_full_slabs);
  LIST_INIT(&pool->pp_part_slabs);
  mtx_init(&pool->pp_mtx, 0);
}

static void destroy_slabs(pool_t *pool, slab_list_t *slabs) {
  slab_t *slab, *next;

  LIST_FOREACH_SAFE (slab, slabs, ph_link, next) {
    klog("destroy_slab: pool = %p, slab = %p", pool, slab);

    pool->pp_ntotal -= slab->ph_ntotal;
    pool->pp_npages -= slab->ph_size;

    LIST_REMOVE(slab, ph_link);

    for (size_t i = 0; i < slab->ph_size; i += PAGESIZE) {
      vm_page_t *pg = kva_find_page((vaddr_t)slab + i);
      assert(pg != NULL);
      assert(pg->slab == slab);
      pg->slab = NULL;
    }

    kmem_free(slab, slab->ph_size);
  }
}

static void pool_dtor(pool_t *pool) {
  destroy_slabs(pool, &pool->pp_empty_slabs);
  destroy_slabs(pool, &pool->pp_part_slabs);
  destroy_slabs(pool, &pool->pp_full_slabs);
  klog("destroyed pool '%s' at %p", pool->pp_desc, pool);
}

void _pool_init(pool_t *pool, pool_init_t *args) {
  const char *desc = args->desc;
  size_t size = args->size;
  size_t alignment = max(args->alignment, PI_ALIGNMENT);
  size_t slabsize =
    args->slabsize ? roundup2(args->slabsize, PAGESIZE) : PAGESIZE;
  assert(desc);
  assert(size);
  assert(powerof2(alignment));
  assert(slabsize);

  pool_ctor(pool);
  pool->pp_desc = desc;
  pool->pp_alignment = alignment;
  pool->pp_slabsize = slabsize;
#if KASAN
  /* the alignment is within the redzone */
  pool->pp_itemsize = size;
  pool->pp_redzone = align(size + KASAN_POOL_REDZONE_SIZE, alignment) - size;
#else /* !KASAN */
  /* no redzone, we have to align the size itself */
  pool->pp_itemsize = align(size, alignment);
#endif
  kasan_quar_init(&pool->pp_quarantine, (quar_free_t)_pool_free);
  klog("initialized '%s' pool at %p (item size = %d)", pool->pp_desc, pool,
       pool->pp_itemsize);
  WITH_MTX_LOCK (&pool_list_lock)
    TAILQ_INSERT_TAIL(&pool_list, pool, pp_link);
}

void init_pool(void) {
  INVOKE_CTORS(pool_ctor_table);
}

void pool_add_page(pool_t *pool, void *page, size_t size) {
  SCOPED_MTX_LOCK(&pool->pp_mtx);
  add_slab(pool, page, size);
}

pool_t *_pool_create(pool_init_t *args) {
  pool_t *pool = kmalloc(M_POOL, sizeof(pool_t), M_ZERO | M_NOWAIT);
  _pool_init(pool, args);
  return pool;
}

void pool_destroy(pool_t *pool) {
  WITH_MTX_LOCK (&pool_list_lock)
    TAILQ_REMOVE(&pool_list, pool, pp_link);
  WITH_MTX_LOCK (&pool->pp_mtx)
    /* Lock needed as the quarantine may call _pool_free! */
    kasan_quar_releaseall(&pool->pp_quarantine);
  pool_dtor(pool);
  kfree(M_POOL, pool);
}
