#include "mem_pool.h"
#include <stddef.h>
#include <stdlib.h>
#include <stdio.h>
#include <jemalloc/jemalloc.h>


void *kvs_malloc(size_t size) {
#if MEMORY_SELECT_MALLOC == MEMORY_USE_DEFAULT
	return malloc(size);
#elif MEMORY_SELECT_MALLOC == MEMORY_USE_MYMALLOC
    return mp_alloc(size);
#elif MEMORY_SELECT_MALLOC == MEMORY_USE_JEMALLOC
    return je_malloc(size);
#endif
}

void kvs_free(void *ptr) {
#if MEMORY_SELECT_MALLOC == MEMORY_USE_DEFAULT
	free(ptr);
#elif MEMORY_SELECT_MALLOC == MEMORY_USE_MYMALLOC
    mp_free(ptr);
#elif MEMORY_SELECT_MALLOC == MEMORY_USE_JEMALLOC
    je_free(ptr);
#endif
}

#if MEMORY_SELECT_MALLOC == MEMORY_USE_MYMALLOC
mp_pool_t global_mempool = {0};

// >= x 的第一个2^n
uint32_t mp_ceil_pow2_u32(uint32_t x) {
    if (x <= 1) return 1;
    x--;
    x |= x >> 1;
    x |= x >> 2;
    x |= x >> 4;
    x |= x >> 8;
    x |= x >> 16;
    return x + 1;
}

// return cid 
uint16_t mp_class_id_for(size_t size) {
    if (size < 8) size = 8;
    if (size > (1u << MP_MAX_SHIFT)) return MP_LARGE_CLASS;

    uint32_t s = (uint32_t)size;
    uint32_t p2 = mp_ceil_pow2_u32(s);  /* 8/16/32... */

    /* 计算 log2(p2) */
    uint16_t shift = 0;
    while ((1u << shift) < p2) shift++;

    return (uint16_t)(shift - MP_MIN_SHIFT);
}

size_t mp_user_size_from_class(uint16_t cid) {
    return (size_t)1u << (cid + MP_MIN_SHIFT);
}

int mp_add_page(mp_pool_t* p, void* mem) {
    mp_page_t* pg = (mp_page_t*)je_malloc(sizeof(mp_page_t));
    if (!pg) return -1;
    pg->mem = mem;
    pg->next = p->pages;
    p->pages = pg;
    return 0;
}

int mp_grow(mp_pool_t* p, uint16_t cid) {
    void* page = je_malloc(MP_PAGE_SIZE);
    if (!page) return -1;
    if (mp_add_page(p, page) != 0) { je_free(page); return -1; }

    // cid对应的空间大小
    const size_t user_sz  = mp_user_size_from_class(cid);
    const size_t block_sz = sizeof(mp_blk_hdr_t) + user_sz;
    size_t n = MP_PAGE_SIZE / block_sz;
    if (n == 0) { je_free(page); return -1; }

    uint8_t* cur = (uint8_t*)page;

    mp_blk_hdr_t* first_hdr = (mp_blk_hdr_t*)cur;
    first_hdr->magic = MP_MAGIC;
    first_hdr->class_id = cid;

    mp_free_node_t* first_node = (mp_free_node_t*)(first_hdr + 1);
    mp_free_node_t* head = first_node;
    mp_free_node_t* tail = first_node;

    cur += block_sz;
    n--;

    for (size_t i = 0; i < n; ++i) {
        mp_blk_hdr_t* h = (mp_blk_hdr_t*)cur;
        h->magic = MP_MAGIC;
        h->class_id = cid;

        mp_free_node_t* node = (mp_free_node_t*)(h + 1);
        tail->next = node;
        tail = node;

        cur += block_sz;
    }

    tail->next = p->buckets[cid].free_list;
    p->buckets[cid].free_list = head;
    return 0;
}

void  mp_init(mp_pool_t* p){
    if (!p) return;
    for (int i = 0; i < MP_NBUCKETS; ++i) p->buckets[i].free_list = NULL;
    p->pages = NULL;
}

void  mp_destroy(mp_pool_t* p){
    if (!p) return;
    mp_page_t* pg = p->pages;
    while (pg) {
        mp_page_t* nxt = pg->next;
        je_free(pg->mem);
        je_free(pg);
        pg = nxt;
    }
    p->pages = NULL;
    for (int i = 0; i < MP_NBUCKETS; ++i) p->buckets[i].free_list = NULL;
}

void* mp_alloc(size_t size){
    mp_pool_t* p = &global_mempool;
    if (!p) return NULL;
    if (size == 0) size = 1;

    uint16_t cid = mp_class_id_for(size);
    if (cid == MP_LARGE_CLASS) {
        mp_blk_hdr_t* h = (mp_blk_hdr_t*)je_malloc(sizeof(mp_blk_hdr_t) + size);
        if (!h) return NULL;
        h->magic = MP_MAGIC;
        h->class_id = MP_LARGE_CLASS;
        return (void*)(h + 1);    
    }

    mp_free_node_t* node = p->buckets[cid].free_list;
    if (!node) {
        if (mp_grow(p, cid) != 0) return NULL;
        node = p->buckets[cid].free_list;
        if (!node) return NULL;
    }

    p->buckets[cid].free_list = node->next;
    return (void*)node;
}

void  mp_free(void* ptr){
    mp_pool_t* p = &global_mempool;
    if (!p || !ptr) return;

    mp_blk_hdr_t* h = ((mp_blk_hdr_t*)ptr) - 1;
    if (h->magic != MP_MAGIC) {
        /* 非法指针：极简处理（你也可以 assert/abort） */
        return;
    }

    if (h->class_id == MP_LARGE_CLASS) {
        /* 大对象：直接 free 整块（从 header 起） */
        je_free(h);
        return;
    }

    if (!p) return;
    uint16_t cid = h->class_id;
    if (cid >= MP_NBUCKETS) return; /* 极简：不做更多校验 */

    mp_free_node_t* node = (mp_free_node_t*)ptr;
    node->next = p->buckets[cid].free_list;
    p->buckets[cid].free_list = node;
}
#endif