/**
 * @file allocator.c
 * @brief Memory allocator implementation
 *
 * Multi-Vendor GPU Aggregation Layer for Linux
 *
 * This module implements the core memory allocation and management.
 * It handles system memory allocations and interfaces with backend-specific
 * allocators (DMA-BUF, GPU-specific, etc.)
 */

#include "memory_internal.h"
#include "mvgal/mvgal_gpu.h"
#include "mvgal/mvgal_log.h"
#include <stdlib.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>

// Global memory state
static mvgal_memory_state_t g_memory_state = {0};

/**
 * @brief Get the global memory state
 */
mvgal_memory_state_t *mvgal_memory_get_state(void) {
    return &g_memory_state;
}

/**
 * @brief Initialize memory module
 */
mvgal_error_t mvgal_memory_module_init(void) {
    if (g_memory_state.initialized) {
        return MVGAL_ERROR_ALREADY_INITIALIZED;
    }
    
    pthread_mutex_init(&g_memory_state.lock, NULL);
    g_memory_state.buffers = NULL;
    g_memory_state.total_allocated = 0;
    g_memory_state.total_freed = 0;
    
    // Initialize DMA-BUF heaps
    g_memory_state.system_heap.fd = -1;
    g_memory_state.system_heap.name = "system";
    g_memory_state.gpu_heap_count = 0;
    for (uint32_t i = 0; i < MVGAL_MAX_GPUS; i++) {
        g_memory_state.gpu_heaps[i].fd = -1;
        g_memory_state.gpu_heaps[i].name = NULL;
    }
    
    // Initialize GPU cache
    g_memory_state.gpu_count = 0;
    
    // Try to open system DMA-BUF heap
    int fd = open("/dev/dma_heap/system", O_RDWR);
    if (fd >= 0) {
        g_memory_state.system_heap.fd = fd;
        MVGAL_LOG_INFO("Opened system DMA-BUF heap at /dev/dma_heap/system");
    } else {
        MVGAL_LOG_DEBUG("System DMA-BUF heap not available (this is normal on older kernels)");
    }
    
    // Scan for GPUs
    int32_t count = mvgal_gpu_get_count();
    if (count > 0) {
        g_memory_state.gpu_count = (uint32_t)count;
        mvgal_gpu_enumerate(g_memory_state.gpus, MVGAL_MAX_GPUS);
        MVGAL_LOG_INFO("Memory module detected %d GPUs", g_memory_state.gpu_count);
    } else {
        MVGAL_LOG_WARN("No GPUs detected by memory module");
    }
    
    // Initialize unified heap regions
    mvgal_heap_init_regions(&g_memory_state);
    
    g_memory_state.initialized = true;
    MVGAL_LOG_INFO("Memory module initialized");
    
    return MVGAL_SUCCESS;
}

/**
 * @brief Shutdown memory module
 */
void mvgal_memory_module_shutdown(void) {
    if (!g_memory_state.initialized) {
        return;
    }
    
    pthread_mutex_lock(&g_memory_state.lock);
    
    // Close DMA-BUF heap
    if (g_memory_state.system_heap.fd >= 0) {
        close(g_memory_state.system_heap.fd);
        g_memory_state.system_heap.fd = -1;
    }
    
    // Close GPU heaps
    for (uint32_t i = 0; i < g_memory_state.gpu_heap_count; i++) {
        if (g_memory_state.gpu_heaps[i].fd >= 0) {
            close(g_memory_state.gpu_heaps[i].fd);
            g_memory_state.gpu_heaps[i].fd = -1;
        }
    }
    
    // Free remaining buffers
    struct mvgal_buffer *buffer = g_memory_state.buffers;
    while (buffer != NULL) {
        struct mvgal_buffer *next = buffer->next;
        mvgal_buffer_free_internal(buffer);
        buffer = next;
    }
    
    g_memory_state.buffers = NULL;
    g_memory_state.initialized = false;
    
    pthread_mutex_unlock(&g_memory_state.lock);
    pthread_mutex_destroy(&g_memory_state.lock);
    
    MVGAL_LOG_INFO("Memory module shut down");
}

/**
 * @brief Increment buffer reference count
 */
void mvgal_buffer_retain(struct mvgal_buffer *buffer) {
    if (buffer == NULL) {
        return;
    }
    atomic_fetch_add(&buffer->refcount, 1);
    MVGAL_LOG_DEBUG("Buffer %p retained, refcount: %u", (void *)buffer, 
                   atomic_load(&buffer->refcount));
}

/**
 * @brief Decrement buffer reference count and possibly free
 */
void mvgal_buffer_release(struct mvgal_buffer *buffer) {
    if (buffer == NULL) {
        return;
    }
    
    uint32_t old_count = atomic_fetch_sub(&buffer->refcount, 1);
    if (old_count == 1) {
        mvgal_buffer_free_internal(buffer);
    }
}

/**
 * @brief Allocate buffer internal
 */
mvgal_error_t mvgal_buffer_allocate_internal(
    const mvgal_memory_alloc_info_t *alloc_info,
    struct mvgal_buffer **buffer_out
) {
    if (alloc_info == NULL || buffer_out == NULL) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }
    
    if (alloc_info->size == 0) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }
    
    // Align size to alignment requirement
    size_t aligned_size = alloc_info->size;
    if (alloc_info->alignment > 0) {
        aligned_size = (aligned_size + alloc_info->alignment - 1) & 
                      ~(alloc_info->alignment - 1);
    }
    
    struct mvgal_buffer *buffer = calloc(1, sizeof(struct mvgal_buffer));
    if (buffer == NULL) {
        MVGAL_LOG_ERROR("Failed to allocate buffer structure");
        return MVGAL_ERROR_OUT_OF_MEMORY;
    }
    
    // Set up buffer state
    buffer->size = alloc_info->size;
    buffer->alignment = alloc_info->alignment;
    buffer->flags = alloc_info->flags;
    buffer->memory_type = alloc_info->memory_type;
    buffer->sharing_mode = alloc_info->sharing_mode;
    buffer->access = alloc_info->access;
    buffer->gpu_mask = alloc_info->gpu_mask;
    buffer->priority = alloc_info->priority;
    buffer->host_ptr = NULL;
    buffer->host_offset = 0;
    buffer->dmabuf_fd = -1;
    buffer->dmabuf_offset = 0;
    buffer->dmabuf_owner = false;
    buffer->fence = NULL;
    buffer->callback = NULL;
    buffer->callback_user_data = NULL;
    buffer->original_ptr = NULL;
    buffer->should_free = false;
    buffer->gpu_binding_count = 0;
    buffer->preferred_gpu = 0;
    buffer->last_used_gpu = UINT32_MAX;
    atomic_init(&buffer->migration_in_progress, 0);
    buffer->migration_fence = NULL;

    buffer->alias_parent = NULL;
    buffer->alias_count = 0;
    buffer->alias_offset = 0;
    buffer->alias_size = 0;
    buffer->alias_flags = MVGAL_BUFFER_ALIAS_FLAG_NONE;
    buffer->is_alias = false;
    buffer->alias_detached = false;
    for (uint32_t i = 0; i < 16; i++) {
        buffer->aliases[i] = NULL;
    }
    
    atomic_init(&buffer->refcount, 1);
    for (uint32_t i = 0; i < MVGAL_MAX_GPUS; i++) {
        atomic_init(&buffer->access_count[i], 0);
    }
    atomic_init(&buffer->bytes_read, 0);
    atomic_init(&buffer->bytes_written, 0);
    atomic_init(&buffer->bytes_transferred, 0);
    
    // Try DMA-BUF allocation first if requested or if GPU mask has multiple GPUs
    if ((alloc_info->flags & MVGAL_MEMORY_FLAG_DMA_BUF) ||
        (alloc_info->sharing_mode == MVGAL_MEMORY_SHARING_DMA_BUF) ||
        (alloc_info->gpu_mask != 0 && __builtin_popcount(alloc_info->gpu_mask) > 1)) {
        
        if (mvgal_dmabuf_is_supported(alloc_info->gpu_mask)) {
            int fd = -1;
            void *ptr = NULL;
            
            mvgal_error_t err = mvgal_dmabuf_allocate(
                aligned_size,
                alloc_info->alignment,
                alloc_info->flags,
                alloc_info->gpu_mask,
                &fd,
                &ptr
            );
            
            if (err == MVGAL_SUCCESS && fd >= 0 && ptr != NULL) {
                buffer->host_ptr = ptr;
                buffer->dmabuf_fd = fd;
                buffer->dmabuf_owner = true;
                buffer->backend = MVGAL_BUFFER_BACKEND_DMA_BUF;
                buffer->state = MVGAL_BUFFER_STATE_ALLOCATED | 
                               MVGAL_BUFFER_STATE_MAPPED;
                
                // Map to all GPUs in the mask
                for (uint32_t i = 0; i < MVGAL_MAX_GPUS; i++) {
                    if (alloc_info->gpu_mask & (1 << i)) {
                        // For now, just store the CPU pointer
                        // GPU-specific addressing would be handled by the driver
                        buffer->gpu_bindings[i].gpu_index = i;
                        buffer->gpu_bindings[i].mapping = ptr;
                        buffer->gpu_bindings[i].readable = true;
                        buffer->gpu_bindings[i].writable = 
                            (alloc_info->access & MVGAL_MEMORY_ACCESS_WRITE) != 0;
                        buffer->gpu_binding_count++;
                    }
                }
                
                MVGAL_LOG_DEBUG("Allocated DMA-BUF buffer: size=%zu, fd=%d, ptr=%p",
                               aligned_size, fd, ptr);
                *buffer_out = buffer;
                return MVGAL_SUCCESS;
            }
            
            // DMA-BUF allocation failed, fall through to system allocation
            MVGAL_LOG_DEBUG("DMA-BUF allocation failed, falling back to system memory");
        }
    }
    
    // Fall back to system memory allocation
    void *ptr = NULL;
    
    if (alloc_info->flags & MVGAL_MEMORY_FLAG_HOST_VALID) {
        // Host-visible allocations should be directly CPU-accessible and
        // freeable through the normal system allocator path.
        if (alloc_info->alignment > sizeof(void *) && alloc_info->alignment != 0) {
            int ret = posix_memalign(&ptr, alloc_info->alignment, aligned_size);
            if (ret != 0) {
                ptr = NULL;
            }
        } else {
            ptr = malloc(aligned_size);
        }
    } else {
        // Use regular malloc
        ptr = malloc(aligned_size);
    }
    
    if (ptr == NULL) {
        MVGAL_LOG_ERROR("Failed to allocate %zu bytes", aligned_size);
        free(buffer);
        return MVGAL_ERROR_OUT_OF_MEMORY;
    }
    
    // Zero-initialize if requested
    if (alloc_info->flags & MVGAL_MEMORY_FLAG_ZERO_INITIALIZED) {
        memset(ptr, 0, aligned_size);
    }
    
    buffer->host_ptr = ptr;
    buffer->original_ptr = ptr;
    buffer->should_free = true;
    buffer->backend = MVGAL_BUFFER_BACKEND_SYSTEM;
    buffer->state = MVGAL_BUFFER_STATE_ALLOCATED;
    
    if (alloc_info->flags & MVGAL_MEMORY_FLAG_HOST_VALID) {
        buffer->state |= MVGAL_BUFFER_STATE_MAPPED;
    }
    
    MVGAL_LOG_DEBUG("Allocated system buffer: size=%zu, ptr=%p", 
                   aligned_size, ptr);
    
    *buffer_out = buffer;
    return MVGAL_SUCCESS;
}

/**
 * @brief Free buffer internal
 */
void mvgal_buffer_free_internal(struct mvgal_buffer *buffer) {
    if (buffer == NULL) {
        return;
    }
    
    MVGAL_LOG_DEBUG("Freeing buffer: %p, ptr=%p, backend=%d, is_alias=%d", 
                   (void *)buffer, buffer->host_ptr, buffer->backend, buffer->is_alias);
    
    // If this is a parent with aliases, detach all aliases first
    if (buffer->alias_count > 0) {
        for (uint32_t i = 0; i < buffer->alias_count; i++) {
            struct mvgal_buffer *alias = buffer->aliases[i];
            if (alias != NULL) {
                alias->alias_parent = NULL;
                alias->alias_detached = true;
                alias->is_alias = false;
                alias->should_free = true;  /* Alias owns its copy now */
                MVGAL_LOG_DEBUG("Alias %p detached from parent %p", (void *)alias, (void *)buffer);
            }
        }
        buffer->alias_count = 0;
    }
    
    // If this is an alias, remove from parent's alias list
    if (buffer->is_alias && buffer->alias_parent != NULL && !buffer->alias_detached) {
        struct mvgal_buffer *parent = buffer->alias_parent;
        if (parent != NULL) {
            for (uint32_t i = 0; i < parent->alias_count; i++) {
                if (parent->aliases[i] == buffer) {
                    parent->aliases[i] = parent->aliases[parent->alias_count - 1];
                    parent->alias_count--;
                    break;
                }
            }
            mvgal_buffer_release(parent); // Drop parent ref held by alias
        }
    }
    
    // Unbind from all GPUs
    mvgal_buffer_unbind(buffer);
    
    // Free DMA-BUF if we own it
    if (buffer->dmabuf_owner && buffer->dmabuf_fd >= 0) {
        mvgal_dmabuf_free(buffer->dmabuf_fd, buffer->host_ptr, buffer->size);
        buffer->dmabuf_fd = -1;
        buffer->host_ptr = NULL;
    }
    
    // Free system memory if we own it
    if (buffer->should_free && buffer->original_ptr != NULL) {
        free(buffer->original_ptr);
        buffer->original_ptr = NULL;
    }
    
    // Release fence if any
    if (buffer->fence != NULL) {
        mvgal_fence_destroy(buffer->fence);
        buffer->fence = NULL;
    }
    
    // Remove from global list
    pthread_mutex_lock(&g_memory_state.lock);
    if (buffer->prev != NULL) {
        buffer->prev->next = buffer->next;
    } else {
        g_memory_state.buffers = buffer->next;
    }
    if (buffer->next != NULL) {
        buffer->next->prev = buffer->prev;
    }
    pthread_mutex_unlock(&g_memory_state.lock);
    
    g_memory_state.total_freed += buffer->size;
    
    free(buffer);
}

/**
 * @brief Bind buffer to a GPU
 */
mvgal_error_t mvgal_buffer_bind_to_gpu(
    struct mvgal_buffer *buffer,
    uint32_t gpu_index,
    bool writable
) {
    if (buffer == NULL || gpu_index >= MVGAL_MAX_GPUS) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }
    
    // Check if GPU exists
    if (gpu_index >= g_memory_state.gpu_count) {
        return MVGAL_ERROR_GPU_NOT_FOUND;
    }
    
    // Check if already bound
    for (uint32_t i = 0; i < buffer->gpu_binding_count; i++) {
        if (buffer->gpu_bindings[i].gpu_index == gpu_index) {
            // Already bound, update flags
            buffer->gpu_bindings[i].writable = writable;
            buffer->gpu_bindings[i].readable = true; // Must be readable if writable
            return MVGAL_SUCCESS;
        }
    }
    
    // Check if GPU is in the mask
    if (!(buffer->gpu_mask & (1 << gpu_index))) {
        buffer->gpu_mask |= (1 << gpu_index);
    }
    
    // Add binding
    if (buffer->gpu_binding_count < MVGAL_MAX_GPUS) {
        mvgal_gpu_binding_t *binding = &buffer->gpu_bindings[buffer->gpu_binding_count];
        binding->gpu_index = gpu_index;
        binding->mapping = buffer->host_ptr; // For now, use CPU pointer
        binding->writable = writable;
        binding->readable = true;
        binding->dmabuf_fd = buffer->dmabuf_fd;
        buffer->gpu_binding_count++;
        
        MVGAL_LOG_DEBUG("Buffer %p bound to GPU %u (writable=%d)", 
                       (void *)buffer, gpu_index, writable);
        return MVGAL_SUCCESS;
    }
    
    MVGAL_LOG_ERROR("Cannot bind buffer to GPU %u: too many bindings", gpu_index);
    return MVGAL_ERROR_MEMORY;
}

/**
 * @brief Unbind buffer from all GPUs
 */
void mvgal_buffer_unbind(struct mvgal_buffer *buffer) {
    if (buffer == NULL) {
        return;
    }
    
    for (uint32_t i = 0; i < buffer->gpu_binding_count; i++) {
        if (buffer->gpu_bindings[i].mapping != NULL &&
            buffer->gpu_bindings[i].mapping != buffer->host_ptr) {
            // Unmap GPU-specific mapping if different from CPU pointer
            // In a real implementation, this would call GPU-specific unmap
            buffer->gpu_bindings[i].mapping = NULL;
        }
    }
    
    buffer->gpu_binding_count = 0;
}

/**
 * @brief Check if DMA-BUF is supported for the given GPU mask
 */
bool mvgal_dmabuf_is_supported(uint32_t gpu_mask) {
    if (gpu_mask == 0) {
        return false;
    }
    
    // Check each GPU in the mask
    for (uint32_t i = 0; i < g_memory_state.gpu_count; i++) {
        if (gpu_mask & (1 << i)) {
            const mvgal_gpu_descriptor_t *gpu = &g_memory_state.gpus[i];
            // Check if GPU supports DMA-BUF
            if (!(gpu->features & MVGAL_FEATURE_DMA_BUF)) {
                return false;
            }
        }
    }
    
    return true;
}

/* ------------------------------------------------------------------ */
/*  Buffer alias operations                                           */
/* ------------------------------------------------------------------ */

mvgal_error_t mvgal_memory_create_alias(
    void *context,
    mvgal_buffer_t parent,
    uint64_t offset,
    size_t size,
    mvgal_buffer_alias_flags_t alias_flags,
    mvgal_buffer_t *buffer
) {
    (void)context;

    if (parent == NULL || buffer == NULL) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    struct mvgal_buffer *p = (struct mvgal_buffer *)parent;

    /* Validate offset and size against parent */
    if (offset >= p->size) {
        MVGAL_LOG_ERROR("Alias offset %lu exceeds parent size %zu",
                        (unsigned long)offset, p->size);
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    size_t alias_size = size;
    if (alias_size == 0) {
        alias_size = p->size - offset;
    }
    if (offset + alias_size > p->size) {
        MVGAL_LOG_ERROR("Alias region [%lu, %lu) exceeds parent size %zu",
                        (unsigned long)offset, (unsigned long)(offset + alias_size), p->size);
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    /* Check parent alias limit */
    if (p->alias_count >= 16) {
        MVGAL_LOG_ERROR("Parent buffer %p has too many aliases (max 16)", (void *)p);
        return MVGAL_ERROR_MEMORY;
    }

    /* Allocate alias buffer structure */
    struct mvgal_buffer *alias = calloc(1, sizeof(struct mvgal_buffer));
    if (alias == NULL) {
        MVGAL_LOG_ERROR("Failed to allocate alias buffer");
        return MVGAL_ERROR_OUT_OF_MEMORY;
    }

    /* Copy parent buffer properties */
    alias->size = alias_size;
    alias->alignment = p->alignment;
    alias->flags = p->flags;
    alias->memory_type = p->memory_type;
    alias->sharing_mode = p->sharing_mode;
    alias->access = p->access;
    alias->gpu_mask = p->gpu_mask;
    alias->priority = p->priority;
    alias->backend = p->backend;

    /* Shared host pointer with offset */
    alias->host_ptr = (p->host_ptr != NULL)
                      ? (void *)((char *)p->host_ptr + offset)
                      : NULL;
    alias->host_offset = offset;
    alias->dmabuf_fd = p->dmabuf_fd;
    alias->dmabuf_offset = p->dmabuf_offset + offset;
    alias->dmabuf_owner = false;  /* Parent owns the DMA-BUF */
    alias->state = p->state;
    alias->original_ptr = (p->host_ptr != NULL)
                          ? (void *)((char *)p->host_ptr + offset)
                          : NULL;
    alias->should_free = false;   /* Parent owns the memory */

    /* Set up alias tracking */
    alias->alias_parent = p;
    alias->alias_offset = offset;
    alias->alias_size = alias_size;
    alias->alias_flags = alias_flags;
    alias->is_alias = true;
    alias->alias_detached = false;
    alias->alias_count = 0;

    atomic_init(&alias->refcount, 1);
    for (uint32_t i = 0; i < MVGAL_MAX_GPUS; i++) {
        atomic_init(&alias->access_count[i], 0);
    }
    atomic_init(&alias->bytes_read, 0);
    atomic_init(&alias->bytes_written, 0);
    atomic_init(&alias->bytes_transferred, 0);

    /* Copy GPU bindings (pointing to same memory) */
    for (uint32_t i = 0; i < p->gpu_binding_count; i++) {
        alias->gpu_bindings[i] = p->gpu_bindings[i];
    }
    alias->gpu_binding_count = p->gpu_binding_count;

    /* Register alias in parent's alias list */
    p->aliases[p->alias_count++] = alias;
    mvgal_buffer_retain(p);  /* Parent held alive while alias exists */

    /* Restrict access based on alias flags */
    if (alias_flags & MVGAL_BUFFER_ALIAS_FLAG_READ_ONLY) {
        for (uint32_t i = 0; i < alias->gpu_binding_count; i++) {
            alias->gpu_bindings[i].writable = false;
        }
        alias->access &= ~MVGAL_MEMORY_ACCESS_WRITE;
    }
    if (alias_flags & MVGAL_BUFFER_ALIAS_FLAG_WRITE_ONLY) {
        alias->access &= ~MVGAL_MEMORY_ACCESS_READ;
    }

    /* Add to global buffer list for tracking */
    pthread_mutex_lock(&g_memory_state.lock);
    alias->next = g_memory_state.buffers;
    alias->prev = NULL;
    if (g_memory_state.buffers != NULL) {
        g_memory_state.buffers->prev = alias;
    }
    g_memory_state.buffers = alias;
    pthread_mutex_unlock(&g_memory_state.lock);

    *buffer = alias;

    MVGAL_LOG_DEBUG("Alias created: %p -> parent %p offset=%lu size=%zu flags=%u",
                    (void *)alias, (void *)p, (unsigned long)offset, alias_size, alias_flags);

    return MVGAL_SUCCESS;
}

mvgal_error_t mvgal_memory_get_alias_info(
    mvgal_buffer_t buffer,
    mvgal_buffer_alias_info_t *info
) {
    if (buffer == NULL || info == NULL) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    struct mvgal_buffer *b = (struct mvgal_buffer *)buffer;

    if (!b->is_alias) {
        return MVGAL_ERROR_NOT_FOUND;
    }

    info->parent = b->alias_parent;
    info->alias = buffer;
    info->offset = b->alias_offset;
    info->size = b->alias_size;
    info->flags = (mvgal_memory_flags_t)b->alias_flags;
    info->gpu_mask = b->gpu_mask;

    return MVGAL_SUCCESS;
}

mvgal_error_t mvgal_memory_alias_detach(
    mvgal_buffer_t buffer
) {
    if (buffer == NULL) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    struct mvgal_buffer *b = (struct mvgal_buffer *)buffer;

    if (!b->is_alias || b->alias_detached) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    struct mvgal_buffer *parent = b->alias_parent;
    if (parent == NULL) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    /* Allocate private copy of the aliased region */
    size_t copy_size = b->alias_size;
    void *private_copy = malloc(copy_size);
    if (private_copy == NULL) {
        MVGAL_LOG_ERROR("Failed to allocate private copy for alias detach (%zu bytes)", copy_size);
        return MVGAL_ERROR_OUT_OF_MEMORY;
    }

    /* Copy data from shared memory */
    if (parent->host_ptr != NULL) {
        memcpy(private_copy, parent->host_ptr, copy_size);
    }

    /* Remove from parent's alias list */
    for (uint32_t i = 0; i < parent->alias_count; i++) {
        if (parent->aliases[i] == b) {
            parent->aliases[i] = parent->aliases[parent->alias_count - 1];
            parent->alias_count--;
            break;
        }
    }
    mvgal_buffer_release(parent); /* Drop parent ref */

    /* Convert alias to independent buffer */
    b->alias_parent = NULL;
    b->alias_detached = true;
    b->is_alias = false;
    b->host_ptr = private_copy;
    b->host_offset = 0;
    b->original_ptr = private_copy;
    b->should_free = true;
    b->alias_offset = 0;
    b->alias_size = 0;
    b->alias_flags = MVGAL_BUFFER_ALIAS_FLAG_NONE;

    MVGAL_LOG_DEBUG("Alias %p detached from parent %p, private copy at %p",
                    (void *)b, (void *)parent, private_copy);

    return MVGAL_SUCCESS;
}

/* ------------------------------------------------------------------ */
/*  Unified heap helpers                                              */
/* ------------------------------------------------------------------ */

void mvgal_heap_init_regions(mvgal_memory_state_t *state)
{
    uint32_t idx = 0;

    // Region 0: System (host) memory
    state->regions[idx].gpu_index = UINT32_MAX;
    state->regions[idx].base_address = 0;
    state->regions[idx].size = 0;           // System RAM is virtually unlimited
    state->regions[idx].free_bytes = 0;     // Not tracked for system heap
    state->regions[idx].allocation_count = 0;
    state->regions[idx].is_device_memory = false;
    idx++;

    // Regions 1..N: Per-GPU VRAM
    for (uint32_t i = 0; i < state->gpu_count && idx <= MVGAL_MAX_GPUS; i++) {
        const mvgal_gpu_descriptor_t *gpu = &state->gpus[i];
        if (gpu->vram_total == 0) {
            continue;
        }
        state->regions[idx].gpu_index = i;
        state->regions[idx].base_address = 0;  // GPU VRAM base not exposed by descriptor
        state->regions[idx].size = gpu->vram_total;
        state->regions[idx].free_bytes = gpu->vram_free;
        state->regions[idx].allocation_count = 0;
        state->regions[idx].is_device_memory = true;
        idx++;
    }
    state->region_count = idx;

    MVGAL_LOG_INFO("Unified heap initialized with %u regions", idx);
}

void mvgal_heap_update_free(mvgal_memory_state_t *state, uint32_t region_index)
{
    if (region_index >= state->region_count) {
        return;
    }

    // System heap (index 0) doesn't track free bytes
    if (region_index == 0) {
        return;
    }

    // Recalculate free bytes based on active allocations.
    // In a full implementation, this would scan VRAM allocation tables.
    // For now, we trust the initial free_bytes minus tracked allocations.
    uint32_t gpu_idx = state->regions[region_index].gpu_index;
    if (gpu_idx < state->gpu_count) {
        uint64_t total_size = state->regions[region_index].size;
        uint64_t used = 0;

        pthread_mutex_lock(&state->lock);
        struct mvgal_buffer *buf = state->buffers;
        while (buf != NULL) {
            if (buf->backend == MVGAL_BUFFER_BACKEND_GPU &&
                (buf->gpu_mask & (1 << gpu_idx))) {
                used += buf->size;
            }
            buf = buf->next;
        }
        pthread_mutex_unlock(&state->lock);

        state->regions[region_index].free_bytes =
            (used < total_size) ? (total_size - used) : 0;
    }
}

uint32_t mvgal_heap_select_best(
    const mvgal_memory_state_t *state,
    uint32_t gpu_mask,
    mvgal_memory_flags_t flags)
{
    (void)flags;

    if (state->region_count == 0) {
        return UINT32_MAX;
    }

    // If no GPU mask, default to system heap (index 0)
    if (gpu_mask == 0) {
        return 0;
    }

    // Find the GPU with the most free space that matches the mask
    uint32_t best_region = UINT32_MAX;
    uint64_t best_free = 0;

    for (uint32_t i = 1; i < state->region_count; i++) {
        uint32_t gpu_idx = state->regions[i].gpu_index;
        if (gpu_idx < MVGAL_MAX_GPUS && (gpu_mask & (1 << gpu_idx))) {
            if (state->regions[i].free_bytes > best_free) {
                best_free = state->regions[i].free_bytes;
                best_region = i;
            }
        }
    }

    // Fallback to system heap if no VRAM region found
    return (best_region != UINT32_MAX) ? best_region : 0;
}

/* ------------------------------------------------------------------ */
/*  Public API: Unified heap query                                    */
/* ------------------------------------------------------------------ */

mvgal_error_t mvgal_memory_get_unified_heap(mvgal_unified_heap_t *heap)
{
    if (heap == NULL) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    mvgal_memory_state_t *state = mvgal_memory_get_state();
    if (!state->initialized) {
        return MVGAL_ERROR_NOT_INITIALIZED;
    }

    // Recalculate free bytes for each GPU region
    for (uint32_t i = 0; i < state->region_count; i++) {
        mvgal_heap_update_free(state, i);
    }

    // Copy regions into the output struct
    uint32_t count = state->region_count;
    if (count > MVGAL_MAX_GPUS + 1) {
        count = MVGAL_MAX_GPUS + 1;
    }

    heap->region_count = count;
    heap->total_size = 0;
    heap->total_free = 0;

    for (uint32_t i = 0; i < count; i++) {
        heap->regions[i] = state->regions[i];
        heap->total_size += state->regions[i].size;
        heap->total_free += state->regions[i].free_bytes;
    }

    return MVGAL_SUCCESS;
}

/* ------------------------------------------------------------------ */
/*  Public API: Buffer migration                                      */
/* ------------------------------------------------------------------ */

mvgal_error_t mvgal_memory_migrate(
    mvgal_buffer_t buffer,
    uint32_t target_gpu_index,
    mvgal_fence_t fence)
{
    if (buffer == NULL) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    struct mvgal_buffer *b = (struct mvgal_buffer *)buffer;

    if (target_gpu_index >= MVGAL_MAX_GPUS) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    // Check if already resident on the target
    for (uint32_t i = 0; i < b->gpu_binding_count; i++) {
        if (b->gpu_bindings[i].gpu_index == target_gpu_index &&
            b->gpu_bindings[i].mapping != NULL) {
            // Already resident, signal fence immediately if provided
            if (fence != NULL) {
                mvgal_fence_signal_with_status(fence, MVGAL_SUCCESS);
            }
            return MVGAL_SUCCESS;
        }
    }

    // Prevent concurrent migrations
    if (atomic_exchange(&b->migration_in_progress, 1) != 0) {
        return MVGAL_ERROR_BUSY;
    }

    // Store the migration fence so the completion path can signal it
    b->migration_fence = fence;

    // Bind the buffer to the target GPU (triggers DMA-BUF share or copy)
    mvgal_error_t err = mvgal_buffer_bind_to_gpu(b, target_gpu_index, true);
    if (err != MVGAL_SUCCESS) {
        atomic_store(&b->migration_in_progress, 0);
        b->migration_fence = NULL;
        if (fence != NULL) {
            mvgal_fence_signal_with_status(fence, err);
        }
        return err;
    }

    // Update the last-used GPU hint
    b->last_used_gpu = target_gpu_index;

    // Update preferred GPU mask
    b->preferred_gpu |= (1 << target_gpu_index);

    // Migration is logically complete.
    // In a full implementation with async DMA engines, the fence would
    // be signaled by a completion callback. For now, signal it here.
    atomic_store(&b->migration_in_progress, 0);
    if (fence != NULL) {
        mvgal_fence_signal_with_status(fence, MVGAL_SUCCESS);
    }

    MVGAL_LOG_DEBUG("Buffer %p migrated to GPU %u", (void *)b, target_gpu_index);
    return MVGAL_SUCCESS;
}

/* ------------------------------------------------------------------ */
/*  Public API: Page-fault handler                                    */
/* ------------------------------------------------------------------ */

mvgal_error_t mvgal_memory_handle_fault(
    mvgal_buffer_t buffer,
    uint32_t faulting_gpu_index,
    uint64_t address)
{
    if (buffer == NULL) {
        return MVGAL_ERROR_INVALID_ARGUMENT;
    }

    struct mvgal_buffer *b = (struct mvgal_buffer *)buffer;
    (void)address;

    MVGAL_LOG_INFO("Page fault for buffer %p on GPU %u (address 0x%lx)",
                   (void *)b, faulting_gpu_index, (unsigned long)address);

    // If there's already a migration in progress, wait for it
    if (atomic_load(&b->migration_in_progress) != 0) {
        if (b->migration_fence != NULL) {
            mvgal_fence_wait(b->migration_fence, UINT32_MAX);
        }
        // After migration completes, check if we're on the right GPU
        for (uint32_t i = 0; i < b->gpu_binding_count; i++) {
            if (b->gpu_bindings[i].gpu_index == faulting_gpu_index) {
                return MVGAL_SUCCESS;
            }
        }
    }

    // Trigger on-demand migration to the faulting GPU
    mvgal_error_t err = mvgal_memory_migrate(buffer, faulting_gpu_index, NULL);
    if (err != MVGAL_SUCCESS) {
        MVGAL_LOG_ERROR("Page fault migration failed for buffer %p to GPU %u: %d",
                        (void *)b, faulting_gpu_index, err);
        return err;
    }

    MVGAL_LOG_DEBUG("Page fault handled: buffer %p migrated to GPU %u",
                    (void *)b, faulting_gpu_index);
    return MVGAL_SUCCESS;
}
