// smallsegment_v6.c - SmallSegment v6 実装（Phase v6-2）

#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <stdint.h>
#include <unistd.h>
#include "box/smallsegment_v6_box.h"

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

// TLS segment slot (metadata stored in TLS, not in mapped memory)
typedef struct {
    SmallSegmentV6 seg;
    int in_use;
    void* mmap_base;  // Actual mmap base (for munmap)
    size_t mmap_size; // Actual mmap size (for munmap)
} TLSSegmentSlotV6;

static __thread TLSSegmentSlotV6 g_tls_segment_v6;

/// Acquire 2MiB aligned segment for current thread (called once per thread)
/// Allocates a 2MiB segment with proper alignment via mmap
/// @return: Segment pointer on success, NULL on failure
SmallSegmentV6* small_segment_v6_acquire_for_thread(void) {
    TLSSegmentSlotV6* slot = &g_tls_segment_v6;

    if (slot->in_use) {
        return &slot->seg;  // Already acquired
    }

    // Allocate 2MiB aligned segment
    // Use mmap with MAP_ANONYMOUS which typically gives aligned addresses for large allocations
    void* mem = mmap(NULL, SMALL_SEGMENT_V6_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;
    void* mmap_base = mem;
    size_t mmap_size = SMALL_SEGMENT_V6_SIZE;

    // Check if we got 2MiB alignment
    if ((addr & (SMALL_SEGMENT_V6_SIZE - 1)) != 0) {
        // Not aligned, need to reallocate with overallocation
        munmap(mem, SMALL_SEGMENT_V6_SIZE);

        // Allocate 4MiB to ensure we can find a 2MiB aligned region
        size_t alloc_size = SMALL_SEGMENT_V6_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_V6_SIZE - 1) & ~(SMALL_SEGMENT_V6_SIZE - 1);

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

        // Keep track of actual mmap base for munmap
        mmap_base = mem;
        mmap_size = alloc_size;
    }

    // Initialize segment structure in TLS (not in mapped memory)
    SmallSegmentV6* seg = &slot->seg;
    slot->in_use = 1;
    slot->mmap_base = mmap_base;
    slot->mmap_size = mmap_size;

    seg->base = addr;
    seg->num_pages = SMALL_PAGES_PER_SEGMENT;
    seg->owner_tid = (uint32_t)getpid();  // Simple owner tracking
    seg->magic = SMALL_SEGMENT_V6_MAGIC;

    // Initialize all page metadata
    for (uint32_t i = 0; i < seg->num_pages; i++) {
        SmallPageMetaV6* 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;
}

/// Release segment and unmap memory (cleanup)
/// @param seg: Segment to release
void small_segment_v6_release(SmallSegmentV6* seg) {
    if (!seg) return;
    if (seg->magic != SMALL_SEGMENT_V6_MAGIC) return;

    TLSSegmentSlotV6* slot = &g_tls_segment_v6;
    if (seg != &slot->seg) return;  // Not our segment

    seg->magic = 0;  // Invalidate
    munmap(slot->mmap_base, slot->mmap_size);

    slot->in_use = 0;
    slot->mmap_base = NULL;
    slot->mmap_size = 0;
}

/// O(1) ptr -> page_meta lookup using segment metadata in TLS
/// Maps any pointer within TLS segment to its page metadata
/// @param ptr: Pointer to lookup (USER or BASE pointer)
/// @return: Page metadata pointer if found and valid, NULL otherwise
SmallPageMetaV6* small_page_meta_v6_of(void* ptr) {
    if (unlikely(!ptr)) {
        return NULL;
    }

    TLSSegmentSlotV6* slot = &g_tls_segment_v6;

    // Check if segment is initialized
    if (unlikely(!slot->in_use)) {
        return NULL;
    }

    SmallSegmentV6* seg = &slot->seg;

    // Check if ptr is within our segment range
    if (unlikely(!small_ptr_in_segment_v6(seg, ptr))) {
        return NULL;
    }

    // Calculate page index using macro
    uintptr_t addr = (uintptr_t)ptr;
    size_t page_idx = SMALL_V6_PAGE_IDX(seg, addr);
    if (unlikely(page_idx >= seg->num_pages)) {
        return NULL;
    }

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

    // Validate that this page is actually in use
    if (unlikely(!small_page_v6_valid(page))) {
        return NULL;
    }

    return page;
}