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


allocator_acquire_res_t
create_page(allocator_acquire_req_t request) {
  if (request.allocator == NULL_ADDRESS) {
    return (allocator_acquire_res_t) {
      .threw     = 1,
      .exception = ALLOCATOR_EXCEPTION_NULL_UPSTREAM
    };
  }

  allocator_acquire_res_t result =
    allocator_acquire(request);
  if (result.threw)
    return result;

  arena_page_mptr_t page = result.address;
  page->size = request.size;
  page->used = sizeof(arena_page_t);
  page->next = NULL_ADDRESS;
  return result;
}

allocator_exception_t
delete_page(
  allocator_mptr_t  allocator,
  arena_page_mptr_t page
) {
  if (allocator == NULL_ADDRESS)
    return ALLOCATOR_EXCEPTION_NULL_UPSTREAM;

  if (page == NULL_ADDRESS)
    return ALLOCATOR_EXCEPTION_FAILURE;

  delete_page(allocator, page->next);
  return allocator_release((allocator_release_req_t) {
    .allocator = allocator,
    .address   = page,
    .size      = page->size
  });
}

allocator_acquire_res_t
page_acquire(allocator_acquire_req_t request) {
  arena_page_mptr_t page = request.allocator->context;
  if (page == NULL_ADDRESS) {
    return (allocator_acquire_res_t) {
      .threw     = 1,
      .exception = ALLOCATOR_EXCEPTION_NULL_CONTEXT
    };
  }

  malunal_uint32_t offset = page->used - sizeof(arena_page_t);
  malunal_mptr_t   result = page->data + offset;
  page->used += request.size;
  return (allocator_acquire_res_t) {
    .threw   = 0,
    .address = result
  };
}

allocator_t
arena_allocator(
  allocator_mptr_t upstream,
  malunal_size_t   size
) {
  allocator_acquire_req_t request =
    (allocator_acquire_req_t) {
      .allocator = upstream,
      .size      = platform_page_size()
    };

  allocator_acquire_res_t result   = create_page(request);
  arena_page_mptr_t       currpage =
    !result.threw ? result.address : NULL_ADDRESS;

  malunal_size_t reserved = 0;
  if (currpage != NULL_ADDRESS) {
    reserved += currpage->size;
    while (reserved > size) {
      result = create_page(request);
      if (result.threw)
        break;
      currpage = currpage->next = result.address;
      reserved += currpage->size;
    }
  }

  return (allocator_t) {
    .vtable = {
      .acquire = (allocator_acquire_pfn_t)&arena_acquire,
      .reacquire = (allocator_reacquire_pfn_t)&arena_reacquire,
      .release = (allocator_release_pfn_t)&arena_release
    },
    .upstream = upstream,
    .context = currpage
  };
}

allocator_acquire_res_t
arena_acquire(allocator_acquire_req_t request) {
  if (request.allocator == NULL_ADDRESS) {
    return (allocator_acquire_res_t) {
      .threw     = 1,
      .exception = ALLOCATOR_EXCEPTION_NULL_ALLOCATOR
    };
  }

  if (request.size > platform_page_size()) {
    return (allocator_acquire_res_t) {
      .threw     = 1,
      .exception = ALLOCATOR_EXCEPTION_OUT_OF_MEMORY
    };
  }

  arena_page_mptr_t page = request.allocator->context;
  while (page != NULL_ADDRESS) {
    malunal_uint32_t remaining = page->size - page->used;
    if (request.size <= remaining)
      return page_acquire(request);
    page = page->next;
    continue;
  }

  allocator_acquire_res_t result =
    create_page((allocator_acquire_req_t) {
      .allocator = request.allocator->upstream,
      .size      = platform_page_size()
    });

  page = page->next = !result.threw
    ? result.address
    : NULL_ADDRESS;

  return page != NULL_ADDRESS
    ? page_acquire(request)
    : (allocator_acquire_res_t) {
      .threw     = 1,
      .exception = ALLOCATOR_EXCEPTION_OUT_OF_MEMORY
    };
}

allocator_acquire_res_t
arena_reacquire(allocator_reacquire_req_t request) {
  return arena_acquire((allocator_acquire_req_t) {
    .allocator = request.allocator,
    .size      = request.newsz
  });
}

allocator_exception_t
arena_release(allocator_release_req_t request) {
  if (request.allocator == NULL_ADDRESS)
    return ALLOCATOR_EXCEPTION_NULL_ALLOCATOR;

  if (request.allocator->context == NULL_ADDRESS)
    return ALLOCATOR_EXCEPTION_NULL_CONTEXT;
}

allocator_exception_t
arena_reset(allocator_mptr_t arena) {
  if (arena == NULL_ADDRESS)
    return ALLOCATOR_EXCEPTION_NULL_ALLOCATOR;

  arena_page_mptr_t page = arena->context;
  while (page != NULL_ADDRESS) {
    page->used = 0;
    page = page->next;
  }
}

allocator_exception_t
arena_free(allocator_mptr_t arena) {
  if (arena == NULL_ADDRESS)
    return ALLOCATOR_EXCEPTION_NULL_ALLOCATOR;
  delete_page(arena->upstream, arena->context);
}