#include "pool_tls.h"
#include <string.h>
#include <stdint.h>
#include <stdbool.h>

// Class sizes: 8KB, 16KB, 24KB, 32KB, 40KB, 48KB, 52KB
const size_t POOL_CLASS_SIZES[POOL_SIZE_CLASSES] = {
    8192, 16384, 24576, 32768, 40960, 49152, 53248
};

// TLS state (per-thread)
__thread void* g_tls_pool_head[POOL_SIZE_CLASSES];
__thread uint32_t g_tls_pool_count[POOL_SIZE_CLASSES];

// Fixed refill counts (Phase 1: no learning)
static const uint32_t DEFAULT_REFILL_COUNT[POOL_SIZE_CLASSES] = {
    64, 48, 32, 32, 24, 16, 16  // Larger classes = smaller refill
};

// Forward declare refill function (from Box 2)
extern void* pool_refill_and_alloc(int class_idx);

// Size to class mapping
static inline int pool_size_to_class(size_t size) {
    // Binary search would be overkill for 7 classes
    // Simple linear search with early exit
    if (size <= 8192) return 0;
    if (size <= 16384) return 1;
    if (size <= 24576) return 2;
    if (size <= 32768) return 3;
    if (size <= 40960) return 4;
    if (size <= 49152) return 5;
    if (size <= 53248) return 6;
    return -1;  // Too large for Pool
}

// Ultra-fast allocation (5-6 cycles)
void* pool_alloc(size_t size) {
    // Quick bounds check
    if (size < 8192 || size > 53248) return NULL;

    int class_idx = pool_size_to_class(size);
    if (class_idx < 0) return NULL;

    void* head = g_tls_pool_head[class_idx];

    if (__builtin_expect(head != NULL, 1)) {  // LIKELY
        // Pop from freelist (3-4 instructions)
        g_tls_pool_head[class_idx] = *(void**)head;
        g_tls_pool_count[class_idx]--;

        #if POOL_USE_HEADERS
        // Write header (1 byte before ptr)
        *((uint8_t*)head - POOL_HEADER_SIZE) = POOL_MAGIC | class_idx;
        #endif

        return head;
    }

    // Cold path: refill
    return pool_refill_and_alloc(class_idx);
}

// Ultra-fast free (5-6 cycles)
void pool_free(void* ptr) {
    if (!ptr) return;

    #if POOL_USE_HEADERS
    // Read class from header
    uint8_t header = *((uint8_t*)ptr - POOL_HEADER_SIZE);
    if ((header & 0xF0) != POOL_MAGIC) {
        // Not ours, route elsewhere
        return;
    }
    int class_idx = header & 0x0F;
    if (class_idx >= POOL_SIZE_CLASSES) return;  // Invalid class
    #else
    // Need registry lookup (slower fallback) - not implemented in Phase 1
    return;
    #endif

    // Push to freelist (2-3 instructions)
    *(void**)ptr = g_tls_pool_head[class_idx];
    g_tls_pool_head[class_idx] = ptr;
    g_tls_pool_count[class_idx]++;

    // Phase 1: No drain logic (keep it simple)
}

// Install refilled chain (called by Box 2)
void pool_install_chain(int class_idx, void* chain, int count) {
    if (class_idx < 0 || class_idx >= POOL_SIZE_CLASSES) return;
    g_tls_pool_head[class_idx] = chain;
    g_tls_pool_count[class_idx] = count;
}

// Get refill count for a class
int pool_get_refill_count(int class_idx) {
    if (class_idx < 0 || class_idx >= POOL_SIZE_CLASSES) return 0;
    return DEFAULT_REFILL_COUNT[class_idx];
}

// Thread init/cleanup
void pool_thread_init(void) {
    memset(g_tls_pool_head, 0, sizeof(g_tls_pool_head));
    memset(g_tls_pool_count, 0, sizeof(g_tls_pool_count));
}

void pool_thread_cleanup(void) {
    // Phase 1: No cleanup (keep it simple)
    // TODO: Drain back to global pool
}