#include "pool_refill.h"
#include "pool_tls.h"
#include <sys/mman.h>
#include <stdint.h>
#include <errno.h>

// Get refill count from Box 1
extern int pool_get_refill_count(int class_idx);

// Refill and return first block
void* pool_refill_and_alloc(int class_idx) {
    int count = pool_get_refill_count(class_idx);
    if (count <= 0) return NULL;

    // Batch allocate from existing Pool backend
    void* chain = backend_batch_carve(class_idx, count);
    if (!chain) return NULL;  // OOM

    // Pop first block for return
    void* ret = chain;
    chain = *(void**)chain;
    count--;

    #if POOL_USE_HEADERS
    // Write header for the block we're returning
    *((uint8_t*)ret - POOL_HEADER_SIZE) = POOL_MAGIC | class_idx;
    #endif

    // Install rest in TLS (if any)
    if (count > 0 && chain) {
        pool_install_chain(class_idx, chain, count);
    }

    return ret;
}

// Backend batch carve - Phase 1: Direct mmap allocation
void* backend_batch_carve(int class_idx, int count) {
    if (class_idx < 0 || class_idx >= POOL_SIZE_CLASSES || count <= 0) {
        return NULL;
    }

    // Get the class size
    size_t block_size = POOL_CLASS_SIZES[class_idx];

    // For Phase 1: Allocate a single large chunk via mmap
    // and carve it into blocks
    #if POOL_USE_HEADERS
    size_t total_block_size = block_size + POOL_HEADER_SIZE;
    #else
    size_t total_block_size = block_size;
    #endif

    // Allocate enough for all requested blocks
    size_t total_size = total_block_size * count;

    // Round up to page size
    size_t page_size = 4096;
    total_size = (total_size + page_size - 1) & ~(page_size - 1);

    // Allocate memory via mmap
    void* chunk = mmap(NULL, total_size, PROT_READ | PROT_WRITE,
                       MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (chunk == MAP_FAILED) {
        return NULL;
    }

    // Carve into blocks and chain them
    void* head = NULL;
    void* tail = NULL;
    char* ptr = (char*)chunk;

    for (int i = 0; i < count; i++) {
        #if POOL_USE_HEADERS
        // Skip header space - user data starts after header
        void* user_ptr = ptr + POOL_HEADER_SIZE;
        #else
        void* user_ptr = ptr;
        #endif

        // Chain the blocks
        if (!head) {
            head = user_ptr;
            tail = user_ptr;
        } else {
            *(void**)tail = user_ptr;
            tail = user_ptr;
        }

        // Move to next block
        ptr += total_block_size;

        // Stop if we'd go past the allocated chunk
        if ((ptr + total_block_size) > ((char*)chunk + total_size)) {
            break;
        }
    }

    // Terminate chain
    if (tail) {
        *(void**)tail = NULL;
    }

    return head;
}