#include "platform_bootstrap.h"

#include <stdbool.h>

#include "esp_heap_caps.h"
#include "esp_log.h"
#include "mbedtls/platform.h"
#include "nvs_flash.h"

static bool s_initialized;
static const char *TAG = "platform_bootstrap";

static void *tls_calloc_prefer_psram(size_t n, size_t size) {
    void *ptr = heap_caps_calloc(n, size, MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT);
    if (ptr == NULL) {
        ptr = heap_caps_calloc(n, size, MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT);
    }
    return ptr;
}

static void tls_free(void *ptr) {
    heap_caps_free(ptr);
}

esp_err_t platform_bootstrap_init(void) {
    if (s_initialized) {
        return ESP_OK;
    }

    esp_err_t err = nvs_flash_init();
    if (err == ESP_ERR_NVS_NO_FREE_PAGES || err == ESP_ERR_NVS_NEW_VERSION_FOUND) {
        ESP_ERROR_CHECK(nvs_flash_erase());
        err = nvs_flash_init();
    }
    ESP_ERROR_CHECK(err);

#if CONFIG_SPIRAM_USE_MALLOC
    // Keep very small allocations on internal RAM, push larger generic mallocs to PSRAM.
    heap_caps_malloc_extmem_enable(4096);
    ESP_LOGI(TAG, "extmem malloc threshold set to 4096 bytes");
#endif

#if defined(MBEDTLS_PLATFORM_MEMORY)
    int tls_rc = mbedtls_platform_set_calloc_free(tls_calloc_prefer_psram, tls_free);
    if (tls_rc == 0) {
        ESP_LOGI(TAG, "mbedtls calloc/free redirected to PSRAM-preferred allocator");
    } else {
        ESP_LOGW(TAG, "mbedtls calloc/free redirect failed: rc=%d", tls_rc);
    }
#endif

    s_initialized = true;
    return ESP_OK;
}