#include "allocator.h"

#include <assert.h>
#include <string.h>

#ifndef KIT_DISABLE_SYSTEM_MALLOC
#  include <stdlib.h>
#endif

static void *kit_allocate_default_(i32 request, i64 size,
                                   i64 previous_size, void *pointer) {
#ifndef KIT_DISABLE_SYSTEM_MALLOC
  switch (request) {
    case KIT_ALLOCATE:
    case KIT_ALLOCATE_ZERO: {
      assert(size >= 0);
      assert(previous_size == 0);
      assert(pointer == NULL);

      if (size <= 0)
        return NULL;

      void *p = malloc(size);

      if (p != NULL && request == KIT_ALLOCATE_ZERO)
        memset(p, 0, size);

      return p;
    }

    case KIT_REALLOCATE:
    case KIT_REALLOCATE_ZERO: {
      assert(size >= 0);
      assert(previous_size != 0 || pointer == NULL);
      assert(previous_size == 0 || pointer != NULL);

      if (previous_size == 0 && pointer != NULL)
        return NULL;
      if (previous_size != 0 && pointer == NULL)
        return NULL;
      if (size == previous_size)
        return pointer;

      u8 *p = NULL;

      if (size > 0) {
        p = (u8 *) malloc(size);

        if (p != NULL) {
          if (size > 0 && previous_size > 0)
            memcpy(p, pointer,
                   size < previous_size ? size : previous_size);
          if (request == KIT_REALLOCATE_ZERO && size > previous_size)
            memset(p + previous_size, 0, size - previous_size);
        }
      }

      free(pointer);

      return p;
    }

    case KIT_DEALLOCATE:
      assert(size == 0);
      assert(pointer != NULL);
      if (pointer != NULL)
        free(pointer);
      return NULL;

    case KIT_DEALLOCATE_ALL:
      //  Do nothing.
      //
      return NULL;

    default:;
  }
#endif

  assert(0);
  return NULL;
}

static void *kit_allocate_from_buffer_(kit_allocator_t *alloc,
                                       i32 request, i64 size,
                                       i64   previous_size,
                                       void *pointer) {
  assert(alloc != NULL);
  assert(pointer == NULL || pointer < alloc->data);

  if (alloc == NULL)
    return NULL;

  switch (request) {
    case KIT_ALLOCATE:
    case KIT_ALLOCATE_ZERO: {
      assert(size >= 0);
      assert(previous_size == 0);
      assert(pointer == NULL);

      if (size <= 0)
        return NULL;

      if (alloc->size < size)
        return NULL;

      void *p = alloc->data;
      alloc->bytes += size;
      alloc->size -= size;

      if (request == KIT_ALLOCATE_ZERO)
        memset(p, 0, size);

      return p;
    }

    case KIT_REALLOCATE:
    case KIT_REALLOCATE_ZERO: {
      assert(size >= 0);
      assert(previous_size != 0 || pointer == NULL);
      assert(previous_size == 0 || pointer != NULL);

      if (size <= 0)
        return NULL;
      if (size <= previous_size)
        return pointer;
      if (previous_size == 0 && pointer != NULL)
        return NULL;
      if (previous_size != 0 && pointer == NULL)
        return NULL;

      if ((u8 *) pointer + previous_size == alloc->data) {
        if (alloc->size < size - previous_size)
          return NULL;
        alloc->bytes += size - previous_size;
        alloc->size -= size - previous_size;
        return pointer;
      }

      if (alloc->size < size)
        return NULL;

      u8 *p = alloc->bytes;
      alloc->bytes += size;
      alloc->size -= size;

      if (previous_size > 0)
        memcpy(p, pointer, previous_size);
      if (request == KIT_REALLOCATE_ZERO)
        memset(p + previous_size, 0, size - previous_size);

      return p;
    }

    case KIT_DEALLOCATE:
    case KIT_DEALLOCATE_ALL: return NULL;

    default:;
  }

  assert(0);
  return NULL;
}

#ifndef KIT_ENABLE_CUSTOM_ALLOC_DISPATCH
void *kit_alloc_dispatch(kit_allocator_t *alloc, i32 request,
                         i64 size, i64 previous_size, void *pointer) {
  if (alloc == NULL)
    return kit_allocate_default_(request, size, previous_size,
                                 pointer);

  switch (alloc->type) {
    case KIT_ALLOC_TYPE_DEFAULT:
      return kit_allocate_default_(request, size, previous_size,
                                   pointer);

    case KIT_ALLOC_TYPE_BUFFER:
      return kit_allocate_from_buffer_(alloc, request,
                                       //  alignment
                                       ((size + 7) / 8) * 8,
                                       previous_size, pointer);

    default:;
  }

  return NULL;
}
#endif

kit_allocator_t kit_alloc_default(void) {
  return (kit_allocator_t) { .type = KIT_ALLOC_TYPE_DEFAULT,
                             .size = 0,
                             .data = NULL };
}

kit_allocator_t kit_alloc_buffer(i64 size, void *buffer) {
  return (kit_allocator_t) { .type = KIT_ALLOC_TYPE_BUFFER,
                             .size = size,
                             .data = buffer };
}