#include "alloc_dispatch.h"
#include "mempool.h"
#include <stddef.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <stdatomic.h>
#include <jemalloc/jemalloc.h>

mp_pool_t global_mempool;

static atomic_int g_memory_mode         = ATOMIC_VAR_INIT(MEMLEAK_DETECT_OFF);
static atomic_int g_memleak_detect_mode = ATOMIC_VAR_INIT(MEMLEAK_DETECT_OFF);

// 设置内存池类型
void kvs_set_alloc_type(AllocatorType mode) {
    atomic_store(&g_memory_mode, mode);
}

// 获取当前内存池类型
AllocatorType kvs_get_alloc_type(void) {
    return atomic_load(&g_memory_mode);
}

// 设置内存泄漏检测模式
void kvs_set_memleak_detect(MemLeakDetectMode mode) {
    atomic_store(&g_memleak_detect_mode, mode);
}

// 获取当前内存泄漏检测模式
MemLeakDetectMode kvs_get_memleak_detect(void) {
    return atomic_load(&g_memleak_detect_mode);
}



void *kvs_malloc_impl(size_t size){

    switch (atomic_load(&g_memleak_detect_mode)){
        case ALLOC_MALLOC:
            return malloc(size);
        case ALLOC_MYPOOL:
            return mp_alloc(&global_mempool, size);
        case ALLOC_JEMALLOC:
            return je_malloc(size);
        case ALLOC_OTHER:
        default:
            return malloc(size);        
    }
    return NULL;
}

void kvs_free_impl(void *ptr) {
    switch (atomic_load(&g_memleak_detect_mode)){
        case ALLOC_MALLOC:
	        free(ptr);
            break;
        case ALLOC_MYPOOL:{
            int ret = mp_free(&global_mempool, ptr);
            if(ret == MEMPOOL_DOUBLE_FREE){
                printf("double free %p\n", ptr);
            }
            break;
        }
        case ALLOC_JEMALLOC:
            je_free(ptr);
            break;
        case ALLOC_OTHER:
        default:
            free(ptr);
            break;
    }
}

void *nMalloc(size_t size, const char * filename, const char *func, int line){
    void *ptr = kvs_malloc_impl(size);

    if(atomic_load(&g_memleak_detect_mode) == MEMLEAK_DETECT_ON) {
        char buff[128];
        snprintf(buff, 128, "./mem_leak/%p.mem", ptr);
        FILE* fp = fopen(buff, "w");
        if(!fp){
            kvs_free(ptr);
            return NULL;
        }

        fprintf(fp, "[+] [%s : %s : %d] %p: %ld malloc\n", filename, func, line, ptr, size);
        fflush(fp);
        fclose(fp);

    }
    return ptr;
}

void nFree(void *ptr, const char * filename, const char *func, int line){
    if(!ptr) {
        return ;
    }
    
    if(atomic_load(&g_memleak_detect_mode) == MEMLEAK_DETECT_ON) {    
        char buff[128];
        snprintf(buff, 128, "./mem_leak/%p.mem", ptr);

        if(unlink(buff) < 0) {
            return ;
        }

    }

    kvs_free_impl(ptr);
}