/**
 * @file memory_internal.h
 * @brief Internal memory management types and declarations
 *
 * Multi-Vendor GPU Aggregation Layer for Linux
 *
 * This header defines internal structures for the memory module.
 * Not part of the public API.
 */

#ifndef MVGAL_MEMORY_INTERNAL_H
#define MVGAL_MEMORY_INTERNAL_H

#include <stddef.h>
#include <stdatomic.h>
#include <pthread.h>
#include "mvgal/mvgal_types.h"
#include "mvgal/mvgal_gpu.h"
#include "mvgal/mvgal_memory.h"

#ifdef __cplusplus
extern "C" {
#endif

/**
 * @brief Maximum number of GPUs supported
 */
#ifndef MVGAL_MAX_GPUS
#define MVGAL_MAX_GPUS 16
#endif

/**
 * @brief Buffer state flags
 */
typedef enum {
    MVGAL_BUFFER_STATE_ALLOCATED = 1 << 0,   ///< Buffer memory allocated
    MVGAL_BUFFER_STATE_MAPPED = 1 << 1,       ///< Mapped to CPU address space
    MVGAL_BUFFER_STATE_EXPORTED = 1 << 2,    ///< Exported as DMA-BUF
    MVGAL_BUFFER_STATE_IMPORTED = 1 << 3,    ///< Imported from DMA-BUF
    MVGAL_BUFFER_STATE_REPLICATED = 1 << 4,  ///< Replicated to multiple GPUs
    MVGAL_BUFFER_STATE_LAZY = 1 << 5,        ///< Lazy allocation pending
    MVGAL_BUFFER_STATE_ERROR = 1 << 6,       ///< Error state
} mvgal_buffer_state_flags_t;

/**
 * @brief Buffer backend type
 */
typedef enum {
    MVGAL_BUFFER_BACKEND_SYSTEM,    ///< System memory (malloc/mmap)
    MVGAL_BUFFER_BACKEND_DMA_BUF,   ///< DMA-BUF allocated
    MVGAL_BUFFER_BACKEND_DMA_BUF_IMPORTED, ///< Imported DMA-BUF
    MVGAL_BUFFER_BACKEND_GPU,        ///< GPU-specific allocation
    MVGAL_BUFFER_BACKEND_WRAPPED,    ///< Wrapped existing memory
    MVGAL_BUFFER_BACKEND_P2P,        ///< P2P accessible memory
} mvgal_buffer_backend_type_t;

/**
 * @brief Internal fence structure
 */
struct mvgal_fence {
    pthread_mutex_t mutex;
    pthread_cond_t cond;
    bool signaled;
    atomic_int refcount;
    mvgal_error_t status;  ///< Status when signaled (MVGAL_SUCCESS or error)
    void *user_data;       ///< Optional user data
};

/**
 * @brief Internal semaphore structure
 */
struct mvgal_semaphore {
    pthread_mutex_t mutex;
    pthread_cond_t cond;
    uint64_t count;
    atomic_int refcount;
};

/**
 * @brief GPU memory binding
 */
typedef struct {
    uint32_t gpu_index;
    uint64_t gpu_address;    ///< GPU virtual address
    void *mapping;          ///< CPU mapping for this GPU (if any)
    bool writable;          ///< Can be written by this GPU
    bool readable;          ///< Can be read by this GPU
    int dmabuf_fd;          ///< DMA-BUF FD for this GPU (if applicable)
} mvgal_gpu_binding_t;

/**
 * @brief Internal buffer structure
 */
struct mvgal_buffer {
    // Reference counting
    atomic_uint refcount;
    
    // State
    mvgal_buffer_state_flags_t state;
    mvgal_buffer_backend_type_t backend;
    
    // Allocation info
    size_t size;
    size_t alignment;
    mvgal_memory_flags_t flags;
    mvgal_memory_type_t memory_type;
    mvgal_memory_sharing_mode_t sharing_mode;
    mvgal_memory_access_flags_t access;
    uint32_t gpu_mask;  ///< Bitmask of GPUs that can access
    uint32_t priority;
    
    // Memory locations
    void *host_ptr;                 ///< CPU pointer (if mapped)
    size_t host_offset;            ///< Offset within host_ptr
    int dmabuf_fd;                  ///< DMA-BUF file descriptor (-1 if none)
    uint64_t dmabuf_offset;         ///< Offset within DMA-BUF
    bool dmabuf_owner;              ///< True if we created the DMA-BUF
    
    // GPU bindings
    mvgal_gpu_binding_t gpu_bindings[MVGAL_MAX_GPUS];
    uint32_t gpu_binding_count;
    
    // Synchronization
    struct mvgal_fence *fence;     ///< Optional fence for this buffer
    
    // Callbacks
    mvgal_memory_callback_t callback;
    void *callback_user_data;
    
    // Statistics
    atomic_ullong bytes_read;
    atomic_ullong bytes_written;
    atomic_ullong bytes_transferred;
    atomic_uint access_count[MVGAL_MAX_GPUS];
    
    // Original allocation (for wrapped memory)
    void *original_ptr;           ///< Original pointer if wrapped
    bool should_free;              ///< Whether to free original_ptr on destroy

    // Migration tracking
    uint32_t preferred_gpu;              ///< Preferred GPU for this buffer (bitmask)
    uint32_t last_used_gpu;              ///< GPU index that last accessed this buffer
    atomic_uint migration_in_progress;   ///< Non-zero if migration is active
    struct mvgal_fence *migration_fence; ///< Fence signaled when migration completes

    // Alias tracking
    struct mvgal_buffer *alias_parent;   ///< Parent buffer (NULL if root)
    struct mvgal_buffer *aliases[16];    ///< Aliases pointing to this buffer
    uint32_t alias_count;                ///< Number of aliases
    uint64_t alias_offset;               ///< Offset within parent buffer
    size_t alias_size;                   ///< Size of alias region (0 = full parent)
    mvgal_buffer_alias_flags_t alias_flags; ///< Alias creation flags
    bool is_alias;                        ///< True if this is an alias buffer
    bool alias_detached;                  ///< True if alias was detached (owns copy)
    
    // Next/prev for tracking (optional)
    struct mvgal_buffer *next;
    struct mvgal_buffer *prev;
};

/**
 * @brief DMA-BUF heap handle
 */
typedef struct {
    int fd;
    const char *name;
} mvgal_dmabuf_heap_t;

/**
 * @brief Memory module global state
 */
typedef struct {
    pthread_mutex_t lock;
    bool initialized;
    struct mvgal_buffer *buffers;  ///< Head of buffer list
    uint64_t total_allocated;
    uint64_t total_freed;
    
    // DMA-BUF heaps
    mvgal_dmabuf_heap_t system_heap;
    mvgal_dmabuf_heap_t gpu_heaps[MVGAL_MAX_GPUS];
    uint32_t gpu_heap_count;
    
    // GPU information (cached from gpu_manager)
    mvgal_gpu_descriptor_t gpus[MVGAL_MAX_GPUS];
    uint32_t gpu_count;

    // Unified heap regions (populated by mvgal_heap_init_regions)
    mvgal_heap_region_t regions[MVGAL_MAX_GPUS + 1];
    uint32_t region_count;
} mvgal_memory_state_t;

/**
 * @brief Get the global memory state
 */
mvgal_memory_state_t *mvgal_memory_get_state(void);

/**
 * @brief Initialize memory module
 */
mvgal_error_t mvgal_memory_module_init(void);

/**
 * @brief Shutdown memory module
 */
void mvgal_memory_module_shutdown(void);

/**
 * @brief Allocate buffer internal
 */
mvgal_error_t mvgal_buffer_allocate_internal(
    const mvgal_memory_alloc_info_t *alloc_info,
    struct mvgal_buffer **buffer_out
);

/**
 * @brief Free buffer internal
 */
void mvgal_buffer_free_internal(struct mvgal_buffer *buffer);

/**
 * @brief Increment buffer reference count
 */
void mvgal_buffer_retain(struct mvgal_buffer *buffer);

/**
 * @brief Decrement buffer reference count and possibly free
 */
void mvgal_buffer_release(struct mvgal_buffer *buffer);

/**
 * @brief Replicate buffer to multiple GPUs
 */
mvgal_error_t mvgal_buffer_replicate(
    struct mvgal_buffer *buffer,
    uint32_t gpu_count,
    const uint32_t *gpu_indices,
    mvgal_fence_t fence
);

/**
 * @brief Get GPU address for a buffer
 */
mvgal_error_t mvgal_buffer_get_gpu_address(
    struct mvgal_buffer *buffer,
    uint32_t gpu_index,
    uint64_t *address
);

/**
 * @brief Check if DMA-BUF is supported for the given GPU mask
 */
bool mvgal_dmabuf_is_supported(uint32_t gpu_mask);

/**
 * @brief Allocate memory using DMA-BUF heap
 */
mvgal_error_t mvgal_dmabuf_allocate(
    size_t size,
    size_t alignment,
    mvgal_memory_flags_t flags,
    uint32_t gpu_mask,
    int *fd_out,
    void **ptr_out
);

/**
 * @brief Free DMA-BUF memory
 */
void mvgal_dmabuf_free(int fd, void *ptr, size_t size);

/**
 * @brief Import existing DMA-BUF
 */
mvgal_error_t mvgal_dmabuf_import(
    int fd,
    size_t size,
    mvgal_memory_flags_t flags,
    struct mvgal_buffer **buffer_out
);

/**
 * @brief Export buffer as DMA-BUF
 */
mvgal_error_t mvgal_dmabuf_export(
    struct mvgal_buffer *buffer,
    int *fd_out
);

/**
 * @brief Copy between GPU-visible buffers.
 */
mvgal_error_t mvgal_dmabuf_copy_gpu_to_gpu(
    struct mvgal_buffer *src_buffer,
    uint64_t src_offset,
    struct mvgal_buffer *dst_buffer,
    uint64_t dst_offset,
    size_t size,
    uint32_t src_gpu_index,
    uint32_t dst_gpu_index
);

/**
 * @brief Select the preferred GPU copy method.
 */
mvgal_memory_copy_method_t mvgal_dmabuf_get_copy_method(
    uint32_t src_gpu_index,
    uint32_t dst_gpu_index
);

/**
 * @brief Map DMA-BUF to CPU address space
 */
mvgal_error_t mvgal_dmabuf_map(int fd, size_t size, size_t offset, void **ptr_out);

/**
 * @brief Unmap DMA-BUF from CPU address space
 */
void mvgal_dmabuf_unmap(int fd, void *ptr, size_t size);

/**
 * @brief Signal a fence with completion status.
 */
mvgal_error_t mvgal_fence_signal_with_status(mvgal_fence_t fence, mvgal_error_t status);

/**
 * @brief Bind buffer to a GPU
 */
mvgal_error_t mvgal_buffer_bind_to_gpu(
    struct mvgal_buffer *buffer,
    uint32_t gpu_index,
    bool writable
);

/**
 * @brief Unbind buffer from all GPUs
 */
void mvgal_buffer_unbind(struct mvgal_buffer *buffer);

/* ------------------------------------------------------------------ */
/*  Unified heap helpers (internal)                                   */
/* ------------------------------------------------------------------ */

/**
 * @brief Initialize heap regions from GPU descriptors
 *
 * Called during module init to populate the unified heap view.
 * Maps GPU descriptors to heap regions with VRAM size estimates.
 *
 * @param state Memory state with populated GPU descriptors
 */
void mvgal_heap_init_regions(mvgal_memory_state_t *state);

/**
 * @brief Update free bytes for a specific region
 *
 * Recalculates free bytes for a GPU's VRAM region based on
 * active allocations. Called after alloc/free operations.
 *
 * @param state Memory state
 * @param region_index Index into state's region list (0 = system, 1+ = GPU)
 */
void mvgal_heap_update_free(mvgal_memory_state_t *state, uint32_t region_index);

/**
 * @brief Select best heap region for an allocation
 *
 * Picks the optimal region based on GPU mask, flags, and free space.
 * Returns the region index into the unified heap's regions array.
 *
 * @param state Memory state
 * @param gpu_mask Bitmask of desired GPUs
 * @param flags Allocation flags
 * @return Region index, or UINT32_MAX if no suitable region
 */
uint32_t mvgal_heap_select_best(
    const mvgal_memory_state_t *state,
    uint32_t gpu_mask,
    mvgal_memory_flags_t flags
);

/**
 * @brief Get optimal memory placement for cross-vendor transfer
 *
 * Determines whether to use P2P, DMA-BUF, or host staging
 * based on source and destination GPU vendor and capabilities.
 *
 * @param src_gpu Source GPU index
 * @param dst_gpu Destination GPU index
 * @param size Transfer size
 * @return Recommended copy method
 */
mvgal_memory_copy_method_t mvgal_memory_get_placement_strategy(
    uint32_t src_gpu,
    uint32_t dst_gpu,
    size_t size
);

/**
 * @brief Optimize buffer placement for multi-GPU access
 *
 * Migrates or replicates a buffer to optimize for upcoming access pattern.
 * Used for gaming (replicate frequently accessed buffers) and
 * rendering (keep large buffers on primary GPU).
 *
 * @param buffer Buffer to optimize
 * @param access_pattern Expected access pattern (gaming/rendering/compute)
 * @param gpu_count Number of GPUs that will access
 * @param gpu_indices Array of GPU indices
 * @return MVGAL_SUCCESS on success
 */
mvgal_error_t mvgal_buffer_optimize_placement(
    struct mvgal_buffer *buffer,
    mvgal_workload_type_t access_pattern,
    uint32_t gpu_count,
    const uint32_t *gpu_indices
);

#ifdef __cplusplus
}
#endif

#endif // MVGAL_MEMORY_INTERNAL_H
