/* SPDX-License-Identifier: GPL-3.0-only */
#include "ssh_memory.h"
#include "esp_heap_caps.h"

#include <assert.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

/* All payloads have guarded rounded capacity and normal malloc alignment. */
#define ALIGNMENT _Alignof(max_align_t)
#define GUARD (2U * sizeof(max_align_t))
#define BLOCKS 8U
#define LIMIT 4096U

typedef struct {
    unsigned char *raw;
    unsigned char *base;
    size_t capacity;
    bool internal;
} block_t;

static block_t blocks[BLOCKS];
static bool fail_psram;
static bool fail_internal;
static unsigned allocations, releases, queries, psram_attempts, internal_attempts;
static unsigned live, peak_live;
static size_t last_request;

static block_t *lookup(void *pointer)
{
    assert(pointer != NULL);
    for (size_t i = 0; i < BLOCKS; ++i) {
        if (blocks[i].base == pointer) return &blocks[i];
    }
    assert(!"not a live allocation base");
    abort();
}

static void bytes_are(const unsigned char *p, size_t size, unsigned char value)
{
    for (size_t i = 0; i < size; ++i) assert(p[i] == value);
}

static void guards(const block_t *block)
{
    bytes_are(block->raw, GUARD, 0xD3);
    bytes_are(block->base + block->capacity, GUARD, 0xD3);
}

void *heap_caps_malloc_prefer(size_t size, size_t count, ...)
{
    va_list args;
    va_start(args, count);
    assert(count == 2U);
    assert(va_arg(args, unsigned int) == (MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT));
    assert(va_arg(args, unsigned int) == (MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT));
    va_end(args);
    ++allocations;
    last_request = size;
    /* The module must pass size unchanged, including zero and SIZE_MAX. */
    if (size == 0U) return NULL;
    ++psram_attempts;
    bool internal = fail_psram || size > LIMIT;
    if (internal) {
        ++internal_attempts;
        if (fail_internal || size > LIMIT) return NULL;
    }
    size_t capacity = ((size + ALIGNMENT - 1U) / ALIGNMENT) * ALIGNMENT;
    for (size_t i = 0; i < BLOCKS; ++i) {
        block_t *block = &blocks[i];
        if (block->base != NULL) continue;
        block->raw = malloc(GUARD + capacity + GUARD);
        assert(block->raw != NULL);
        block->base = block->raw + GUARD;
        block->capacity = capacity;
        block->internal = internal;
        memset(block->raw, 0xD3, GUARD + capacity + GUARD);
        memset(block->base, 0xA5, capacity);
        assert((uintptr_t)block->base % ALIGNMENT == 0U);
        ++live;
        if (live > peak_live) peak_live = live;
        return block->base;
    }
    assert(!"fake heap exhausted");
    return NULL;
}

size_t heap_caps_get_allocated_size(void *pointer)
{
    ++queries;
    block_t *block = lookup(pointer);
    guards(block);
    return block->capacity;
}

void heap_caps_free(void *pointer)
{
    block_t *block = lookup(pointer);
    guards(block);
    /* Inspect BEFORE real free: no reads through dangling pointers. */
    bytes_are(block->base, block->capacity, 0);
    free(block->raw);
    memset(block, 0, sizeof(*block));
    ++releases;
    --live;
}

static void test_null_zero(void)
{
    unsigned before = queries;
    ssh_memory_free(NULL);
    assert(queries == before && releases == 0U);
    assert(ssh_memory_malloc(0) == NULL && last_request == 0U);
    unsigned calls = allocations;
    assert(ssh_memory_realloc(NULL, 0) == NULL);
    assert(allocations == calls + 1U && queries == before);
    void *p = ssh_memory_realloc(NULL, 7);
    assert(p != NULL && last_request == 7U);
    assert(ssh_memory_realloc(p, 0) == NULL && live == 0U);
}

static void test_retained_capacity(void)
{
    unsigned char *p = ssh_memory_malloc(17);
    block_t *block = lookup(p);
    size_t capacity = block->capacity;
    assert(capacity > 17U);
    memset(p, 0x71, capacity);
    unsigned calls = allocations;
    assert(ssh_memory_realloc(p, capacity) == p);
    bytes_are(p, capacity, 0x71);
    assert(ssh_memory_realloc(p, 17) == p);
    bytes_are(p, 17, 0x71);
    bytes_are(p + 17, capacity - 17, 0);
    assert(ssh_memory_realloc(p, 5) == p);
    bytes_are(p, 5, 0x71);
    bytes_are(p + 5, capacity - 5, 0);
    /* Logical regrowth within retained capacity allocates nothing. */
    assert(ssh_memory_realloc(p, capacity - 1U) == p);
    bytes_are(p, 5, 0x71);
    bytes_are(p + 5, capacity - 5, 0);
    assert(block->capacity == capacity && allocations == calls);
    guards(block);
    ssh_memory_free(p);
}

static void test_growth_and_failure(void)
{
    unsigned char *p = ssh_memory_malloc(17);
    size_t capacity = lookup(p)->capacity;
    for (size_t i = 0; i < capacity; ++i) p[i] = (unsigned char)(i + 1U);
    fail_psram = fail_internal = true;
    unsigned freed = releases;
    assert(ssh_memory_realloc(p, capacity + 1U) == NULL);
    assert(releases == freed && live == 1U);
    for (size_t i = 0; i < capacity; ++i) assert(p[i] == (unsigned char)(i + 1U));
    guards(lookup(p));
    fail_internal = false;
    unsigned char *q = ssh_memory_realloc(p, capacity + 1U);
    assert(q != NULL && lookup(q)->internal);
    assert(last_request == capacity + 1U && releases == freed + 1U);
    assert(live == 1U && peak_live == 2U);
    for (size_t i = 0; i < capacity; ++i) assert(q[i] == (unsigned char)(i + 1U));
    /* The new suffix isn't promised zero; ensure no over-copy either. */
    bytes_are(q + capacity, lookup(q)->capacity - capacity, 0xA5);
    fail_psram = false;
    size_t old_capacity = lookup(q)->capacity;
    memset(q, 0x69, old_capacity);
    unsigned char *r = ssh_memory_realloc(q, old_capacity + 19U);
    assert(r != NULL && !lookup(r)->internal);
    bytes_are(r, old_capacity, 0x69);
    ssh_memory_free(r);
}

static void test_sizes_alignment_and_preference(void)
{
    for (size_t size = 1; size <= 129; ++size) {
        fail_psram = (size % 2U) != 0U;
        unsigned external_before = psram_attempts;
        unsigned internal_before = internal_attempts;
        void *p = ssh_memory_malloc(size);
        assert(last_request == size && lookup(p)->internal == fail_psram);
        assert(psram_attempts == external_before + 1U);
        assert(internal_attempts == internal_before + (fail_psram ? 1U : 0U));
        assert((uintptr_t)p % ALIGNMENT == 0U);
        ssh_memory_free(p);
    }
    fail_psram = fail_internal = true;
    assert(ssh_memory_malloc(33) == NULL);
    fail_psram = fail_internal = false;
    assert(ssh_memory_malloc(SIZE_MAX) == NULL && last_request == SIZE_MAX);
    unsigned char *p = ssh_memory_malloc(9);
    size_t capacity = lookup(p)->capacity;
    memset(p, 0x81, capacity);
    assert(ssh_memory_realloc(p, SIZE_MAX) == NULL && last_request == SIZE_MAX);
    bytes_are(p, capacity, 0x81);
    guards(lookup(p));
    ssh_memory_free(p);
}

int main(void)
{
    test_null_zero();
    test_retained_capacity();
    test_growth_and_failure();
    test_sizes_alignment_and_preference();
    assert(live == 0U);
    puts("PASS ssh_memory: null/zero, rounded extent, retained shrink/equal, growth, failure, caps, alignment, guards");
    return 0;
}