#include <memory.h>
#include "malunal/allocators/arena.h"

typedef struct {
  malunal_uint32_t size;
  malunal_uint32_t used;
  malunal_mptr_t   next;
  malunal_uint8_t  bytes[];
} region_t;

typedef       region_t* region_mptr_t;
typedef const region_t* region_iptr_t;

typedef struct {
  allocator_t   allocator;
  region_mptr_t context;
} impl_t;

typedef       impl_t* impl_mptr_t;
typedef const impl_t* impl_iptr_t;
_Static_assert(
  sizeof(arena_allocator_t) == sizeof(impl_t),
  "Arena allocator must be the size of its implementation"
);

error_t
create_region(region_mptr_t* region) {
  allocator_mptr_t allocator = platform_allocator();
  malunal_size_t   rounded   = align_to_page(MALUNAL_ARENA_REGION_SIZE);
  malunal_mptr_t*  address   = (malunal_mptr_t*)region;
  error_t          result    = allocator_acquire(allocator, rounded, address);

  if (result.domain != null)
    return result;

  (*region)->size = rounded;
  (*region)->used = sizeof(region_t);
  (*region)->next = null;
  return NO_ERROR;
}

error_t
delete_region(region_mptr_t region) {
  if (region == null)
    return NO_ERROR;

  error_t result = delete_region(region->next);
  if (result.domain != null)
    return result;

  region->next = null;
  allocator_mptr_t allocator = platform_allocator();
  return allocator_dispose(allocator, region, region->size);
}

error_t
region_acquire(
  region_mptr_t   region,
  malunal_size_t  size,
  malunal_mptr_t* out
) {
  malunal_uint32_t offset =
    region->used - sizeof(region_t);
  *out = region->bytes + offset;
  region->used += size;
  return NO_ERROR;
}

static const
allocator_vtable_t arena_allocator_vtable = {
  .acquire = (allocator_acquire_pfn_t)&arena_allocator_acquire,
  .dispose = (allocator_dispose_pfn_t)&arena_allocator_dispose
};

error_t
arena_allocator_size(
  arena_allocator_iptr_t allocator,
  malunal_uint32_t*      out
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_iptr_t    self   = (impl_iptr_t)allocator;
  region_iptr_t  region = self->context;
  malunal_size_t total  = 0;
  while (region != null) {
    total += region->size;
    region = region->next;
  }

  *out = total;
  return NO_ERROR;
}

error_t
arena_allocator_used(
  arena_allocator_iptr_t allocator,
  malunal_uint32_t*      out
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_iptr_t    self   = (impl_iptr_t)allocator;
  region_iptr_t  region = self->context;
  malunal_size_t total  = 0;
  while (region != null) {
    total += region->used;
    region = region->next;
  }

  *out = total;
  return NO_ERROR;
}

error_t
arena_allocator_init(
  malunal_size_t         capacity,
  arena_allocator_mptr_t allocator
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  region_mptr_t region = null;
  error_t       result = create_region(&region);
  if (result.domain != null)
    return result;

  impl_mptr_t self = (impl_mptr_t)allocator;
  self->allocator = (allocator_t){ &arena_allocator_vtable };
  self->context   = region;
  return NO_ERROR;
}

error_t
arena_allocator_acquire(
  arena_allocator_mptr_t allocator,
  malunal_size_t         size,
  malunal_mptr_t*        out
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_mptr_t   self   = (impl_mptr_t)allocator;
  region_mptr_t region = self->context;
  while (region != null) {
    if (region->size - region->used >= size)
      return region_acquire(region, size, out);
    region = region->next;
  }

  error_t result = create_region((region_mptr_t*)&region->next);
  return result.domain == null
    ? region_acquire(region->next, size, out)
    : result;
}

error_t
arena_allocator_dispose(
  arena_allocator_mptr_t allocator,
  malunal_mptr_t         address,
  malunal_size_t         size
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_mptr_t      self = (impl_mptr_t)allocator;
  region_mptr_t    reg  = self->context;
  malunal_uint8_t* addr = address;
  malunal_uint8_t* beg  = reg->bytes;
  malunal_uint8_t* end  = reg->bytes + reg->size - sizeof(region_t);
  while (true) {
    if (addr < beg || addr > end)
      return (error_t) {
        .domain = &ERROR_DOMAIN_ALLOCATOR_T,
        .code   = ALLOCATOR_ERROR_NOT_MY_ADDRESS
      };

    reg = reg->next;
    if (reg == null)
      break;
    beg = reg->bytes;
    end = reg->bytes + reg->size - sizeof(region_t);
  }
  return NO_ERROR;
}

error_t
arena_allocator_reset(
  arena_allocator_mptr_t allocator
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_mptr_t   self   = (impl_mptr_t)allocator;
  region_mptr_t region = self->context;
  while (region != null) {
    region->used = sizeof(region_t);
    region = region->next;
  }
}

error_t
arena_allocator_free(
  arena_allocator_mptr_t allocator
) {
  if (allocator == null)
    return (error_t) {
      .domain = &ERROR_DOMAIN_ALLOCATOR_T,
      .code   = ALLOCATOR_ERROR_NULL_ALLOCATOR
    };

  impl_mptr_t self   = (impl_mptr_t)allocator;
  error_t     result = delete_region(self->context);
  if (result.domain != null)
    return result;

  self->context = null;
  return NO_ERROR;
}