// SPDX-License-Identifier: MPL-2.0
#include "memory/memory.h"
#include "kernel/log.h"
#include "kernel/test.h"
#include "kernel/efifb.h"
sh_memory_MEMORY_SERVICES *memory_services=SH_NULLPTR;
SH_STATUS sh_memory_init_subsystem(sh_conf_BOOT_CONFIG *boot_config,sh_memory_MEMORY_SERVICES *services) {
  if (boot_config==SH_NULLPTR || services==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_log_log("Trying to obtain memory map...",SH_LOG_SOURCE_MEMS);
  SH_STATUS status=sh_page_copy_memory_map();
  if (status==SH_STATUS_SUCCESS) {
    sh_log_log("Successfully copied memory map.",SH_LOG_SOURCE_PAGE);
    sh_log_log("Successfully obtained memory map.",SH_LOG_SOURCE_MEMS);
  } else {
    sh_log_error("Can't copy memory map.",SH_LOG_SOURCE_PAGE);
    sh_log_fatal("Can't obtain memory map.",SH_LOG_SOURCE_MEMS);
    while (SH_TRUE) {};
  }
  sh_log_log("Checking memory map...",SH_LOG_SOURCE_MEMS);
  if (status==SH_STATUS_SUCCESS) {
    sh_log_log("Successfully checked memory map.",SH_LOG_SOURCE_MEMS);
  } else if (status==SH_STATUS_INVALID_SIGNATURE) {
    sh_log_fatal("Can't check memory map.",SH_LOG_SOURCE_MEMS);
    while (SH_TRUE) {};
  } else if (status==SH_STATUS_MMAP_BUFFER_OVERFLOW) {
    sh_log_fatal("Memory map is overflowing. Can't safely continue.",SH_LOG_SOURCE_MEMS);
    while (SH_TRUE) {};
  }
  sh_log_log("Initializing current page table pool...",SH_LOG_SOURCE_MEMS);
  status=sh_page_init_ptp(boot_config->page_table_pool_pa,sh_page_get_boot_ptp_va(),boot_config->page_table_allocator_level,&services->kernel_ptp);
  if (status==SH_STATUS_PT_POOL_NO_BITMAP_INIT) {
    sh_log_fatal("Couldn't initialize current page table pool.",SH_LOG_SOURCE_MEMS);
    while (SH_TRUE) {};
  } else if (status==SH_STATUS_PT_POOL_NO_PAGE_SET) {
    sh_log_fatal("Couldn't initialize current page table pool.",SH_LOG_SOURCE_MEMS);
    while (SH_TRUE) {};
  }
  sh_log_log("Successfully initialized current page table pool.",SH_LOG_SOURCE_MEMS);
  status=sh_efifb_init_fb(&boot_config->fb_config,&services->kernel_ptp);
  if (sh_status_error(status)) {
    sh_log_ffatal(SH_LOG_SOURCE_MEMS,"Couldn't initialize EFI FB subsystem. Status: %8s\n",status);
    while (SH_TRUE) {};
  }
  sh_efifb_init_bar();
  sh_log_log("Successfully initialized EFI FB subsystem.",SH_LOG_SOURCE_MEMS);
  sh_log_log("Parsing memory map...",SH_LOG_SOURCE_MEMS);
  status=sh_page_analyse_memory_map(&services->kernel_ptp);
  if (sh_status_error(status)) {
    sh_log_ffatal(SH_LOG_SOURCE_MEMS,"Couldn't parse memory map. Error code: %8s\n",status);
    while (SH_TRUE) {};
  }
  sh_log_log("Successfully parsed memory map.",SH_LOG_SOURCE_MEMS);
  sh_log_mem_stats(SH_LOG_SOURCE_MEMS);
  sh_log_log("Page managment subsystem successfully initialized",SH_LOG_SOURCE_MEMS);
  status=sh_slab_reg_phys_alloc_init(&services->slab_reg_phys,&services->kernel_ptp);
  if (sh_status_error(status)) {
    sh_log_fatal("Couldn't initialize physical region object slab allocator.",SH_LOG_SOURCE_MEMS);
    while (SH_TRUE) {};
  }
  sh_log_log("Successfully initialize physical region object slab allocator.",SH_LOG_SOURCE_MEMS);
  sh_log_log("Allocating first slab...",SH_LOG_SOURCE_MEMS);
  sh_slab_reg_phys_SLAB_STRUCT *first_slab_rp=SH_NULLPTR;
  status=sh_slab_reg_phys_add_slab(&services->slab_reg_phys,&services->kernel_ptp,&first_slab_rp);
  if (status!=SH_STATUS_SUCCESS) {
    sh_log_ffatal(SH_LOG_SOURCE_MEMS,"Couldn't allocate first slab. Error code: %8s\n",status);
    while (SH_TRUE) {};
  } else {
    sh_log_log("Successfully allocated first slab. Verifying data...",SH_LOG_SOURCE_MEMS);
    if (first_slab_rp->free_bitmap[0]!=1 ||
        first_slab_rp->slab_index!=0 ||
        first_slab_rp->next!=SH_NULLPTR ||
        first_slab_rp->prev!=SH_NULLPTR ||
        first_slab_rp->used_count!=0) {
          sh_log_error("One of the attributes inside first slab is wrong.",SH_LOG_SOURCE_MEMS);
          sh_log_fatal("Stopping for safety.",SH_LOG_SOURCE_MEMS);
          while (SH_TRUE) {};
        } else {
          sh_log_log("All attributes are good.",SH_LOG_SOURCE_MEMS);
        }
  }
  status=sh_slab_reg_virt_alloc_init(&services->slab_reg_virt,&services->kernel_ptp);
  if (sh_status_error(status)) {
    sh_log_fatal("Couldn't initialize virtual region object slab allocator.",SH_LOG_SOURCE_MEMS);
    while (SH_TRUE) {};
  }
  sh_log_log("Successfully initialize virtual region object slab allocator.",SH_LOG_SOURCE_MEMS);
  sh_log_log("Allocating first slab...",SH_LOG_SOURCE_MEMS);
  sh_slab_reg_virt_SLAB_STRUCT *first_slab_rv=SH_NULLPTR;
  status=sh_slab_reg_virt_add_slab(&services->slab_reg_virt,&services->kernel_ptp,&first_slab_rv);
  if (status!=SH_STATUS_SUCCESS) {
    sh_log_ffatal(SH_LOG_SOURCE_MEMS,"Couldn't allocate first slab. Error code: %8s\n",status);
    while (SH_TRUE) {};
  } else {
    sh_log_log("Successfully allocated first slab. Verifying data...",SH_LOG_SOURCE_MEMS);
    if (first_slab_rv->free_bitmap[0]!=1 ||
        first_slab_rv->slab_index!=0 ||
        first_slab_rv->next!=SH_NULLPTR ||
        first_slab_rv->prev!=SH_NULLPTR ||
        first_slab_rv->used_count!=0) {
          sh_log_error("One of the attributes inside first slab is wrong.",SH_LOG_SOURCE_MEMS);
          sh_log_fatal("Stopping for safety.",SH_LOG_SOURCE_MEMS);
          while (SH_TRUE) {};
        } else {
          sh_log_log("All attributes are good.",SH_LOG_SOURCE_MEMS);
        }
  }
  sh_pba_PAGE_BLOCK_ALLOCATOR slab_radix_node_pba;
  status=sh_pba_init(&slab_radix_node_pba,SH_VMEM_LAYOUT_SLAB_RADIX_NODE_VA,SH_VMEM_LAYOUT_SLAB_RADIX_NODE_SIZE_BYTES/SH_PAGE_SIZE,SH_SLAB_RADIX_NODE_SLAB_DATA_PAGES);
  if (sh_status_error(status)) {
    sh_log_fatal("Couldn't initialize PBA for radix nodes slab allocator.",SH_LOG_SOURCE_MEMS);
    while (SH_TRUE) {};
  }
  sh_log_log("Successfully initialize PBA for radix nodes slab allocator.",SH_LOG_SOURCE_MEMS);
  status=sh_slab_radix_node_alloc_init(&services->slab_radix_node,&slab_radix_node_pba);
  if (sh_status_error(status)) {
    sh_log_fatal("Couldn't initialize radix nodes slab allocator.",SH_LOG_SOURCE_MEMS);
    while (SH_TRUE) {};
  }
  sh_log_log("Successfully initialize radix nodes slab allocator.",SH_LOG_SOURCE_MEMS);
  sh_log_log("Allocating first slab...",SH_LOG_SOURCE_MEMS);
  sh_slab_radix_node_SLAB *first_slab_rn=SH_NULLPTR;
  status=sh_slab_radix_node_add_slab(&services->slab_radix_node,&services->kernel_ptp,&first_slab_rn);
  if (status!=SH_STATUS_SUCCESS) {
    sh_log_ffatal(SH_LOG_SOURCE_MEMS,"Couldn't allocate first slab. Error code: %8s\n",status);
    while (SH_TRUE) {};
  } else {
    sh_log_log("Successfully allocated first slab. Verifying data...",SH_LOG_SOURCE_MEMS);
    if (first_slab_rn->free_bitmap[0]!=0x3FFFF ||
        first_slab_rn->slab_index!=0 ||
        first_slab_rn->next_slab!=SH_NULLPTR ||
        first_slab_rn->prev_slab!=SH_NULLPTR ||
        first_slab_rn->next_partial!=SH_NULLPTR ||
        first_slab_rn->prev_partial!=SH_NULLPTR ||
        first_slab_rn->used_count!=0 ) {
          sh_log_error("One of the attributes inside first slab is wrong.",SH_LOG_SOURCE_MEMS);
          sh_log_fatal("Stopping for safety.",SH_LOG_SOURCE_MEMS);
          while (SH_TRUE) {};
        } else {
          sh_log_log("All attributes are good.",SH_LOG_SOURCE_MEMS);
        }
  }
  if (boot_config->test_benchmark) {
    status=sh_test_slabs_benchmark(&services->slab_reg_phys,&services->slab_reg_virt,&services->slab_radix_node,&services->kernel_ptp);
    if (status!=SH_STATUS_SUCCESS) {
      sh_log_fatal("Couldn't benchmark one of the slab allocator.",SH_LOG_SOURCE_MEMS);
      while (SH_TRUE) {};
    }
    status=sh_test_radix_benchmark(&services->slab_radix_node,&services->kernel_ptp);
    if (status!=SH_STATUS_SUCCESS) {
      sh_log_fatal("Couldn't benchmark radix trees subsystem.",SH_LOG_SOURCE_MEMS);
      while (SH_TRUE) {};
    }
  }
  sh_log_log("Creating Pez physical plane...",SH_LOG_SOURCE_MEMS);
  status=sh_pez_init_physical_plane((sh_uint8*)sh_page_get_physical_bitmap_ptr(),sh_page_get_physical_memory_amount_pages(),&services->slab_reg_phys,&services->slab_radix_node,&services->kernel_ptp,&services->physical_plane);
  if (status!=SH_STATUS_SUCCESS) {
    sh_log_ffatal(SH_LOG_SOURCE_MEMS,"Couldn't create Pez physical plane. Error code: %8s\n",status);
    while (SH_TRUE) {};
  }
  sh_log_log("Succesfully created Pez physical plane.",SH_LOG_SOURCE_MEMS);
  if (boot_config->test_benchmark) {
    status=sh_test_pez_benchmark_physical(&services->physical_plane);
    if (status!=SH_STATUS_SUCCESS) {
      sh_log_fatal("Couldn't benchmark Pez physical plane.",SH_LOG_SOURCE_MEMS);
      while (SH_TRUE) {};
    }
  }
  sh_log_log("Creating Pez virtual plane for kernel heap...",SH_LOG_SOURCE_MEMS);
  status=sh_pez_init_virtual_plane(SH_VMEM_LAYOUT_HEAP_BIG_VA,&services->slab_reg_virt,&services->slab_radix_node,&services->kernel_ptp,&services->kernel_ptp,&services->virtual_plane_kernel_heap);
  if (status!=SH_STATUS_SUCCESS) {
    sh_log_ffatal(SH_LOG_SOURCE_MEMS,"Couldn't create Pez virtual plane for kernel heap. Error code: %8s\n",status);
    while (SH_TRUE) {};
  }
  sh_log_log("Succesfully created Pez virtual plane for kernel heap.",SH_LOG_SOURCE_MEMS);
  sh_pez_set_available();
  status=sh_heap_init_heap(&services->physical_plane,&services->virtual_plane_kernel_heap,&services->kernel_ptp,&services->kernel_heap);
  if (sh_status_error(status)) {
    sh_log_ffatal(SH_LOG_SOURCE_MEMS,"Couldn't initialize kernel heap. Error code: %8s",status);
    while (SH_TRUE) {};
  }
  sh_uint64 slab_data_pages[8]=SH_SLAB_GENERIC_SLAB_DATA_PAGES;
  for (sh_iter64 i=0;i<8;i++) {
    status=sh_pba_init(&services->kernel_heap.pba[i],SH_VMEM_LAYOUT_HEAP_SLAB_LEVEL_0_VA+i*SH_VMEM_LAYOUT_HEAP_SLAB_SPACING,((SH_VMEM_LAYOUT_HEAP_SLAB_LEVEL_0_SIZE_BYTES/SH_PAGE_SIZE)/slab_data_pages[i])*slab_data_pages[i],slab_data_pages[i]);
    if (sh_status_error(status)) {
      sh_log_ffatal(SH_LOG_SOURCE_MEMS,"Couldn't create PBA for heap allocations level %8u. Error code: %8s\n",i,status);
      while (SH_TRUE);
    }
    status=sh_slab_generic_alloc_init((sh_uint8)i,&services->kernel_heap.slabs_allocator[i],&services->kernel_heap.pba[i]);
    if (sh_status_error(status)) {
      sh_log_ffatal(SH_LOG_SOURCE_MEMS,"Couldn't create slab allocator for heap allocations level %8u. Error code: %8s\n",i,status);
      while (SH_TRUE);
    }
    sh_slab_generic_SLAB *first_slab_generic;
    status=sh_slab_generic_add_slab(&services->kernel_heap.slabs_allocator[i],&services->kernel_ptp,&first_slab_generic);
    if (status!=SH_STATUS_SUCCESS) {
      sh_log_ffatal(SH_LOG_SOURCE_MEMS,"Couldn't allocate first slab for heap allocations level %8u. Error code: %8s\n",i,status);
      while (SH_TRUE) {};
    } else {
      if (first_slab_generic->free_bitmap[0]!=services->kernel_heap.slabs_allocator[i].bitmap_init ||
          first_slab_generic->slab_index!=0 ||
          first_slab_generic->next_slab!=SH_NULLPTR ||
          first_slab_generic->prev_slab!=SH_NULLPTR ||
          first_slab_generic->next_partial!=SH_NULLPTR ||
          first_slab_generic->prev_partial!=SH_NULLPTR ||
          first_slab_generic->used_count!=0 ) {
            sh_log_error("One of the attributes inside first slab is wrong.",SH_LOG_SOURCE_MEMS);
            sh_log_fatal("Stopping for safety.",SH_LOG_SOURCE_MEMS);
            while (SH_TRUE) {};
          }
    }
    sh_log_flog(SH_LOG_SOURCE_MEMS,"Succesfully created heap slab allocator with level %8u.\n",i);
  }
  sh_heap_load_default_heap(&services->kernel_heap);
  sh_log_log("Succesfully created kernel heap.",SH_LOG_SOURCE_MEMS);
  if (boot_config->test_benchmark) {
    status=sh_test_malloc_benchmark();
    if (status!=SH_STATUS_SUCCESS) {
      sh_log_fatal("Couldn't benchmark Pez physical plane.",SH_LOG_SOURCE_MEMS);
      while (SH_TRUE) {};
    }
  }
  memory_services=services;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_memory_identity_map(sh_page_PHYSICAL_ADDRESS phys_start,sh_uint64 page_count,sh_uint64 flags) {
  if (phys_start==0 || page_count==0) return SH_STATUS_INVALID_PARAMETER;
  if (memory_services==SH_NULLPTR) return SH_STATUS_MEMORY_SUBSYSTEM_NOT_INITIALIZED;
  SH_STATUS status=sh_page_is_va_range_mapped_ptp(&memory_services->kernel_ptp,(sh_page_VIRTUAL_ADDRESS)phys_start,page_count*SH_PAGE_SIZE);
  if (status!=SH_STATUS_VA_NOT_MAPPED) {
    if (status==SH_STATUS_VA_MAPPED) {
      return SH_STATUS_ERROR_VA_FULLY_MAPPED;
    } else if (status==SH_STATUS_VA_PARTIALLY_MAPPED) {
      return SH_STATUS_ERROR_VA_PARTIALLY_MAPPED;
    }
  } else if (sh_status_error(status)) {
    return status;
  }
  status=sh_page_map_contiguous_pages_range_ptp(&memory_services->kernel_ptp,(sh_page_VIRTUAL_ADDRESS)phys_start,phys_start,flags,page_count*SH_PAGE_SIZE);
  if (sh_status_error(status)) {
    return status;
  }
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_memory_devs_query(char *sub_path,sh_devs_RESULT *result) {
  if (sub_path==SH_NULLPTR || result==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_uint64 path_len=sh_string_len(sub_path);
  if (path_len==0 || path_len>256) return SH_STATUS_INVALID_PARAMETER;
  if (memory_services==SH_NULLPTR) return SH_STATUS_DEVS_NOT_INITIALIZED;
  if (sh_string_compare(sub_path,"/free-pages",11)) {
    if (path_len==11) {
      result->type=SH_DEVS_VALUE;
      result->value=memory_services->physical_plane.free_pages;
      return SH_STATUS_SUCCESS;
    } else {
      return SH_STATUS_NOT_FOUND;
    }
  } else if (sh_string_compare(sub_path,"/used-pages",11)) {
    if (path_len==11) {
      result->type=SH_DEVS_VALUE;
      result->value=memory_services->physical_plane.used_pages;
      return SH_STATUS_SUCCESS;
    } else {
      return SH_STATUS_NOT_FOUND;
    }
  } else if (sh_string_compare(sub_path,"/total-pages",12)) {
    if (path_len==12) {
      result->type=SH_DEVS_VALUE;
      result->value=sh_page_get_physical_memory_amount_pages();
      return SH_STATUS_SUCCESS;
    } else {
      return SH_STATUS_NOT_FOUND;
    }
  } else if (sh_string_compare(sub_path,"/installed-pages",16)) {
    if (path_len==16) {
      result->type=SH_DEVS_VALUE;
      result->value=sh_page_get_installed_memory_bytes()/SH_PAGE_SIZE;
      return SH_STATUS_SUCCESS;
    } else {
      return SH_STATUS_NOT_FOUND;
    }
  } else {
    return SH_STATUS_NOT_FOUND;
  }
}
