#include "pool_refill.h"
#include "pool_tls.h"
#include "pool_tls_arena.h"
#include "pool_tls_remote.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;

    // Refill boundary integration: try draining remote frees first.
    // If we can satisfy from remote queue, avoid backend carve (OS pressure).
    {
        void* rchain = NULL;
        int rgot = pool_remote_pop_chain(class_idx, count, &rchain);
        if (rgot > 0 && rchain) {
            // Pop one to return, install the rest into TLS
            void* ret = rchain;
            rchain = *(void**)rchain;
            rgot--;
            if (rgot > 0 && rchain) {
                pool_install_chain(class_idx, rchain, rgot);
            }
            #if POOL_USE_HEADERS
            *((uint8_t*)ret - POOL_HEADER_SIZE) = POOL_MAGIC | class_idx;
            #endif
            return ret;
        }
    }

    // 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.5a: TLS Arena with chunk carving
void* backend_batch_carve(int class_idx, int count) {
    if (class_idx < 0 || class_idx >= POOL_SIZE_CLASSES || count <= 0) {
        return NULL;
    }

    // Allocate blocks array on stack
    void* blocks[64];  // Max refill count is 64
    if (count > 64) count = 64;

    // Carve from TLS Arena (Phase 1.5a)
    int carved = arena_batch_carve(class_idx, blocks, count);
    if (carved == 0) {
        return NULL;  // OOM
    }

    // Chain the carved blocks
    void* head = NULL;
    void* tail = NULL;

    for (int i = 0; i < carved; i++) {
        void* block = blocks[i];

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

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

    return head;
}