module Core.Utility;

import Core.Logger;
import Core.Thread;
import Core.Types;
import Core.Std;
import Core.New;
import Core.ToString;

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

[[noreturn]] void Core::exitWithHandler(
    int value, const std::source_location& l) noexcept {
    if(value != 0) {
        const char* file = getShortFileName(l.file_name());
        logError("Exit from {}:{} with value {}", file, l.line(), value);
    }
    if(exitHandler != nullptr) {
        exitHandler(value, exitData);
    }
    exit(value);
}

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

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

static void* exitOnNull(void* p, size_t n) noexcept {
    if(p == nullptr) {
        Core::logError("Out of memory, requested '{}' bytes", n);
        Core::exitWithHandler(1);
    }
    return p;
}

static void* realAllocate(size_t n) noexcept {
    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) noexcept {
    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) noexcept {
    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;
    u32 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 std::source_location& l, size_t n) noexcept {
    i->next = nullptr;
    i->previous = nullptr;
    i->size = n;
    i->line = l.line();
    Core::StringBase(i->buffer, sizeof(i->buffer))
        .add(Core::getShortFileName(l.file_name()));
    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) noexcept {
    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::allocateRaw(size_t n, const std::source_location& l) noexcept {
    n += sizeof(MemoryInfo);
    void* p = realAllocate(n + sizeof(CANARY));
    addMemoryInfo(static_cast<MemoryInfo*>(p), l, n);
    return static_cast<char*>(p) + sizeof(MemoryInfo);
}

void* Core::zeroAllocateRaw(size_t n, const std::source_location& l) noexcept {
    void* p = allocateRaw(n, l);
    memset(p, 0, n);
    return p;
}

void* Core::reallocateRaw(
    void* p, size_t n, const std::source_location& l) noexcept {
    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), l, n - sizeof(CANARY));
    return static_cast<char*>(np) + sizeof(MemoryInfo);
}

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

void Core::deallocateRaw(void* p, const std::source_location& l) noexcept {
    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)) {
        logError(
            "Free at {}:{} from {}:{} violated pre canary",
            getShortFileName(l.file_name()), l.line(), rp->buffer, rp->line);
        exitWithHandler(1, l);
    } else if(checkCanary(reinterpret_cast<char*>(rp) + rp->size)) {
        logError(
            "Free at {}:{} from {}:{} violated post canary",
            getShortFileName(l.file_name()), l.line(), rp->buffer, rp->line);
        exitWithHandler(1, l);
    }
    removeMemoryInfo(rp);
    realFree(rp);
}

void* operator new(size_t count, const std::source_location& l) noexcept {
    return Core::allocateRaw(count, l);
}

void* operator new[](size_t count, const std::source_location& l) noexcept {
    return Core::allocateRaw(count, l);
}

#else

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

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

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

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

#endif

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