#include "hakmem_super_registry.h"
#include "hakmem_tiny_superslab.h"
#include <string.h>
#include <stdio.h>

// Global registry storage
SuperRegEntry g_super_reg[SUPER_REG_SIZE];
pthread_mutex_t g_super_reg_lock = PTHREAD_MUTEX_INITIALIZER;
int g_super_reg_initialized = 0;

// Initialize registry (call once at startup)
void hak_super_registry_init(void) {
    if (g_super_reg_initialized) return;

    // Zero-initialize all entries
    memset(g_super_reg, 0, sizeof(g_super_reg));

    // Memory fence to ensure initialization is visible to all threads
    atomic_thread_fence(memory_order_release);

    g_super_reg_initialized = 1;
}

// Register SuperSlab (mutex-protected)
// CRITICAL: Call AFTER SuperSlab is fully initialized
// Publish order: ss init → release fence → base write
// Phase 8.3: ACE - lg_size aware registration
int hak_super_register(uintptr_t base, SuperSlab* ss) {
    if (!g_super_reg_initialized) {
        hak_super_registry_init();
    }

    pthread_mutex_lock(&g_super_reg_lock);

    int lg = ss->lg_size;  // Phase 8.3: Get lg_size from SuperSlab
    int h = hak_super_hash(base, lg);

    // Linear probing to find empty slot
    for (int i = 0; i < SUPER_MAX_PROBE; i++) {
        SuperRegEntry* e = &g_super_reg[(h + i) & SUPER_REG_MASK];

        if (e->base == 0) {
            // Found empty slot
            // Step 1: Write SuperSlab pointer and lg_size (atomic for MT-safety)
            atomic_store_explicit(&e->ss, ss, memory_order_release);
            e->lg_size = lg;  // Phase 8.3: Store lg_size for fast lookup

            // Step 2: Release fence (ensures ss/lg_size write is visible before base)
            atomic_thread_fence(memory_order_release);

            // Step 3: Publish base address (makes entry visible to readers)
            atomic_store_explicit((_Atomic uintptr_t*)&e->base, base,
                                 memory_order_release);

            pthread_mutex_unlock(&g_super_reg_lock);
            return 1;
        }

        if (e->base == base && e->lg_size == lg) {
            // Already registered (duplicate registration)
            pthread_mutex_unlock(&g_super_reg_lock);
            return 1;
        }
    }

    // Registry full (probing limit reached)
    pthread_mutex_unlock(&g_super_reg_lock);
    fprintf(stderr, "HAKMEM: SuperSlab registry full! Increase SUPER_REG_SIZE\n");
    return 0;
}

// Unregister SuperSlab (mutex-protected)
// CRITICAL: Call BEFORE munmap to prevent reader segfault
// Unpublish order: base = 0 (release) → munmap outside this function
// Phase 8.3: ACE - Try both lg_sizes (we don't know which one was used)
void hak_super_unregister(uintptr_t base) {
    if (!g_super_reg_initialized) return;

    pthread_mutex_lock(&g_super_reg_lock);

    // Try both 1MB (20) and 2MB (21) alignments
    for (int lg = 20; lg <= 21; lg++) {
        int h = hak_super_hash(base, lg);

        // Linear probing to find matching entry
        for (int i = 0; i < SUPER_MAX_PROBE; i++) {
            SuperRegEntry* e = &g_super_reg[(h + i) & SUPER_REG_MASK];

            if (e->base == base && e->lg_size == lg) {
                // Found entry to remove
                // Step 1: Clear SuperSlab pointer (atomic, prevents TOCTOU race)
                atomic_store_explicit(&e->ss, NULL, memory_order_release);

                // Step 2: Unpublish base (makes entry invisible to readers)
                atomic_store_explicit((_Atomic uintptr_t*)&e->base, 0,
                                     memory_order_release);

                // Step 3: Clear lg_size (optional cleanup)
                e->lg_size = 0;

                pthread_mutex_unlock(&g_super_reg_lock);
                return;
            }

            if (e->base == 0) {
                // Not found in this lg_size, try next
                break;
            }
        }
    }

    pthread_mutex_unlock(&g_super_reg_lock);
    // Not found is not an error (could be duplicate unregister)
}

// Debug: Get registry statistics
void hak_super_registry_stats(SuperRegStats* stats) {
    if (!stats) return;

    stats->total_slots = SUPER_REG_SIZE;
    stats->used_slots = 0;
    stats->max_probe_depth = 0;

    pthread_mutex_lock(&g_super_reg_lock);

    // Count used slots
    for (int i = 0; i < SUPER_REG_SIZE; i++) {
        if (g_super_reg[i].base != 0) {
            stats->used_slots++;
        }
    }

    // Calculate max probe depth
    for (int i = 0; i < SUPER_REG_SIZE; i++) {
        if (g_super_reg[i].base != 0) {
            uintptr_t base = g_super_reg[i].base;
            int lg = g_super_reg[i].lg_size;  // Phase 8.3: Use stored lg_size
            int h = hak_super_hash(base, lg);

            // Find actual probe depth for this entry
            for (int j = 0; j < SUPER_MAX_PROBE; j++) {
                int idx = (h + j) & SUPER_REG_MASK;
                if (g_super_reg[idx].base == base && g_super_reg[idx].lg_size == lg) {
                    if (j > stats->max_probe_depth) {
                        stats->max_probe_depth = j;
                    }
                    break;
                }
            }
        }
    }

    pthread_mutex_unlock(&g_super_reg_lock);
}