// smallsegment_v5.c - SmallSegment v5 Implementation (Phase v5-2)
//
// Purpose: 2MiB segment-based page allocation with O(1) page_meta lookup
// Design: Each segment contains 32 pages (64KiB each) with embedded metadata

#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <stdint.h>
#include <stdio.h>
#include "box/smallsegment_v5_box.h"

#ifndef likely
#define likely(x)   __builtin_expect(!!(x), 1)
#define unlikely(x) __builtin_expect(!!(x), 0)
#endif

// ============================================================================
// Segment Allocation (Phase v5-2)
// ============================================================================

// Thread-local segment (Phase v5-3: single segment per thread for O(1) lookup)
// C6-only v5 uses at most 1 segment per thread - this eliminates slot search
typedef struct {
    SmallSegmentV5 seg;
    int in_use;
    uint32_t used_pages;  // Bitmap: which pages are currently in use
} TLSSegmentSlot;

static __thread TLSSegmentSlot g_tls_segment_v5;  // Single TLS segment

SmallSegmentV5* small_segment_v5_acquire(void) {
    // Phase v5-3: Single segment per thread - no slot search needed
    TLSSegmentSlot* slot = &g_tls_segment_v5;

    if (slot->in_use) {
        return NULL;  // Already have a segment (reuse pages instead)
    }

    // Allocate 2MiB aligned segment via mmap
    // Use MAP_ANONYMOUS which typically gives us aligned addresses for large allocations
    void* mem = mmap(NULL, SMALL_SEGMENT_V5_SIZE,
                     PROT_READ | PROT_WRITE,
                     MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);

    if (mem == MAP_FAILED || mem == NULL) {
        return NULL;
    }

    uintptr_t addr = (uintptr_t)mem;

    // Check if we got 2MiB alignment (common for large mmap allocations)
    // If not, remap with extra space to force alignment
    if ((addr & (SMALL_SEGMENT_V5_SIZE - 1)) != 0) {
        // Not aligned, need to reallocate with overallocation
        munmap(mem, SMALL_SEGMENT_V5_SIZE);

        // Allocate 4MiB to ensure we can find a 2MiB aligned region
        size_t alloc_size = SMALL_SEGMENT_V5_SIZE * 2;
        mem = mmap(NULL, alloc_size,
                   PROT_READ | PROT_WRITE,
                   MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);

        if (mem == MAP_FAILED || mem == NULL) {
            return NULL;
        }

        // Find the aligned address within this region
        uintptr_t raw_addr = (uintptr_t)mem;
        addr = (raw_addr + SMALL_SEGMENT_V5_SIZE - 1) & ~(SMALL_SEGMENT_V5_SIZE - 1);

        // Verify the aligned address is within our mapping
        if (addr < raw_addr || addr + SMALL_SEGMENT_V5_SIZE > raw_addr + alloc_size) {
            munmap(mem, alloc_size);
            return NULL;  // Alignment calculation error
        }

        // We keep the whole 4MiB mapping to avoid complex munmap logic
        // This wastes some memory but ensures correctness
    }

    // Debug: Verify address is aligned
    if ((addr & (SMALL_SEGMENT_V5_SIZE - 1)) != 0) {
        fprintf(stderr, "[V5_SEG] ERROR: Address 0x%lx not aligned to 0x%lx\n",
                (unsigned long)addr, (unsigned long)SMALL_SEGMENT_V5_SIZE);
        if (addr != (uintptr_t)mem) {
            munmap(mem, SMALL_SEGMENT_V5_SIZE * 2);
        } else {
            munmap(mem, SMALL_SEGMENT_V5_SIZE);
        }
        return NULL;
    }

    // Use TLS slot for metadata (no malloc needed)
    SmallSegmentV5* seg = &slot->seg;
    slot->in_use = 1;
    slot->used_pages = 0;  // Initially no pages are allocated

    // Initialize segment metadata
    seg->base = addr;
    seg->num_pages = SMALL_SEGMENT_V5_NUM_PAGES;
    seg->owner_tid = 0;  // Will be set by caller if needed
    seg->magic = SMALL_SEGMENT_V5_MAGIC;

    // Initialize all page metadata
    for (uint32_t i = 0; i < seg->num_pages; i++) {
        SmallPageMetaV5* m = &seg->page_meta[i];
        m->free_list = NULL;
        m->used = 0;
        m->capacity = 0;
        m->class_idx = 0;
        m->flags = 0;
        m->page_idx = (uint16_t)i;
        m->segment = seg;
    }

    return seg;
}

void small_segment_v5_release(SmallSegmentV5* seg) {
    if (!seg) return;

    // Verify magic before releasing
    if (seg->magic != SMALL_SEGMENT_V5_MAGIC) {
        return;  // Invalid segment, don't release
    }

    // Clear magic to prevent use-after-free
    seg->magic = 0;

    // Release the 2MiB backing memory
    munmap((void*)seg->base, SMALL_SEGMENT_V5_SIZE);

    // Phase v5-3: Single segment - direct reset
    g_tls_segment_v5.in_use = 0;
    g_tls_segment_v5.used_pages = 0;
}

// ============================================================================
// Page Allocation from Segment (Phase v5-2 fix)
// ============================================================================

SmallPageMetaV5* small_segment_v5_alloc_page(void) {
    TLSSegmentSlot* slot = &g_tls_segment_v5;

    // Phase v5-3: Single segment - direct access, no slot search
    if (slot->in_use && slot->used_pages != 0xFFFFFFFF) {
        // Segment has free pages - use __builtin_ctz for O(1) free page find
        uint32_t free_mask = ~slot->used_pages;
        if (free_mask) {
            uint32_t page_idx = (uint32_t)__builtin_ctz(free_mask);
            slot->used_pages |= (1U << page_idx);
            return &slot->seg.page_meta[page_idx];
        }
    }

    // Need new segment
    if (!slot->in_use) {
        SmallSegmentV5* seg = small_segment_v5_acquire();
        if (!seg) {
            return NULL;
        }
        // First page
        slot->used_pages |= 1U;
        return &seg->page_meta[0];
    }

    return NULL;  // Segment full (shouldn't happen with page recycling)
}

void small_segment_v5_free_page(SmallPageMetaV5* page) {
    if (!page || !page->segment) {
        return;
    }

    // Phase v5-3: Single segment - direct bitmap clear
    TLSSegmentSlot* slot = &g_tls_segment_v5;
    if (slot->in_use && page->segment == &slot->seg) {
        slot->used_pages &= ~(1U << page->page_idx);
        // Keep segment for reuse even if empty
    }
}

// ============================================================================
// O(1) Page Metadata Lookup (Phase v5-3: Single segment, no search)
// ============================================================================

SmallPageMetaV5* small_segment_v5_page_meta_of(void* ptr) {
    if (unlikely(!ptr)) {
        return NULL;
    }

    // Phase v5-3: Single segment - direct TLS access, no loop
    TLSSegmentSlot* slot = &g_tls_segment_v5;
    if (unlikely(!slot->in_use)) {
        return NULL;
    }

    SmallSegmentV5* seg = &slot->seg;
    uintptr_t addr = (uintptr_t)ptr;
    uintptr_t seg_base = seg->base;

    // Check if ptr is within this segment's range
    if (unlikely(addr < seg_base || addr >= seg_base + SMALL_SEGMENT_V5_SIZE)) {
        return NULL;
    }

    // Compute page index via shift (O(1))
    size_t page_idx = (addr - seg_base) >> SMALL_SEGMENT_V5_PAGE_SHIFT;

    // Bounds check (should always pass if within segment)
    if (unlikely(page_idx >= seg->num_pages)) {
        return NULL;
    }

    SmallPageMetaV5* page = &seg->page_meta[page_idx];

    // Validate that this page is actually in use
    if (unlikely(page->capacity == 0)) {
        return NULL;
    }

    return page;
}