// SPDX-License-Identifier: MPL-2.0
#include "memory/heap.h"
#include "kernel/log.h"
#include "memory/vmem_layout.h"
static sh_heap_KERNEL_HEAP *default_heap=SH_NULLPTR;
void sh_heap_load_default_heap(sh_heap_KERNEL_HEAP *heap) {
  default_heap=heap;
}
sh_heap_KERNEL_HEAP *sh_heap_get_default_heap() {
  return default_heap;
}
SH_STATUS sh_heap_init_heap(sh_pez_PHYSICAL_PLANE *phys_plane,sh_pez_VIRTUAL_PLANE *virt_plane,sh_page_PAGE_TABLE_POOL *kernel_ptp,sh_heap_KERNEL_HEAP *kernel_heap) {
  if (phys_plane==SH_NULLPTR || virt_plane==SH_NULLPTR || kernel_ptp==SH_NULLPTR || kernel_heap==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  kernel_heap->phys_plane=phys_plane;
  kernel_heap->virt_plane=virt_plane;
  kernel_heap->kernel_ptp=kernel_ptp;
  SH_STATUS status=sh_radix_tree_init(phys_plane->slab_radix_node,kernel_ptp,&kernel_heap->alloc_size_tree,8);
  if (sh_status_error(status)) return status;
  return SH_STATUS_SUCCESS;
}
static void sh_heap_internal_crash(const char *str,SH_STATUS status) {
  sh_log_ffatal(SH_LOG_SOURCE_HEAP,str,status);
  sh_log_mem_stats(SH_LOG_SOURCE_HEAP);
  while (SH_TRUE) {};
  return;
}
SH_STATUS sh_heap_allocate_pages(sh_uint32 pages_count,sh_page_VIRTUAL_ADDRESS *address) {
  if (pages_count==0 || address==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  if (default_heap==SH_NULLPTR) return SH_STATUS_HEAP_NOT_INITIALIZED;
  sh_page_PHYSICAL_ADDRESS pa;
  SH_STATUS status=sh_pez_alloc_physical_pages(default_heap->phys_plane,pages_count,&pa);
  if (sh_status_error(status)) return status;
  sh_page_PHYSICAL_ADDRESS va;
  status=sh_pez_alloc_virtual_pages(default_heap->virt_plane,pages_count,&va);
  if (sh_status_error(status)) {
    status=sh_pez_free_physical_pages(default_heap->phys_plane,&pa,pages_count);
    if (sh_status_error(status)) sh_heap_internal_crash("Heap alloc error: couldn't free in emergency physical pages after virtual alloc error. Status: %8s\n",status);
    return status;
  }
  if (pa%SH_PAGE_SIZE!=0 || va%SH_PAGE_SIZE!=0) sh_heap_internal_crash("Heap alloc error: physical address or virtual address isn't page aligned. Status: %8s\n",SH_STATUS_PEZ_CORRUPTED);
  status=sh_page_map_contiguous_pages_range_ptp(default_heap->kernel_ptp,va,pa,SH_PAGE_PRESENT | SH_PAGE_RW | SH_PAGE_NX,pages_count*SH_PAGE_SIZE);
  if (status!=SH_STATUS_SUCCESS) sh_heap_internal_crash("Heap alloc error: couldn't map all pages. Status: %8s",status);
  sh_uint32 page_index=(sh_uint32)((va/SH_PAGE_SIZE)-SH_VMEM_LAYOUT_HEAP_BIG_VA);
  status=sh_radix_tree_insert_value(default_heap->phys_plane->slab_radix_node,default_heap->kernel_ptp,&default_heap->alloc_size_tree,page_index,pages_count);
  if (status!=SH_STATUS_SUCCESS) sh_heap_internal_crash("Heap alloc error: couldn't insert alloc record inside alloc_size_tree. Status: %8s\n",status);
  *address=va;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_heap_free_pages(sh_page_VIRTUAL_ADDRESS va) {
  if (default_heap==SH_NULLPTR) return SH_STATUS_HEAP_NOT_INITIALIZED;
  if (va%SH_PAGE_SIZE!=0) return SH_STATUS_INVALID_PARAMETER;
  sh_uint32 page_index=(sh_uint32)((va/SH_PAGE_SIZE)-SH_VMEM_LAYOUT_HEAP_BIG_VA);
  sh_page_VIRTUAL_ADDRESS pages_count_raw=0;
  SH_STATUS status=sh_radix_tree_get_value(&default_heap->alloc_size_tree,page_index,&pages_count_raw);
  if (status!=SH_STATUS_SUCCESS) {
    // sh_log_flog(SH_LOG_SOURCE_HEAP,"va=0x%x\n",va);
    sh_heap_internal_crash("Heap free error: alloc record not found in alloc_size_tree. Status: %8s\n",status);
  }
  sh_uint32 pages_count=(sh_uint32)pages_count_raw;
  sh_page_PHYSICAL_ADDRESS pa;
  status=sh_page_ptp_va_to_pa(default_heap->kernel_ptp,va,&pa);
  if (sh_status_error(status)) sh_heap_internal_crash("Heap free error: couldn't find associated physical address. Status: %8s\n",status);
  status=sh_page_unmap_contiguous_pages_range_ptp(default_heap->kernel_ptp,va,pages_count*SH_PAGE_SIZE);
  if (sh_status_error(status)) sh_heap_internal_crash("Heap free error: unmap failed. Status: %8s\n",status);
  status=sh_pez_free_virtual_pages(default_heap->virt_plane,&va,pages_count);
  if (sh_status_error(status)) sh_heap_internal_crash("Heap free error: virt_plane free failed. Status: %8s\n",status);
  status=sh_pez_free_physical_pages(default_heap->phys_plane,&pa,pages_count);
  if (sh_status_error(status)) sh_heap_internal_crash("Heap free error: phys_plane free failed. Status: %8s\n",status);
  status=sh_radix_tree_delete_value(default_heap->phys_plane->slab_radix_node,&default_heap->alloc_size_tree,page_index);
  if (sh_status_error(status)) sh_heap_internal_crash("Heap free error: record deletion failed. Status: %8s\n",status);
  return SH_STATUS_SUCCESS;
}
static inline sh_uint32 sh_heap_internal_get_index(sh_uint32 n) {
  if (n<=8) return 0;
  return (sh_uint32)31-(sh_uint32)__builtin_clz(n-1)-(sh_uint32)2;
}
SH_STATUS sh_heap_allocate_object(sh_uint32 size_bytes,sh_page_VIRTUAL_ADDRESS *address) {
  if (default_heap==SH_NULLPTR) return SH_STATUS_HEAP_NOT_INITIALIZED;
  if (size_bytes==0 || address==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  if (size_bytes>1024) return SH_STATUS_INVALID_PARAMETER;
  sh_uint32 slab_alloc_index=sh_heap_internal_get_index(size_bytes);
  void *ptr;
  SH_STATUS status=sh_slab_generic_alloc(&default_heap->slabs_allocator[slab_alloc_index],default_heap->kernel_ptp,&ptr);
  if (sh_status_error(status)) return status;
  sh_page_VIRTUAL_ADDRESS va=(sh_uint64)ptr;
  *address=va;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_heap_free_object(sh_page_VIRTUAL_ADDRESS va) {
  if (default_heap==SH_NULLPTR) return SH_STATUS_HEAP_NOT_INITIALIZED;
  sh_uint64 slab_alloc_index=(va-SH_VMEM_LAYOUT_HEAP_SLAB_LEVEL_0_VA)/0x00000C0000000000;
  if (slab_alloc_index>7) return SH_STATUS_INVALID_PARAMETER;
  return sh_slab_generic_dealloc(&default_heap->slabs_allocator[slab_alloc_index],(void*)va);
}
