#include <stdint.h>
#include <stddef.h>
#include <stdbool.h>
#include "bitmap.h"
#include "util/util.h"
#include "kprintf.h"

void *bitmap_toptr(BitMap *bm, size_t block) {
  uint8_t *ptr = (uint8_t *)(bm->mem_start + (block * BITMAP_BLOCK_SIZE));
  return (void *)ptr;
}

size_t bitmap_toblock(BitMap *bm, void *ptr) {
  uint8_t *p = ptr;
  return (size_t)(p - bm->mem_start) / BITMAP_BLOCK_SIZE;
}

size_t bitmap_toblock_roundup(BitMap *bm, void *ptr) {
  uint8_t *p = ptr;
  return (size_t)_DIV_ROUNDUP((size_t)(p - bm->mem_start), BITMAP_BLOCK_SIZE);
}

size_t bitmap_calcsize(size_t total) {
  size_t nblocks = _DIV_ROUNDUP(total, BITMAP_BLOCK_SIZE);
  size_t nbytes = _DIV_ROUNDUP(nblocks, 8);
  return nbytes;
}

bool bitmap_get(BitMap *bm, size_t block) {
  size_t elem = block / BITMAP_BLOCKS_PER_BYTE;
  size_t off = block % BITMAP_BLOCKS_PER_BYTE;
  return (bm->map[elem] & (1 << off)) != 0;
}

void bitmap_set(BitMap *bm, size_t block, bool value) {
  size_t elem = block / BITMAP_BLOCKS_PER_BYTE;
  size_t off = block % BITMAP_BLOCKS_PER_BYTE;
  if (value) {
    bm->map[elem] |= (1 << off);
  } else {
    bm->map[elem] &= ~(1 << off);
  }
}

void bitmap_markblocks(BitMap *bm, size_t start, size_t size, bool value) {
  if (!value && start < bm->last_deep_frag) {
    bm->last_deep_frag = start;
  }

  for (size_t i = start; i < start + size; i++) {
    bitmap_set(bm, i, value);
  }

  bm->alloc_blocks += value ? size : -size;
}

void bitmap_markregion(BitMap *bm, void *baseptr, size_t size, bool is_used) {
  size_t base, size1;

  if (is_used) {
    base = bitmap_toblock(bm, baseptr);
    size1 = _DIV_ROUNDUP(size, BITMAP_BLOCK_SIZE);
  } else {
    base = bitmap_toblock(bm, baseptr);
    size1 = size / BITMAP_BLOCK_SIZE;
  }
  bitmap_markblocks(bm, base, size1, is_used);
}

size_t bitmap_freeregion(BitMap *bm, size_t blocks) {
  size_t curregstart = bm->last_deep_frag;
  size_t curregsize = 0;
  for (size_t i = curregstart; i < bm->nblocks; i++) {
    if (bitmap_get(bm, i)) {
      curregsize = 0;
      curregstart = i + 1;
    } else {
      if (blocks == 1) {
        bm->last_deep_frag = curregstart + 1;
      }

      curregsize++;
      if (curregsize >= blocks) {
        return curregstart;
      }
    }
  }
  return BITMAP_INVALID_BLOCK;
}

void *bitmap_alloc(BitMap *bm, size_t blocks) {
  if (!blocks) {
    return NULL;
  }

  size_t pickedreg = bitmap_freeregion(bm, blocks);
  if (pickedreg == BITMAP_INVALID_BLOCK) {
    return NULL;
  }

  bitmap_markblocks(bm, pickedreg, blocks, 1);
  return bitmap_toptr(bm, pickedreg);
}

void bitmap_free(BitMap *bm, void *base, size_t blocks) {
  bitmap_markregion(bm, base, BITMAP_BLOCK_SIZE * blocks, 0);
}
