#include "mempool.h"



static inline void* page_payload(mp_page_t *pg){
    return (void*)((char*)pg + sizeof(mp_page_t));
}

static inline void* large_page_payload(mp_large_t *large){
    return (void*)((char*)large + sizeof(mp_large_t));
}

static inline mp_page_t* ptr_to_page(void* p){
    return (mp_page_t*)((uintptr_t)p & ~(MEMPOOL_PAGE_SIZE - 1));
}

static inline int size_to_index(size_t size){
    if(size == 0 || size > MEMPOOL_BLOCK_MAX_SIZE) return -1;

    size_t aligned = (size + 7) & ~7;  // 向上8字节对齐
    return (aligned >> 3) - 1;
}


// bitmap 操作函数
static inline void bitmap_set(uint64_t *bitmap, uint16_t index){
    bitmap[index / 64] |= (1ULL << (index % 64));
}

static inline void bitmap_clear(uint64_t *bitmap, uint16_t index){
    bitmap[index / 64] &= ~(1ULL << (index % 64));
}

static inline int bitmap_test(uint64_t *bitmap, uint16_t index){
    return (bitmap[index / 64] & (1ULL << (index % 64))) != 0;
}

static inline void bitmap_clear_all(uint64_t *bitmap, size_t size){
    memset(bitmap, 0, size * sizeof(uint64_t));
}

// 根据指针计算在页中的块索引
static inline uint16_t ptr_to_block_index(mp_page_t *pg, void *ptr){
    char *base = (char*)page_payload(pg);
    char *p = (char*)ptr;
    size_t offset = p - base;
    return (uint16_t)(offset / pg->owner->block_size);
}


static mp_page_t* mp_page_create(mp_bucket_t *owner){
    if(!owner) return NULL;

    void *mem = NULL;
    mem = aligned_alloc(MEMPOOL_PAGE_SIZE, MEMPOOL_PAGE_SIZE);
    if(!mem) return NULL;

    mp_page_t *pg = (mp_page_t*) mem;

    size_t usable = MEMPOOL_PAGE_SIZE - sizeof(mp_page_t);
    uint16_t cap = (uint16_t)(usable / owner->block_size);
    if(cap == 0) {
        free(mem);
        return NULL;
    }

    pg->owner = owner;
    pg->capacity = cap;
    pg->free_count = cap;
    pg->prev = NULL;
    pg->next = NULL;

    bitmap_clear_all(pg->bitmap, 16);

    char *p = (char*)page_payload(pg);
    for(uint16_t i = 0;i < cap - 1; ++ i){
        *(void**)p = (void*)(p + owner->block_size);
        p = p + owner->block_size;
    }
    *(void**)p = NULL;
    pg->free_list = page_payload(pg);

    return pg;
}

static void *mp_page_alloc(mp_page_t *pg){
    if(!pg || !pg->free_list || pg->free_count == 0) return NULL;

    void *ret = pg->free_list;
    pg->free_list = *(void**)ret;
    pg->free_count --;

    // 标记该块为已分配
    uint16_t index = ptr_to_block_index(pg, ret);
    bitmap_set(pg->bitmap, index);

    return ret;
}

static int mp_page_free(mp_page_t *pg, void *ptr){
    if(!pg || !ptr) return MEMPOOL_INVALID_INPUT;

    // 检查是否是 double free
    uint16_t index = ptr_to_block_index(pg, ptr);
    if(!bitmap_test(pg->bitmap, index)){
        // 该块未被分配，可能是 double free
        return MEMPOOL_DOUBLE_FREE;
    }

    // 标记该块为空闲
    bitmap_clear(pg->bitmap, index);

    *(void**)ptr    = pg->free_list;
    pg->free_list   = ptr;
    pg->free_count ++;

    return 0;
}


int mp_create(mp_pool_t *pool){
    if(!pool) return MEMPOOL_INVALID_INPUT;
    for(int i = 0;i < MEMPOOL_NUM_CLASSES; ++ i){
        mp_bucket_t *bucket = &pool->buckets[i];
        bucket->block_size = (i+1)*MEMPOOL_ALIGNMENT;
        bucket->page_count = 0;
        bucket->empty_count = 0;
        bucket->empty_pages = bucket->partial_pages = bucket->full_pages = NULL;
    }
    pool->large_list = NULL;
    return 0;
}

int mp_destroy(mp_pool_t *pool){
    if(!pool) return MEMPOOL_INVALID_INPUT;

    // 释放所有大块内存
    mp_large_t *large = pool->large_list;
    while(large){
        mp_large_t *next = large->next;
        free(large);
        large = next;
    }
    
    // 释放所有页面
    for(int i = 0; i < MEMPOOL_NUM_CLASSES; i++){
        mp_bucket_t *bucket = &pool->buckets[i];
        mp_page_t *lists[] = {bucket->empty_pages, bucket->partial_pages, bucket->full_pages};
        for(int j = 0; j < 3; j++){
            mp_page_t *pg = lists[j];
            while(pg){
                mp_page_t *next = pg->next;
                free(pg);
                pg = next;
            }
        }
    }

    return 0;
}

void *mp_alloc(mp_pool_t *pool, size_t size){
    if(!pool) return NULL;
    if(size > MEMPOOL_BLOCK_MAX_SIZE){
        // large_page
        void *mem = malloc(size + sizeof(mp_large_t));
        if(!mem) return NULL;

        mp_large_t *large_page = (mp_large_t*) mem;

        large_page->prev = NULL;
        large_page->next = pool->large_list;
        if (pool->large_list) pool->large_list->prev = large_page;
        pool->large_list = large_page;

        return large_page_payload(large_page);
    }

    int i = size_to_index(size);
    if(i == -1) return NULL;

    mp_bucket_t *bucket = &pool->buckets[i];

    if(!bucket->partial_pages){
        if(!bucket->empty_pages){
            // 无空页，重新分配
            mp_page_t *pg = mp_page_create(bucket);    
            if(!pg) return NULL;
            bucket->partial_pages = pg;
            bucket->page_count ++;
        }else{
            // 有空页，使用空页
            mp_page_t *epage = bucket->empty_pages;
            bucket->empty_pages = epage->next;
            if(bucket->empty_pages) bucket->empty_pages->prev = NULL;
            epage->next = NULL;
            epage->prev = NULL;

            bucket->partial_pages = epage;
            bucket->empty_count --;
        }
    }

    mp_page_t   *pg     = bucket->partial_pages;    
    void        *p      = mp_page_alloc(pg);
    if(!p) return NULL;

    if(pg->free_count == 0){
        // 该页没有空位
        // 从partial链表中取出
        if(pg == bucket->partial_pages){
            bucket->partial_pages = pg->next;
        }
        if(pg->prev != NULL) pg->prev->next = pg->next;
        if(pg->next != NULL) pg->next->prev = pg->prev;
        pg->prev = NULL;
        pg->next = NULL;

        // 移入 fullpage
        if(!bucket->full_pages) {
            bucket->full_pages = pg;
        }else{
            pg->prev = NULL;
            pg->next = bucket->full_pages;
            if (bucket->full_pages) bucket->full_pages->prev = pg;
            bucket->full_pages = pg;
        }
    }else {
        // 不做处理

    }

    return p; 
}

int mp_free(mp_pool_t *pool, void *ptr){
    if(!pool || !ptr) return MEMPOOL_INVALID_INPUT;
    mp_large_t *large = pool->large_list;
    while(large){
        if(ptr == large_page_payload(large)){
            // 从large链表中删除
            if(large == pool->large_list) pool->large_list = large->next;
            if(large->prev) large->prev->next = large->next;
            if(large->next) large->next->prev = large->prev;
            large->prev = NULL;
            large->next = NULL;
            
            free(large);
            return 0;
        }
        large = large->next;
    }

    mp_page_t *pg = ptr_to_page(ptr);

    uint16_t before = pg->free_count;
    int ret = mp_page_free(pg, ptr);
    if(ret != 0){
        printf("mempool error\n");
        return ret;
    }

    if(pg->free_count >= pg->capacity){
        // free后成为完全空闲页
        mp_bucket_t *bucket = pg->owner;
        if(bucket->empty_count > MEMPOOL_CACHE_PAGE){
            // 空闲页很多
            // 摘出partial链表
            if(pg == bucket->partial_pages) bucket->partial_pages = pg->next;
            if(pg->prev != NULL) pg->prev->next = pg->next;
            if(pg->next != NULL) pg->next->prev = pg->prev;
            pg->prev = NULL;
            pg->next = NULL;
            // 释放
            free(pg);
            bucket->page_count --;
        }else{
            // 空闲页并不多
            // 摘出partial链表
            if(pg == bucket->partial_pages) bucket->partial_pages = pg->next;
            if(pg->prev != NULL) pg->prev->next = pg->next;
            if(pg->next != NULL) pg->next->prev = pg->prev;
            pg->prev = NULL;
            pg->next = NULL;
            // 加入empty链表
            pg->next = bucket->empty_pages;
            if(bucket->empty_pages) bucket->empty_pages->prev = pg;
            bucket->empty_pages = pg;
            bucket->empty_count ++;
        }
        
    }else if(before == 0 && pg->free_count == 1){
        // free 后需要从full移动到partial
        mp_bucket_t *bucket = pg->owner;
        // 移出full链表
        if(pg == bucket->full_pages) bucket->full_pages = pg->next;
        if(pg->prev != NULL) pg->prev->next = pg->next;
        if(pg->next != NULL) pg->next->prev = pg->prev;
        pg->prev = NULL;
        pg->next = NULL;

        // 加入partial
        pg->prev = NULL;
        pg->next = bucket->partial_pages;
        if (bucket->partial_pages) bucket->partial_pages->prev = pg;
        bucket->partial_pages = pg;
    }else {
        // 有空余空间的页不做处理
    }

    return 0;
}

int mp_print(mp_pool_t *pool){
    int ret = 0;
    printf("------\n");
    for(int i = 0; i < MEMPOOL_NUM_CLASSES; i++){
        mp_bucket_t *bucket = &pool->buckets[i];
        if(bucket->page_count) ret += bucket->page_count; 
        printf("size:%ld, page:%d, empty:%d\n", bucket->block_size, bucket->page_count, bucket->empty_count);
    }
    printf("------\n");

    printf("page count: %d\n", ret);
    return ret;
}