#include "best_fit_allocator.h"
#include <assert.h>
#include <stdalign.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

static size_t align_up(size_t n, size_t align) {
  return (n + align - 1) & ~(align - 1);
}

struct block {
  size_t size;
  int free;
  struct block *next;
  struct block *prev;
};

struct best_fit_allocator {
  void *buffer;
  size_t pool_size;
  struct block *head;
};

best_fit_allocator_t *bf_allocator_create(size_t pool_size) {
  if (pool_size <= sizeof(struct block))
    return NULL;
  size_t alignment = alignof(max_align_t);

  if (pool_size % alignment != 0) {
    pool_size = ((pool_size + alignment - 1) / alignment) * alignment;
  }

  void *buf = aligned_alloc(alignment, pool_size);
  if (!buf)
    return NULL;
  best_fit_allocator_t *a = malloc(sizeof(*a));

  if (!a) {
    free(buf);
    return NULL;
  }

  a->buffer = buf;
  a->pool_size = pool_size;
  a->head = (struct block *)buf;
  a->head->size = pool_size - sizeof(struct block);
  a->head->free = 1;
  a->head->next = NULL;
  a->head->prev = NULL;

  return a;
}

void bf_allocator_destroy(best_fit_allocator_t *allocator) {
  if (!allocator)
    return;
  free(allocator->buffer);
  free(allocator);
}

void *bf_allocator_alloc(best_fit_allocator_t *allocator, size_t size) {
  if (!allocator || size == 0)
    return NULL;

  size = align_up(size, alignof(max_align_t));
  struct block *best = NULL;

  for (struct block *it = allocator->head; it; it = it->next) {
    if (it->free && it->size >= size) {
      if (!best || it->size < best->size) {
        best = it;
        if (best->size == size)
          break;
      }
    }
  }

  if (!best)
    return NULL;

  if (best->size >= size + sizeof(struct block) + alignof(max_align_t)) {
    uint8_t *raw = (uint8_t *)best;
    struct block *next_blk =
        (struct block *)(raw + sizeof(struct block) + size);

    next_blk->size = best->size - size - sizeof(struct block);
    next_blk->free = 1;
    next_blk->next = best->next;
    next_blk->prev = best;

    if (best->next)
      best->next->prev = next_blk;

    best->next = next_blk;
    best->size = size;
  }

  best->free = 0;
  return (uint8_t *)best + sizeof(struct block);
}

void bf_allocator_free(best_fit_allocator_t *allocator, void *ptr) {
  if (!allocator || !ptr)
    return;

  uint8_t *start = (uint8_t *)allocator->buffer + sizeof(struct block);
  uint8_t *end = (uint8_t *)allocator->buffer + allocator->pool_size;

  assert(ptr >= (void *)start && ptr < (void *)end);
  struct block *blk = (struct block *)((uint8_t *)ptr - sizeof(struct block));
  assert(!blk->free);
  blk->free = 1;

  if (blk->next && blk->next->free) {
    blk->size += sizeof(struct block) + blk->next->size;
    blk->next = blk->next->next;
    if (blk->next)
      blk->next->prev = blk;
  }

  if (blk->prev && blk->prev->free) {
    blk->prev->size += sizeof(struct block) + blk->size;
    blk->prev->next = blk->next;
    if (blk->next)
      blk->next->prev = blk->prev;
  }
}

void bf_allocator_print(const best_fit_allocator_t *allocator) {
  if (!allocator)
    return;

  printf("Block list:\n");
  for (const struct block *it = allocator->head; it; it = it->next) {
    printf("  Block at %p | size: %zu bytes | %s\n", (void *)it, it->size,
           it->free ? "FREE" : "ALLOCATED");
  }
}
