//
// Created by epagris on 2022.10.20..
//

#include <stdbool.h>
#include "memory_pool.h"
#include "utils.h"

MP *mp_init(uint8_t *p, uint32_t size) {
    ASSERT_BAD_ALIGN(p); // check for alignment
    size = FLOOR_TO_4(size); // force alignment on size
    MP *pool = (MP *) p; // fill in properties
    pool->p = p + sizeof(MP); // compute beginning of the allocatable area
    pool->poolSize = size; // save total pool size
    pool->blockRegistry = (MPAllocRecord *) (p + size - sizeof(MPAllocRecord)); // determine block registry address (points to the TOPMOST record, NOT one block further!)
    pool->freeSpace = ((uint8_t *) pool->blockRegistry) - pool->p - 1 * sizeof(MPAllocRecord); // calculate free space size (with SENTRY and first FREE block)

    // place sentry element
    pool->blockRegistry[-1].addrStart = 0;
    pool->blockRegistry[-1].size = 0;
    pool->blockRegistry[-1].type = MPRT_SENTRY;

    // place element record holding info on the single, undivided, unallocated block
    pool->blockRegistry[0].addrStart = pool->p;
    pool->blockRegistry[0].size = pool->freeSpace;
    pool->blockRegistry[0].type = MPRT_FREE;

    pool->blockRecCnt = 2; // sentry and free block

    return pool;
}

// Break allocation table into two halves and insert BELOW the given position an empty record.
// Also maintain global free space counter.
static void mp_break_alloc_table(MP *mp, MPAllocRecord * ar) {
    // shift the registry below the given block
    MPAllocRecord *rec = mp->blockRegistry - (mp->blockRecCnt); // bottom of the FILLED registry (there's one unfilled entry below this point, the new record)
    while (rec != ar) {
        *(rec - 1) = *(rec);
        rec++;
    }

    // decrease last free block (adjacent to block registry) size with the increase in the block registry
    mp->blockRegistry[0].size -= sizeof(MPAllocRecord);

    // decrease free size with the increase in the registry size
    mp->freeSpace -= sizeof(MPAllocRecord);

    // increase record count
    mp->blockRecCnt++;
}

// check that allocation table can grow downwards
static bool mp_can_alloc_table_grow(const MP * mp) {
    // examine, that an allocated block was not blocking registry table downward growth
    return mp->blockRegistry[0].type == MPRT_FREE;
}

uint8_t *mp_alloc(MP *mp, uint32_t size) {
    // make the allocation from the beginning of the smallest suitable (large enough)
    // contiguous block

    // round size to make it divisible by 4
    size = CEIL_TO_4(size);

    // badness = area left unclaimed on the candidate free block after allocating the requested block
    MPAllocRecord *bestBlock = NULL;
    uint32_t leastBadness = ~0;
    MPAllocRecord *recIter = mp->blockRegistry; // allocation record
    while (recIter->type != MPRT_SENTRY && leastBadness > 0) { // at badness = 0 just break, since it's a perfectly fitting block
        // look for suitable free block
        if (recIter->type == MPRT_FREE && recIter->size >= size) {
            uint32_t iterBadness = recIter->size - size;
            if (iterBadness < leastBadness) { // calculate badness
                bestBlock = recIter;
                leastBadness = iterBadness;
            }
        }
        recIter--; // step to next block
    }

    // allocate block
    uint8_t *ptr = NULL;
    if (leastBadness == 0) { // just change block registry class if block perfectly fits
        bestBlock->type = MPRT_ALLOCATED;
        ptr = bestBlock->addrStart;
    } else { // if there are some bytes left between allocated blocks

        // examine, that an allocated block was not blocking registry table downward growth
        if (!mp_can_alloc_table_grow(mp)) {
            return NULL;
        }

        // break the allocation table, shift the registry below best block
        mp_break_alloc_table(mp, bestBlock);

        // store information on allocated
        MPAllocRecord *allocated = bestBlock - 1; // don't copy best block!
        allocated->type = MPRT_ALLOCATED;
        allocated->size = size;
        allocated->addrStart = bestBlock->addrStart;
        ptr = allocated->addrStart;

        // shrink the remaining free block size
        bestBlock->size -= size;
        bestBlock->addrStart += size;
    }

    // decrease free space size with the allocated block size
    mp->freeSpace -= size;

    return ptr;
}

// join adjacent free blocks
static void mp_join_free_blocks(MP *mp) {
    MPAllocRecord *recIter = mp->blockRegistry;
    while (recIter->type != MPRT_SENTRY) {
        if (recIter->type == MPRT_FREE && (recIter - 1)->type == MPRT_FREE) { // if two adjacent free blocks have been found...
            // join the blocks
            recIter->size += (recIter - 1)->size;
            recIter->addrStart = (recIter - 1)->addrStart;

            // shift block below joined blocks one upper
            MPAllocRecord *joinIter = (recIter - 1);
            while (joinIter->type != MPRT_SENTRY) {
                *joinIter = *(joinIter - 1);
                joinIter--;
            }
            mp->freeSpace += sizeof(MPAllocRecord);
            mp->blockRegistry[0].size += sizeof(MPAllocRecord); // grow the last record size
            mp->blockRecCnt--;
        } else {
            recIter--;
        }
    }
}

void mp_free(MP *mp, const uint8_t *p) {
    if (p == NULL) { // don't do anything with a NULL pointer
        return;
    }

    // look for registry record
    bool success = false;
    MPAllocRecord *recIter = mp->blockRegistry;
    while (recIter->type != MPRT_SENTRY) {
        if ((recIter->type == MPRT_ALLOCATED) && (recIter->addrStart == p)) { // ...block found
            recIter->type = MPRT_FREE;
            mp->freeSpace += recIter->size;
            success = true;
            break;
        }
        recIter--;
    }

    if (success) {
        mp_join_free_blocks(mp);
    }

    if (!success) {
        WARNING("Possible double free!\n");
    }
}

void mp_report(MP *mp) {
    INFO("#     TYPE    BEGIN           SIZE\n");
    MPAllocRecord *recIter = mp->blockRegistry;
    uint32_t bi = 0;
    while (recIter->type != MPRT_SENTRY) {
        INFO("%05u %s %p  %u\n", bi, recIter->type == MPRT_ALLOCATED ? "ALLOC  " : "FREE   ",
             recIter->addrStart, recIter->size);
        recIter--;
        bi++;
    }
    INFO("%05u %s\n", bi, "SENTRY");

    INFO("----------------------\n");
    INFO("Used: %u (mgmt: %u)\nFree: %u of %u\n\n", (mp->poolSize - mp->freeSpace),
         (mp->blockRecCnt) * sizeof(MPAllocRecord), mp->freeSpace, mp->poolSize);
}

uint32_t mp_largest_free_block_size(MP *mp) {
    MPAllocRecord *recIter = mp->blockRegistry;
    if (recIter->type == MPRT_ALLOCATED) { // if topmost block is allocated, then registry table cannot be grown
        return 0;
    }

    uint32_t largestFreeSize = 0;
    while (recIter->type != MPRT_SENTRY) {
        largestFreeSize = (recIter->size > largestFreeSize) ? recIter->size : largestFreeSize;
        recIter--;
    }
    return largestFreeSize;
}

void mp_foreach_block(MP *mp, MPForeachFn *fn, void *userData, bool inclFree) {
    MPAllocRecord *recIter = mp->blockRegistry;
    while (recIter->type != MPRT_SENTRY) {
        if (recIter->type != MPRT_FREE || inclFree) {
            fn(mp, recIter, userData);
        }
        recIter--;
    }
}

MPAllocRecord *mp_get_block_by_start_addr(MP *mp, const uint8_t *startAddr) {
    MPAllocRecord *recIter = mp->blockRegistry, *rec = NULL;
    while (recIter->type != MPRT_SENTRY) {
        if (recIter->type == MPRT_ALLOCATED && recIter->addrStart == startAddr) {
            rec = recIter;
        }
        recIter--;
    }
    return rec;
}

bool mp_shrink_block(MP *mp, uint8_t *beginAddr, uint32_t size, MPShrinkStrategy strategy) {
    // verify that it is possible to shrink at all
    if (!mp_can_alloc_table_grow(mp)) {
        return false;
    }

    // check that size is divisible by 4
    if (size & 0b11) {
        MSG("Shrinking is not possible but only by size divisible by 4. (So that 32-bits alignments are maintained.)\n");
        return false;
    }

    // get the block we want to shrink
    MPAllocRecord * block = mp_get_block_by_start_addr(mp, beginAddr);
    if (block == NULL) { // no block found, cannot shrink
        return false;
    }

    // ...block found...

    // It's certain, that (block + 1) is still in the allocation table, since, at least
    // a single FREE block is at the top of the table. (see: mp_can_alloc_table_grow())
    MPAllocRecord * splitBelow = (strategy == MP_SHRINK_BEGIN_FIXED) ? (block + 1) : (block);
    MPAllocRecord * newRec = splitBelow - 1;

    // break up the table
    mp_break_alloc_table(mp, splitBelow);

    // Beginning remains fixed: ]------| -> ]----|--|
    if (strategy == MP_SHRINK_BEGIN_FIXED) {
        // shrink old block size
        block->size -= size;

        // fill new block properties
        newRec->size = size;
        newRec->type = MPRT_FREE;
        newRec->addrStart = block->addrStart + block->size;
    } else { // End remains fixed: |-------[ -> |---|----[
        // fill new block properties
        newRec->size = size;
        newRec->type = MPRT_FREE;
        newRec->addrStart = block->addrStart; // copy the start address

        // shrink old block and move it's beginning
        block->size -= size;
        block->addrStart += size;
    }

    // also release size from the global counter
    mp->freeSpace += size;

    return true;
}