#include "core/Utility.hpp"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <ctime>

#include "core/Logger.hpp"
#include "core/Thread.hpp"

static Core::ExitHandler exitHandler = nullptr;
static void* exitData = nullptr;
static Core::OutOfMemoryHandler outOfMemoryHandler = nullptr;
static void* outOfMemoryData = nullptr;

[[noreturn]] void Core::exitWithHandler(const char* file, int line, int value) {
    if(value != 0) {
        file = getShortFileName(file);
        LOG_ERROR("Exit from #:# with value #", file, line, value);
    }
    if(exitHandler != nullptr) {
        exitHandler(value, exitData);
    }
    exit(value);
}

void Core::setExitHandler(ExitHandler eh, void* data) {
    exitHandler = eh;
    exitData = data;
}

void Core::setOutOfMemoryHandler(OutOfMemoryHandler h, void* data) {
    outOfMemoryHandler = h;
    outOfMemoryData = data;
    std::set_new_handler([]() {
        if(outOfMemoryHandler != nullptr) {
            outOfMemoryHandler(outOfMemoryData);
        } else {
            LOG_ERROR("Out of memory");
            EXIT(1);
        }
    });
}

static void* exitOnNull(void* p, size_t n) {
    if(p == nullptr) {
        LOG_ERROR("Out of memory, requested '#' bytes", n);
        EXIT(1);
    }
    return p;
}

static void* realAllocate(size_t n) {
    void* p = malloc(n);
    while(p == nullptr && outOfMemoryHandler != nullptr) {
        outOfMemoryHandler(outOfMemoryData);
        p = malloc(n);
    }
    return exitOnNull(p, n);
}

static void* realReallocate(void* oldP, size_t n) {
    if(n <= 0) {
        free(oldP);
        return nullptr;
    }
    void* p = realloc(oldP, n);
    if(p == nullptr) {
        while(p == nullptr && outOfMemoryHandler != nullptr) {
            outOfMemoryHandler(outOfMemoryData);
            p = realloc(oldP, n);
        }
    }
    return exitOnNull(p, n);
}

static void realFree(void* p) {
    free(p);
}

#ifdef CHECK_MEMORY
static const u8 CANARY[16] = {0xDE, 0xAD, 0xBE, 0xEF, 0xDE, 0xAD, 0xBE, 0xEF,
                              0xDE, 0xAD, 0xBE, 0xEF, 0xDE, 0xAD, 0xBE, 0xEF};
static Core::Mutex memoryInfoMutex;

struct MemoryInfo {
    MemoryInfo* next;
    MemoryInfo* previous;
    size_t size;
    int line;
    char buffer[64 - 2 * sizeof(void*) - sizeof(size_t) - sizeof(int)];
    char canary[sizeof(CANARY)];
};

static_assert(sizeof(MemoryInfo) == 80, "memory info has invalid size");
static MemoryInfo* headMemoryInfo = nullptr;

static void addMemoryInfo(MemoryInfo* i, const char* file, int line, size_t n) {
    i->next = nullptr;
    i->previous = nullptr;
    i->size = n;
    i->line = line;
    snprintf(i->buffer, sizeof(i->buffer), "%s", Core::getShortFileName(file));
    memcpy(i->canary, CANARY, sizeof(CANARY));
    memcpy(reinterpret_cast<char*>(i) + n, CANARY, sizeof(CANARY));
    Core::MutexGuard mg(memoryInfoMutex);
    if(headMemoryInfo == nullptr) {
        headMemoryInfo = i;
    } else {
        headMemoryInfo->previous = i;
        i->next = headMemoryInfo;
        headMemoryInfo = i;
    }
}

static void removeMemoryInfo(MemoryInfo* info) {
    Core::MutexGuard mg(memoryInfoMutex);
    if(info->previous == nullptr) {
        if(info->next == nullptr) {
            headMemoryInfo = nullptr;
        } else {
            headMemoryInfo = info->next;
            info->next->previous = nullptr;
        }
    } else {
        if(info->next == nullptr) {
            info->previous->next = nullptr;
        } else {
            info->previous->next = info->next;
            info->next->previous = info->previous;
        }
    }
}

void* Core::debugAllocateRaw(const char* file, int line, size_t n) {
    n += sizeof(MemoryInfo);
    void* p = realAllocate(n + sizeof(CANARY));
    addMemoryInfo(static_cast<MemoryInfo*>(p), file, line, n);
    return static_cast<char*>(p) + sizeof(MemoryInfo);
}

void* Core::debugZeroAllocateRaw(const char* file, int line, size_t n) {
    void* p = debugAllocateRaw(file, line, n);
    memset(p, 0, n);
    return p;
}

void* Core::debugReallocateRaw(const char* file, int line, void* p, size_t n) {
    if(n > 0) {
        n += sizeof(MemoryInfo) + sizeof(CANARY);
    }
    void* rp = p;
    if(rp != nullptr) {
        rp = static_cast<char*>(rp) - sizeof(MemoryInfo);
        removeMemoryInfo(static_cast<MemoryInfo*>(rp));
    }
    void* np = realReallocate(rp, n);
    if(np == nullptr) {
        return nullptr;
    }
    addMemoryInfo(static_cast<MemoryInfo*>(np), file, line, n - sizeof(CANARY));
    return static_cast<char*>(np) + sizeof(MemoryInfo);
}

static bool checkCanary(void* p) {
    return memcmp(p, CANARY, sizeof(CANARY)) != 0;
}

void Core::debugDeallocateRaw(void* p) {
    if(p == nullptr) {
        return;
    }
    void* w = static_cast<char*>(p) - sizeof(MemoryInfo);
    MemoryInfo* rp = static_cast<MemoryInfo*>(w);
    rp->buffer[sizeof(rp->buffer) - 1] = '\0'; // end might be broken
    if(checkCanary(rp->canary)) {
        LOG_ERROR("Free at #:# violated pre canary", rp->buffer, rp->line);
        EXIT(1);
    } else if(checkCanary(reinterpret_cast<char*>(rp) + rp->size)) {
        LOG_ERROR("Free at #:# violated post canary", rp->buffer, rp->line);
        EXIT(1);
    }
    removeMemoryInfo(rp);
    realFree(rp);
}

void Core::printMemoryReport() {
    Core::MutexGuard mg(memoryInfoMutex);
    size_t counter = 0;
    for(MemoryInfo* i = headMemoryInfo; i != nullptr; i = i->next) {
        if(i->line < 0) {
            counter++;
        } else {
            LOG_ERROR("#:# was not freed", i->buffer, i->line);
        }
    }
    if(counter > 0) {
        LOG_ERROR("# unknown entries were not freed", counter);
    }
}

void* operator new(size_t count) {
    return Core::debugAllocateRaw("unknown", -1, count);
}

void* operator new(size_t count, const char* file, int line) {
    return Core::debugAllocateRaw(file, line, count);
}

void* operator new[](size_t count) {
    return Core::debugAllocateRaw("unknown", -1, count);
}

void* operator new[](size_t count, const char* file, int line) {
    return Core::debugAllocateRaw(file, line, count);
}

void operator delete(void* p) noexcept {
    Core::debugDeallocateRaw(p);
}

void operator delete(void* p, size_t) noexcept {
    Core::debugDeallocateRaw(p);
}

void operator delete[](void* p) noexcept {
    Core::debugDeallocateRaw(p);
}

void operator delete[](void* p, size_t) noexcept {
    Core::debugDeallocateRaw(p);
}

#else

void* Core::allocateRaw(size_t n) {
    return realAllocate(n);
}

void* Core::zeroAllocateRaw(size_t n) {
    void* p = allocateRaw(n);
    memset(p, 0, n);
    return p;
}

void* Core::reallocateRaw(void* p, size_t n) {
    return realReallocate(p, n);
}

void Core::deallocateRaw(void* p) {
    realFree(p);
}

#endif