/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 *
 *    Copyright 2018 (c) Stefan Profanter, fortiss GmbH
 */

/**
 * This memory manager allows to manually reduce the available RAM.
 *
 * It keeps track of malloc and free calls and counts the available RAM.
 * If the requested memory allocation results in hitting the limit,
 * it will return NULL and prevent the new allocation.
 */

#include "custom_memory_manager.h"
#include <stdlib.h>
#include <stdio.h>
#include <stdint.h>
#include <limits.h>

#include <pthread.h>

pthread_mutex_t mutex;

struct UA_mm_entry {
    size_t size;
    void* address;
    struct UA_mm_entry *next;
    struct UA_mm_entry *prev;
};

unsigned long long totalMemorySize = 0;
unsigned long long memoryLimit = ULONG_MAX;

/* Head and Tail of the double linked list */
struct UA_mm_entry *address_map_first = NULL;
struct UA_mm_entry *address_map_last = NULL;

/**
 * Add address entry to the linked list after it was allocated.
 *
 * @param size Size of the allocated memory block
 * @param addr Address of the allocated memory block
 * @return 1 on success, 0 if it failed
 */
static int addToMap(size_t size, void *addr) {
    struct UA_mm_entry *newEntry = (struct UA_mm_entry*)malloc(sizeof(struct UA_mm_entry));
    if(!newEntry) {
        //printf("MemoryManager: Could not allocate memory");
        return 0;
    }
    newEntry->size = size;
    newEntry->address = addr;
    newEntry->next = NULL;
    pthread_mutex_lock(&mutex);
    newEntry->prev = address_map_last;
    if(address_map_last)
        address_map_last->next = newEntry;
    address_map_last = newEntry;
    totalMemorySize += size;
    if(address_map_first == NULL)
        address_map_first = newEntry;
    pthread_mutex_unlock(&mutex);
    //printf("Total size (malloc): %lld And new address: %p Entry %p\n", totalMemorySize, addr, (void*)newEntry);

    return 1;
}

/**
 * Remove entry from the list before the memory block is freed.
 *
 * @param addr Address of the memory block which is freed.
 *
 * @return 1 on success, 0 if the memory block was not found in the list.
 */
static int removeFromMap(void *addr) {
    if(addr == NULL)
        return 1;

    pthread_mutex_lock(&mutex);

    struct UA_mm_entry *e = address_map_last;

    while (e) {
        if(e->address == addr) {
            if(e == address_map_first)
                address_map_first = e->next;
            if(e == address_map_last)
                address_map_last = e->prev;
            if(e->prev)
                e->prev->next = e->next;
            if(e->next)
                e->next->prev = e->prev;
            totalMemorySize -= e->size;

            //printf("Total size (free): %lld after addr %p and deleting %p\n", totalMemorySize, addr, (void*)e);
            free(e);
            pthread_mutex_unlock(&mutex);
            return 1;
        }
        e = e->prev;
    }
    pthread_mutex_unlock(&mutex);
    //printf("MemoryManager: Entry with address %p not found", addr);
    return 0;
}

static void *UA_memoryManager_malloc(size_t size) {
    if(totalMemorySize + size > memoryLimit)
        return NULL;
    void *addr = malloc(size);
    if(!addr)
        return NULL;
    addToMap(size, addr);
    return addr;
}

static void *UA_memoryManager_calloc(size_t num, size_t size) {
    if(totalMemorySize + (size * num) > memoryLimit)
        return NULL;
    void *addr = calloc(num, size);
    if(!addr)
        return NULL;
    addToMap(size*num, addr);
    return addr;
}

static void *UA_memoryManager_realloc(void *ptr, size_t new_size) {
    removeFromMap(ptr);
    if(totalMemorySize + new_size > memoryLimit)
        return NULL;
    void *addr = realloc(ptr, new_size);
    if(!addr)
        return NULL;
    addToMap(new_size, addr);
    return addr;

}

static void UA_memoryManager_free(void* ptr) {
    removeFromMap(ptr);
    free(ptr);
}

void UA_memoryManager_setLimit(unsigned long long newLimit) {
    memoryLimit = newLimit;
    //printf("MemoryManager: Setting memory limit to %lld\n", newLimit);

    UA_mallocSingleton = UA_memoryManager_malloc;
    UA_freeSingleton = UA_memoryManager_free;
    UA_callocSingleton = UA_memoryManager_calloc;
    UA_reallocSingleton = UA_memoryManager_realloc;
}

int UA_memoryManager_setLimitFromLast4Bytes(const uint8_t *data, size_t size) {
    UA_mallocSingleton = UA_memoryManager_malloc;
    UA_freeSingleton = UA_memoryManager_free;
    UA_callocSingleton = UA_memoryManager_calloc;
    UA_reallocSingleton = UA_memoryManager_realloc;
    if(size <4)
        return 0;
    // just cast the last 4 bytes to uint32
    const uint32_t *newLimit = (const uint32_t*)(uintptr_t)&(data[size-4]);
    uint32_t limit;
    // use memcopy to avoid asan complaining on misaligned memory
    memcpy(&limit, newLimit, sizeof(uint32_t));
    memoryLimit = limit;
    return 1;
}