#include <rtthread.h>
#include <rthw.h>

/**
 * @brief Allocate aligned memory from a memheap
 *
 * @param heap    Pointer to the memheap object
 * @param size    Requested user memory size (in bytes)
 * @param align   Alignment boundary (must be power of 2, e.g., 4, 8, 16)
 * @return void*  Aligned pointer on success, RT_NULL on failure
 *
 * @note The returned pointer must be freed with rt_memheap_free_align()
 */
void *rt_memheap_alloc_align(struct rt_memheap *heap, rt_size_t size, rt_size_t align)
{
    void *raw_ptr;
    void *aligned_ptr;
    uintptr_t *header; // stores raw_ptr for free

    RT_ASSERT(heap != RT_NULL);
    RT_ASSERT(size > 0);
    RT_ASSERT(align > 0 && (align & (align - 1)) == 0); // check power of 2

    // Minimum space: user data + alignment padding + hidden header (sizeof(uintptr_t))
    rt_size_t total_size = size + align - 1 + sizeof(uintptr_t);

    raw_ptr = rt_memheap_alloc(heap, total_size);
    if (raw_ptr == RT_NULL)
        return RT_NULL;

    // Calculate aligned address after hidden header
    uintptr_t addr = (uintptr_t)raw_ptr + sizeof(uintptr_t);
    aligned_ptr = (void *)((addr + align - 1) & ~(align - 1));

    // Store raw_ptr just before aligned_ptr (in the hidden header)
    header = (uintptr_t *)aligned_ptr - 1;
    *header = (uintptr_t)raw_ptr;

    return aligned_ptr;
}

/**
 * @brief Free memory allocated by rt_memheap_alloc_align()
 *
 * @param heap   Pointer to the memheap object
 * @param ptr    Pointer returned by rt_memheap_alloc_align()
 */
void rt_memheap_free_align(struct rt_memheap *heap, void *ptr)
{
    if (ptr == RT_NULL)
        return;

    // Retrieve raw_ptr from hidden header
    uintptr_t *header = (uintptr_t *)ptr - 1;
    void *raw_ptr = (void *)(*header);

    rt_memheap_free(raw_ptr);
}