#include "core/Utility.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <threads.h>
#include <time.h>

#include "ErrorSimulator.h"
#include "core/Logger.h"

extern inline size_t maxSize(size_t a, size_t b);
extern inline size_t minSize(size_t a, size_t b);

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

size_t popCount(u64 u) {
    static const u64 map[16] = {0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4};
    size_t sum = 0;
    for(size_t i = 0; i < sizeof(u64) * 8; i += 4) {
        sum += map[(u >> i) & 0xF];
    }
    return sum;
}

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

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

void setOutOfMemoryHandler(OutOfMemoryHandler h, void* data) {
    outOfMemoryHandler = h;
    outOfMemoryData = data;
}

static void* exitOnNull(void* p, size_t n) {
    if(p == nullptr) {
        LOG_ERROR("Out of memory, requested '%zu' 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};

struct MemoryInfo;
typedef struct MemoryInfo MemoryInfo;

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* info, const char* file, int line, size_t n) {
    info->next = nullptr;
    info->previous = nullptr;
    info->size = n;
    info->line = line;
    snprintf(info->buffer, sizeof(info->buffer), "%s", getShortFileName(file));
    memcpy(info->canary, CANARY, sizeof(CANARY));
    memcpy((char*)info + n - sizeof(CANARY), CANARY, sizeof(CANARY));
    if(headMemoryInfo == nullptr) {
        headMemoryInfo = info;
    } else {
        headMemoryInfo->previous = info;
        info->next = headMemoryInfo;
        headMemoryInfo = info;
    }
}

static void removeMemoryInfo(MemoryInfo* info) {
    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* coreDebugAllocate(const char* file, int line, size_t n) {
    n += sizeof(MemoryInfo) + sizeof(CANARY);
    void* p = RealAllocate(n + sizeof(CANARY));
    addMemoryInfo(p, file, line, n);
    return (char*)p + sizeof(MemoryInfo);
}

void* coreDebugZeroAllocate(const char* file, int line, size_t n) {
    void* p = coreDebugAllocate(file, line, n);
    memset(p, 0, n);
    return p;
}

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

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

void coreFreeDebug(const char* file, int line, void* p) {
    if(p == nullptr) {
        return;
    }
    void* w = (char*)p - sizeof(MemoryInfo);
    MemoryInfo* rp = w;
    if(checkCanary(rp->canary)) {
        file = getShortFileName(file);
        LOG_ERROR("Free at %s:%d violated pre canary", file, line);
        EXIT(1);
    } else if(checkCanary((char*)rp + rp->size - sizeof(CANARY))) {
        file = getShortFileName(file);
        LOG_ERROR("Free at %s:%d violated post canary", file, line);
        EXIT(1);
    }
    removeMemoryInfo(rp);
    RealFree(rp);
}

void printMemoryReport() {
    for(MemoryInfo* i = headMemoryInfo; i != nullptr; i = i->next) {
        LOG_ERROR("%s:%d was not freed", i->buffer, i->line);
    }
}

#else

void* coreAllocate(size_t n) {
    return RealAllocate(n);
}

void* coreZeroAllocate(size_t n) {
    void* p = coreAllocate(n);
    memset(p, 0, n);
    return p;
}

void* coreReallocate(void* p, size_t n) {
    return RealReallocate(p, n);
}

void coreFree(void* p) {
    RealFree(p);
}

#endif

bool sleepNanos(i64 nanos) {
    struct timespec t;
    t.tv_nsec = nanos % 1'000'000'000;
    t.tv_sec = nanos / 1'000'000'000;
    return thrd_sleep(&t, nullptr) != 0;
}

i64 getNanos(void) {
    struct timespec ts;
    if(timespec_get(&ts, TIME_UTC) == 0 || TIME_GET_FAIL) {
        return -1;
    }
    return (i64)ts.tv_sec * 1'000'000'000L + (i64)ts.tv_nsec;
}