// SPDX-License-Identifier: MPL-2.0
#include "memory/slabs/slab_reg_virt.h"
#include "std/mem.h"
#include "kernel/log.h"
#include "memory/pez/pez.h"
SH_STATUS sh_slab_reg_virt_alloc_init(sh_slab_reg_virt_SLAB_ALLOCATOR* slab_alloc,sh_page_PAGE_TABLE_POOL *ptp) {
  if (slab_alloc==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  slab_alloc->max_slabs=SH_SLAB_REG_VIRT_MAX_SLAB;
  slab_alloc->slab_count=0;
  slab_alloc->partial_head=SH_NULLPTR;
  slab_alloc->slabs_data=SH_NULLPTR;
  sh_page_VIRTUAL_ADDRESS header_va;
  SH_STATUS status=sh_page_alloc_contiguous_extended(ptp,SH_SLAB_REG_VIRT_HEADER_LIST_SIZE_BYTES,&header_va,SH_PAGE_RW | SH_PAGE_PRESENT | SH_PAGE_NX,SH_PAGE_KERNEL_PERM_VA_BASE,(SH_PAGE_KERNEL_PERM_VA_END-SH_PAGE_KERNEL_PERM_VA_BASE+1-0x1000));
  if (sh_status_error(status)) return status;
  slab_alloc->slabs_header=(sh_slab_reg_virt_SLAB_STRUCT*)header_va;
  status=sh_mem_set_8((sh_uint8*)(slab_alloc->slabs_header),0x00,SH_SLAB_REG_VIRT_HEADER_LIST_SIZE_BYTES);
  if (sh_status_error(status)) return status;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_reg_virt_add_slab(sh_slab_reg_virt_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL *ptp,sh_slab_reg_virt_SLAB_STRUCT** out_slab) {
  if (alloc==SH_NULLPTR || out_slab==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  if (alloc->slab_count>=alloc->max_slabs) return SH_STATUS_SLAB_ALLOCATOR_FULL;
  sh_page_VIRTUAL_ADDRESS slab_pa;
  SH_STATUS status;
  if (sh_pez_is_available()) {
    sh_pez_PHYSICAL_PLANE *phys_plane=sh_pez_get_reference_phys_plane();
    status=sh_pez_alloc_physical_pages(phys_plane,SH_SLAB_REG_PHYS_SLAB_DATA_PAGES,&slab_pa);
    if (sh_status_error(status)) return status;
  } else {
    status=sh_page_search_physical_contiguous_block_na(SH_SLAB_REG_PHYS_SLAB_DATA_PAGES,&slab_pa);
    if (sh_status_error(status)) return status;
  }
  if (alloc->slabs_data==SH_NULLPTR && alloc->slab_count==0) {
    sh_page_VIRTUAL_ADDRESS alloc_slab_data_va_base=SH_SLAB_REG_VIRT_DATA_VA;
    status=sh_page_is_va_range_mapped_ptp(ptp,alloc_slab_data_va_base,SH_SLAB_REG_VIRT_MAX_SLAB*SH_SLAB_REG_VIRT_SLAB_DATA_SIZE_BYTES);
    if (status!=SH_STATUS_VA_NOT_MAPPED) return status;
    status=sh_page_map_contiguous_pages_range_ptp(ptp,alloc_slab_data_va_base,slab_pa,SH_PAGE_PRESENT | SH_PAGE_NX | SH_PAGE_RW,SH_SLAB_REG_VIRT_SLAB_DATA_SIZE_BYTES);
    if (sh_status_error(status)) return status;
    alloc->slabs_data=(sh_uint8*)alloc_slab_data_va_base;
    sh_slab_reg_virt_SLAB_STRUCT *new_slab_header=alloc->slabs_header;
    new_slab_header->slab_index=alloc->slab_count;
    new_slab_header->used_count=0;
    new_slab_header->prev=SH_NULLPTR;
    new_slab_header->next=SH_NULLPTR;
    for (sh_iter64 i=0;i<16;i++) new_slab_header->free_bitmap[i]=0;
    new_slab_header->free_bitmap[0]|=1ULL;
    new_slab_header->used_count=0;
    status=sh_mem_set_8(alloc->slabs_data,0x00,SH_SLAB_REG_VIRT_SLAB_DATA_SIZE_BYTES);
    if (sh_status_error(status)) return status;
    if (!sh_pez_is_available()) {
      sh_page_MEM_STATS mem_stats;
      status=sh_page_get_memory_stats(&mem_stats);
      if (sh_status_error(status)) return status;
      status=sh_page_set_pages_range_bitmap((sh_uint8*)sh_page_get_physical_bitmap_ptr(),mem_stats.memory_total_pages,slab_pa/SH_PAGE_SIZE,SH_SLAB_REG_VIRT_SLAB_DATA_PAGES,SH_TRUE);
      if (sh_status_error(status)) return status;
    }
    alloc->partial_head=alloc->slabs_header;
    *out_slab=new_slab_header;
    alloc->slab_count++;
  } else if (alloc->slabs_data!=SH_NULLPTR && alloc->slab_count!=0) {
    sh_page_VIRTUAL_ADDRESS alloc_slab_data_va=SH_SLAB_REG_VIRT_DATA_VA+alloc->slab_count*SH_SLAB_REG_VIRT_SLAB_DATA_SIZE_BYTES;
    status=sh_page_is_va_range_mapped_ptp(ptp,alloc_slab_data_va,SH_SLAB_REG_VIRT_SLAB_DATA_SIZE_BYTES);
    if (status!=SH_STATUS_VA_NOT_MAPPED) return status;
    status=sh_page_map_contiguous_pages_range_ptp(ptp,alloc_slab_data_va,slab_pa,SH_PAGE_PRESENT | SH_PAGE_NX | SH_PAGE_RW,SH_SLAB_REG_VIRT_SLAB_DATA_SIZE_BYTES);
    if (sh_status_error(status)) return status;
    sh_slab_reg_virt_SLAB_STRUCT *new_slab_header=alloc->slabs_header+alloc->slab_count;
    new_slab_header->slab_index=alloc->slab_count;
    new_slab_header->used_count=0;
    new_slab_header->prev=SH_NULLPTR;
    new_slab_header->next=alloc->partial_head;
    sh_slab_reg_virt_SLAB_STRUCT *next_slab_header=new_slab_header->next;
    next_slab_header->prev=new_slab_header;
    alloc->partial_head=new_slab_header;
    for (sh_iter64 i=0;i<16;i++) new_slab_header->free_bitmap[i]=0;
    new_slab_header->free_bitmap[0]|=1ULL;
    new_slab_header->used_count=0;
    status=sh_mem_set_8(alloc->slabs_data+alloc->slab_count*SH_SLAB_REG_VIRT_SLAB_DATA_SIZE_BYTES,0x00,SH_SLAB_REG_VIRT_SLAB_DATA_SIZE_BYTES);
    if (sh_status_error(status)) return status;
    if (!sh_pez_is_available()) {
      sh_page_MEM_STATS mem_stats;
      status=sh_page_get_memory_stats(&mem_stats);
      if (sh_status_error(status)) return status;
      status=sh_page_set_pages_range_bitmap((sh_uint8*)sh_page_get_physical_bitmap_ptr(),mem_stats.memory_total_pages,slab_pa/SH_PAGE_SIZE,SH_SLAB_REG_VIRT_SLAB_DATA_PAGES,SH_TRUE);
      if (sh_status_error(status)) return status;
    }
    *out_slab=new_slab_header;
    alloc->slab_count++;
  }
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_reg_virt_get_partial_slab(sh_slab_reg_virt_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL* ptp,sh_slab_reg_virt_SLAB_STRUCT** found_slab) {
  if (alloc==SH_NULLPTR || ptp==SH_NULLPTR || found_slab==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  if (alloc->partial_head==SH_NULLPTR) {
    *found_slab=SH_NULLPTR;
    return SH_STATUS_NOT_FOUND;
  }
  sh_slab_reg_virt_SLAB_STRUCT* slab=alloc->partial_head;
  if (slab->used_count>=SH_SLAB_REG_VIRT_ACTUAL_OBJECTS_PER_SLAB) {
    if (slab->next!=SH_NULLPTR && slab->next->used_count<SH_SLAB_REG_VIRT_ACTUAL_OBJECTS_PER_SLAB) {
      alloc->partial_head=slab->next;
      alloc->partial_head->prev=SH_NULLPTR;
      slab->next=SH_NULLPTR;
      slab->prev=SH_NULLPTR;
      *found_slab=alloc->partial_head;
      return SH_STATUS_SUCCESS;
    } else {
      sh_slab_reg_virt_SLAB_STRUCT* new_slab;
      SH_STATUS status=sh_slab_reg_virt_add_slab(alloc,ptp,&new_slab);
      if (sh_status_error(status)) return status;
      *found_slab=new_slab;
      return SH_STATUS_NEW_SLAB_ADDED;
    }
  }
  *found_slab=slab;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_reg_virt_scan_slabs(sh_slab_reg_virt_SLAB_ALLOCATOR* alloc) {
  if (alloc==SH_NULLPTR || alloc->slabs_data==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_slab_reg_virt_SLAB_STRUCT* first_partial=SH_NULLPTR;
  sh_slab_reg_virt_SLAB_STRUCT* last_partial=SH_NULLPTR;
  sh_log_debug("Slabs rescan triggered.",SH_LOG_SOURCE_SLAB);
  for (sh_uint32 i=0;i<alloc->slab_count;i++) {
    sh_slab_reg_virt_SLAB_STRUCT* slab=alloc->slabs_header+i;
    sh_uint8* slab_base=alloc->slabs_data+i*SH_SLAB_REG_VIRT_SLAB_DATA_SIZE_BYTES;
    slab->used_count=0;
    for (sh_iter64 w=0;w<16;w++) slab->free_bitmap[w]=0;
    slab->free_bitmap[0]|=1ULL;
    for (sh_uint16 obj_idx=1;obj_idx<SH_SLAB_REG_VIRT_OBJECTS_PER_SLAB;obj_idx++) {
      sh_uint8* obj_ptr=slab_base+obj_idx*SH_SLAB_REG_VIRT_OBJECT_SIZE_BYTES;
      int is_free=1;
      for (sh_iter64 b=0;b<SH_SLAB_REG_VIRT_OBJECT_SIZE_BYTES;b++) {
        if (obj_ptr[b]!=0) {
          is_free=0;
          break;
        }
      }
      if (is_free) {
        sh_uint16 word=obj_idx>>6;
        sh_uint16 bit=obj_idx&63;
        slab->free_bitmap[word]&= ~(1ULL<<bit);
      } else {
        sh_uint16 word=obj_idx>>6;
        sh_uint16 bit=obj_idx&63;
        slab->free_bitmap[word]|=(1ULL<<bit);
        slab->used_count++;
      }
    }
    sh_log_fdebug(SH_LOG_SOURCE_SLAB,"Slab index: %4u; used count: %4u\n",slab->slab_index,slab->used_count);
    slab->prev=SH_NULLPTR;
    slab->next=SH_NULLPTR;
    if (slab->used_count>0 && slab->used_count<SH_SLAB_REG_VIRT_ACTUAL_OBJECTS_PER_SLAB) {
      if (first_partial==SH_NULLPTR) {
        first_partial=slab;
        last_partial=slab;
      } else {
        slab->prev=last_partial;
        last_partial->next=slab;
        last_partial=slab;
      }
    }
  }
  alloc->partial_head=first_partial;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_reg_virt_find_free_object(sh_slab_reg_virt_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL *ptp,sh_slab_reg_virt_OBJECT_INDEX* out_ref) {
  if (alloc==SH_NULLPTR || ptp==SH_NULLPTR || out_ref==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_slab_reg_virt_SLAB_STRUCT *partial_slab;
  sh_slab_reg_virt_SLAB_STRUCT *new_slab;
  SH_STATUS status=sh_slab_reg_virt_get_partial_slab(alloc,ptp,&partial_slab);
  if (status==SH_STATUS_RESCAN_NEEDED) {
    status=sh_slab_reg_virt_scan_slabs(alloc);
    if (sh_status_error(status)) return status;
    status=sh_slab_reg_virt_get_partial_slab(alloc,ptp,&partial_slab);
  }
  if (status==SH_STATUS_NOT_FOUND) {
    status=sh_slab_reg_virt_add_slab(alloc,ptp,&new_slab);
    if (sh_status_error(status)) return status;
    status=sh_slab_reg_virt_get_partial_slab(alloc,ptp,&partial_slab);
    if (sh_status_error(status)) return status;
  } else if (sh_status_error(status) && status!=SH_STATUS_NEW_SLAB_ADDED) {
    return status;
  }
  sh_uint64 obj;
  sh_bool obj_found=SH_FALSE;
  for (sh_uint8 w=0;w<16;w++) {
    sh_uint64 word=partial_slab->free_bitmap[w];
    if (word!=SH_UINT64_MAX) {
      if (w==0) {
        for (sh_uint8 bit=0;bit<64;bit++) {
          if (!(word & (1ULL<<bit))) {
            sh_uint64 candidate=(sh_uint64)w*64+bit;
            if (candidate==0) continue;
            obj=candidate;
            obj_found=SH_TRUE;
            break;
          }
        }
      } else {
        for (sh_uint8 bit=0;bit<64;bit++) {
          if (!(word & (1ULL<<bit))) {
            obj=(sh_uint64)w*64+bit;
            obj_found=SH_TRUE;
            break;
          }
        }
      }
    }
    if (obj_found==SH_TRUE) break;
  }
  if (obj_found==SH_FALSE) {
    return SH_STATUS_RESCAN_NEEDED;
  }
  *out_ref=SH_SLAB_REG_VIRT_MAKE_REF(partial_slab->slab_index,obj);
  return SH_STATUS_SUCCESS;
}
void* sh_slab_reg_virt_ref_to_ptr(sh_slab_reg_virt_SLAB_ALLOCATOR* alloc,sh_slab_reg_virt_OBJECT_INDEX ref) {
  if (alloc==SH_NULLPTR) {
    return SH_NULLPTR;
  }
  if (SH_SLAB_REG_VIRT_REF_SLAB(ref)>=alloc->slab_count || SH_SLAB_REG_VIRT_REF_OBJECT(ref)==0) {
    return SH_NULLPTR;
  }
  sh_uint32 slab=SH_SLAB_REG_VIRT_REF_SLAB(ref);
  sh_uint16 object=SH_SLAB_REG_VIRT_REF_OBJECT(ref);
  if (object==0) {
    return SH_NULLPTR;
  }
  return alloc->slabs_data+slab*SH_SLAB_REG_VIRT_SLAB_DATA_SIZE_BYTES+object*SH_SLAB_REG_VIRT_OBJECT_SIZE_BYTES;
}
SH_STATUS sh_slab_reg_virt_alloc(sh_slab_reg_virt_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL *ptp,sh_slab_reg_virt_OBJECT_INDEX* out_index) {
  if (alloc==SH_NULLPTR || ptp==SH_NULLPTR || out_index==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_slab_reg_virt_OBJECT_INDEX ref;
  SH_STATUS status=sh_slab_reg_virt_find_free_object(alloc,ptp,&ref);
  if (sh_status_error(status)) return status;
  sh_uint32 slab_idx=SH_SLAB_REG_VIRT_REF_SLAB(ref);
  sh_uint16 obj_idx=SH_SLAB_REG_VIRT_REF_OBJECT(ref);
  sh_slab_reg_virt_SLAB_STRUCT* slab=alloc->slabs_header+slab_idx;
  sh_uint16 word=obj_idx>>6;
  sh_uint16 bit=obj_idx&63;
  slab->free_bitmap[word]|=(1ULL<<bit);
  sh_pez_REGION_VIRTUAL_OBJECT* obj=(sh_pez_REGION_VIRTUAL_OBJECT*)sh_slab_reg_virt_ref_to_ptr(alloc,ref);
  if (obj==SH_NULLPTR) return SH_STATUS_ERROR_NULLPTR_RETURNED;
  sh_mem_set_8((sh_uint8*)obj,0x0,SH_SLAB_REG_VIRT_OBJECT_SIZE_BYTES);
  slab->used_count++;
  *out_index=ref;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_reg_virt_dealloc(sh_slab_reg_virt_SLAB_ALLOCATOR* alloc,sh_slab_reg_virt_OBJECT_INDEX index) {
  if (alloc==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  if (SH_SLAB_REG_VIRT_REF_OBJECT(index)==0) return SH_STATUS_INVALID_PARAMETER;
  if (SH_SLAB_REG_VIRT_REF_SLAB(index)>=alloc->slab_count) return SH_STATUS_INVALID_PARAMETER;
  sh_uint16 obj_idx=SH_SLAB_REG_VIRT_REF_OBJECT(index);
  sh_uint16 word=obj_idx>>6;
  sh_uint16 bit=obj_idx&63;
  sh_slab_reg_virt_SLAB_STRUCT* slab_header=alloc->slabs_header+SH_SLAB_REG_VIRT_REF_SLAB(index);
  if (!(slab_header->free_bitmap[word] & (1ULL<<bit))) return SH_STATUS_SLAB_SLOT_ALREADY_FREE;
  slab_header->free_bitmap[word]&= ~(1ULL<<bit);
  if (slab_header->used_count==SH_SLAB_REG_VIRT_ACTUAL_OBJECTS_PER_SLAB) {
    slab_header->prev=SH_NULLPTR;
    slab_header->next=alloc->partial_head;
    if (alloc->partial_head!=SH_NULLPTR) {
      alloc->partial_head->prev=slab_header;
    }
    alloc->partial_head=slab_header;
  }
  slab_header->used_count--;
  sh_uint8* obj_ptr=sh_slab_reg_virt_ref_to_ptr(alloc,index);
  if (obj_ptr==SH_NULLPTR) return SH_STATUS_ERROR_NULLPTR_RETURNED;
  sh_mem_set_8(obj_ptr,0x00,SH_SLAB_REG_VIRT_OBJECT_SIZE_BYTES);
  return SH_STATUS_SUCCESS;
}
