#include "heap.h"
#include "FreeRTOS.h"
#include "task.h"
#include "tlsf.h"
#include "elog.h"
#include <sys/lock.h>
#include <string.h>
#include "stm32_lock.h"

// static const char *TAG = "system_memory";
static uint8_t __attribute__((section(".ram_dtcm"))) heap_dtcm[HEAP_DTCM_SIZE_KB * 1024];
static uint8_t __attribute__((section(".ram_axi"))) heap_axi[HEAP_AXI_SIZE_KB * 1024];
static uint8_t __attribute__((section(".ram_ahb"))) heap_ahb[HEAP_APB_SIZE_KB * 1024];
static tlsf_t tlsf_heap_dtcm;
static tlsf_t tlsf_heap_axi;
static tlsf_t tlsf_heap_ahb;
static LockingData_t malloc_lock_dtcm;
static LockingData_t malloc_lock_axi;
static LockingData_t malloc_lock_ahb;

static void tlsf_mem_walker(void* ptr, size_t size, int used, void* user)
{
  tlsf_pool_statistics *statistics = (tlsf_pool_statistics *)user;

  if (used) {
    ++ statistics->used_cnt;
    statistics->used_size += size;
    if(size > statistics->used_biggest_size) {
      statistics->used_biggest_size = size;
    }
  } else {
    ++ statistics->free_cnt;
    statistics->free_size += size;
    if(size > statistics->free_biggest_size) {
      statistics->free_biggest_size = size;
    }
  }
}

void get_tlsf_pool_statistics(tlsf_t tlsf_heap, tlsf_pool_statistics *statistics_out)
{
  memset(statistics_out, 0, sizeof(tlsf_pool_statistics));
  tlsf_walk_pool(tlsf_get_pool(tlsf_heap), tlsf_mem_walker, statistics_out);
}

void *malloc_dtcm(size_t size)
{
  stm32_lock_acquire(&malloc_lock_dtcm);
  void *p = tlsf_malloc(tlsf_heap_dtcm, size);
  stm32_lock_release(&malloc_lock_dtcm);

  // if (xTaskGetSchedulerState() == pdTRUE) {
  //   elog_verbose(TAG, "dtcm: allocating %d bytes memory @ %p", size, p);
  // }
  return p;
}

void free_dtcm(void *p)
{
  if (p) {
    stm32_lock_acquire(&malloc_lock_dtcm);
    tlsf_free(tlsf_heap_dtcm, p);
    stm32_lock_release(&malloc_lock_dtcm);
  }
}

void *realloc_dtcm(void *oldmem, size_t bytes)
{
  if (oldmem == NULL) { //当旧内存地址为空时，此函数等于malloc函数
    return malloc_dtcm(bytes);
  }

  void *p;
  stm32_lock_acquire(&malloc_lock_dtcm);
  p = tlsf_realloc(tlsf_heap_dtcm, oldmem, bytes);
  stm32_lock_release(&malloc_lock_dtcm);

  return p;
}

tlsf_pool_statistics *get_dtcm_heap_statistics(void)
{
  static tlsf_pool_statistics statistics;

  stm32_lock_acquire(&malloc_lock_dtcm);
  get_tlsf_pool_statistics(tlsf_heap_dtcm, &statistics);
  stm32_lock_release(&malloc_lock_dtcm);

  return &statistics;
}

void *malloc_ahb(size_t size)
{
  void *p;

  stm32_lock_acquire(&malloc_lock_ahb);
  p = tlsf_malloc(tlsf_heap_ahb, size);
  stm32_lock_release(&malloc_lock_ahb);

  return p;
}

void free_ahb(void *p)
{
  if (p) {
    stm32_lock_acquire(&malloc_lock_ahb);
    tlsf_free(tlsf_heap_ahb, p);
    stm32_lock_release(&malloc_lock_ahb);
  }
}

void *realloc_ahb(void *oldmem, size_t bytes)
{
  if (oldmem == NULL) { //当旧内存地址为空时，此函数等于malloc函数
    return malloc_ahb(bytes);
  }

  void *p;
  stm32_lock_acquire(&malloc_lock_ahb);
  p = tlsf_realloc(tlsf_heap_ahb, oldmem, bytes);
  stm32_lock_release(&malloc_lock_ahb);

  return p;
}

tlsf_pool_statistics *get_ahb_heap_statistics(void)
{
  static tlsf_pool_statistics statistics;

  stm32_lock_acquire(&malloc_lock_dtcm);
  get_tlsf_pool_statistics(tlsf_heap_ahb, &statistics);
  stm32_lock_release(&malloc_lock_dtcm);

  return &statistics;
}

void *__wrap__malloc_r(struct _reent *ptr, size_t bytes)
{
  UNUSED(ptr); //不使用newlib自带的基于重入的锁

  stm32_lock_acquire(&malloc_lock_axi);
  void *p = tlsf_malloc(tlsf_heap_axi, bytes);
  stm32_lock_release(&malloc_lock_axi);

  // elog_verbose(TAG, "axi: allocating %d bytes memory @ %p", bytes, p);

  return p;
}

void __wrap__free_r(struct _reent *ptr, void *mem)
{
  UNUSED(ptr); //不使用newlib自带的基于重入的锁

  if (mem) {
    stm32_lock_acquire(&malloc_lock_axi);
    tlsf_free(tlsf_heap_axi, mem);
    stm32_lock_release(&malloc_lock_axi);
  }
}

void *__wrap__realloc_r(struct _reent *ptr, void *oldmem, size_t bytes)
{
  UNUSED(ptr); //不使用newlib自带的基于重入的锁

  if (oldmem == NULL) { //当旧内存地址为空时，此函数等于malloc函数
    return __wrap__malloc_r(ptr, bytes);
  }

  stm32_lock_acquire(&malloc_lock_axi);
  void *p = tlsf_realloc(tlsf_heap_axi, oldmem, bytes);
  stm32_lock_release(&malloc_lock_axi);

  return p;
}

tlsf_pool_statistics *get_axi_heap_statistics(void)
{
  static tlsf_pool_statistics statistics;

  stm32_lock_acquire(&malloc_lock_dtcm);
  get_tlsf_pool_statistics(tlsf_heap_axi, &statistics);
  stm32_lock_release(&malloc_lock_dtcm);

  return &statistics;
}

void heap_init(void)
{
  stm32_lock_init(&malloc_lock_dtcm);
  stm32_lock_init(&malloc_lock_axi);
  stm32_lock_init(&malloc_lock_ahb);
  tlsf_heap_dtcm = tlsf_create_with_pool(heap_dtcm, sizeof(heap_dtcm));
  tlsf_heap_axi = tlsf_create_with_pool(heap_axi, sizeof(heap_axi));
  tlsf_heap_ahb = tlsf_create_with_pool(heap_ahb, sizeof(heap_ahb));
}