// SPDX-License-Identifier: MPL-2.0
#include "memory/slabs/slab_generic.h"
#include "std/mem.h"
#include "kernel/log.h"
static sh_uint8 sig[8]=SH_SLAB_GENERIC_SLAB_SIG;
static sh_uint64 object_size_bytes[8]=SH_SLAB_GENERIC_OBJECT_SIZE_BYTES;
static sh_uint64 actual_objects_per_slab[8]=SH_SLAB_GENERIC_ACTUAL_OBJECTS_PER_SLAB;
static sh_uint64 bitmap_init[8]=SH_SLAB_GENERIC_SLAB_BITMAP_INIT;
static sh_uint64 slab_data_pages[8]=SH_SLAB_GENERIC_SLAB_DATA_PAGES;
SH_STATUS sh_slab_generic_alloc_init(sh_uint8 level,struct sh_slab_generic_SLAB_ALLOCATOR* slab_alloc,sh_pba_PAGE_BLOCK_ALLOCATOR *pba) {
  if (slab_alloc==SH_NULLPTR || pba==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  if (level>7) return SH_STATUS_INVALID_PARAMETER;
  slab_alloc->first_slab=SH_NULLPTR;
  slab_alloc->slab_count=0;
  slab_alloc->partial_head=SH_NULLPTR;
  slab_alloc->pba=pba;
  slab_alloc->level=level;
  slab_alloc->object_size_bytes=object_size_bytes[level];
  slab_alloc->object_per_slab=SH_SLAB_GENERIC_OBJECTS_PER_SLAB;
  slab_alloc->actual_object_per_slab=actual_objects_per_slab[level];
  slab_alloc->bitmap_init=bitmap_init[level];
  slab_alloc->slab_pages_count=slab_data_pages[level];
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_generic_add_slab(struct sh_slab_generic_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL *ptp,sh_slab_generic_SLAB** out_slab) {
  if (alloc==SH_NULLPTR || out_slab==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER; // phys_plane isn't checked because it can be SH_NULLPTR
  sh_page_VIRTUAL_ADDRESS slab_va;
  SH_STATUS status=sh_pba_alloc(alloc->pba,ptp,&slab_va);
  if (sh_status_error(status)) return status;
  if (alloc->slab_count==0) {
    sh_slab_generic_SLAB *new_slab=(sh_slab_generic_SLAB*)slab_va;
    for (sh_iter64 i=0;i<sizeof(sig);i++) new_slab->sig[i]=sig[i];
    new_slab->next_slab=SH_NULLPTR;
    new_slab->prev_slab=SH_NULLPTR;
    new_slab->next_partial=SH_NULLPTR;
    new_slab->prev_partial=SH_NULLPTR;
    new_slab->slab_index=alloc->slab_count;
    new_slab->used_count=0;
    for (sh_iter64 i=0;i<8;i++) new_slab->free_bitmap[i]=0;
    new_slab->free_bitmap[0]|=alloc->bitmap_init;
    alloc->first_slab=new_slab;
    alloc->partial_head=new_slab;
    alloc->slab_count++;
    *out_slab=new_slab;
  } else {
    sh_slab_generic_SLAB *new_slab=(sh_slab_generic_SLAB*)slab_va;
    for (sh_iter64 i=0;i<sizeof(sig);i++) new_slab->sig[i]=sig[i];
    new_slab->next_slab=alloc->first_slab;
    new_slab->prev_slab=SH_NULLPTR;
    alloc->first_slab->prev_slab=new_slab;
    alloc->first_slab=new_slab;
    new_slab->next_partial=alloc->partial_head;
    new_slab->prev_partial=SH_NULLPTR;
    alloc->partial_head->prev_partial=new_slab;
    alloc->partial_head=new_slab;
    new_slab->slab_index=alloc->slab_count;
    new_slab->used_count=0;
    for (sh_iter64 i=0;i<8;i++) new_slab->free_bitmap[i]=0;
    new_slab->free_bitmap[0]|=alloc->bitmap_init;
    alloc->slab_count++;
    *out_slab=new_slab;
  }
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_generic_get_partial_slab(struct sh_slab_generic_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL* ptp,sh_slab_generic_SLAB** 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_generic_SLAB *slab=alloc->partial_head;
  if (slab->used_count>=alloc->actual_object_per_slab) {
    if (slab->next_partial!=SH_NULLPTR && slab->next_partial->used_count<alloc->actual_object_per_slab) {
      alloc->partial_head=slab->next_partial;
      alloc->partial_head->prev_partial=SH_NULLPTR;
      slab->next_partial=SH_NULLPTR;
      slab->prev_partial=SH_NULLPTR;
      *found_slab=alloc->partial_head;
      return SH_STATUS_SUCCESS;
    } else {
      sh_slab_generic_SLAB *new_slab;
      SH_STATUS status=sh_slab_generic_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_generic_scan_slabs(struct sh_slab_generic_SLAB_ALLOCATOR* alloc) {
  if (alloc==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_slab_generic_SLAB* current=alloc->first_slab;
  sh_slab_generic_SLAB* last_partial=SH_NULLPTR;
  alloc->partial_head=SH_NULLPTR;
  sh_log_debug("Slabs rescan triggered.",SH_LOG_SOURCE_SLAB);
  while (current!=SH_NULLPTR) {
    if (*(sh_uint64*)current->sig!=SH_SLAB_GENERIC_MAGIC) {
      return SH_STATUS_FOUND_CORRUPTED_SLAB;
    }
    current->next_partial=SH_NULLPTR;
    current->prev_partial=SH_NULLPTR;
    if (current->used_count<alloc->actual_object_per_slab) {
      if (alloc->partial_head==SH_NULLPTR) {
        alloc->partial_head=current;
        last_partial=current;
      } else {
        last_partial->next_partial=current;
        current->prev_partial=last_partial;
        last_partial=current;
      }
    }
    current=current->next_slab;
  }
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_generic_find_free_object(struct sh_slab_generic_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL *ptp,void** out,sh_slab_generic_OBJECT_INDEX_IN_SLAB* index_in_slab) {
  if (alloc==SH_NULLPTR || ptp==SH_NULLPTR || out==SH_NULLPTR || index_in_slab==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_slab_generic_SLAB *partial_slab;
  sh_slab_generic_SLAB *new_slab;
  SH_STATUS status=sh_slab_generic_get_partial_slab(alloc,ptp,&partial_slab);
  if (status==SH_STATUS_RESCAN_NEEDED) {
    status=sh_slab_generic_scan_slabs(alloc);
    if (sh_status_error(status)) {
      return status;
    }
    status=sh_slab_generic_get_partial_slab(alloc,ptp,&partial_slab);
  }
  if (status==SH_STATUS_NOT_FOUND) {
    status=sh_slab_generic_add_slab(alloc,ptp,&new_slab);
    if (sh_status_error(status)) {
      return status;
    }
    status=sh_slab_generic_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_uint16 found_index=0;
  sh_bool found=SH_FALSE;
  for (sh_iter64 w=0;w<8;w++) {
    sh_uint64 word=partial_slab->free_bitmap[w];
    sh_uint64 free_bits=(~word);
    if (free_bits) {
      sh_uint64 bit=(sh_uint64)__builtin_ctzll(free_bits);
      found_index=(sh_uint16)((w<<6)+bit);
      found=SH_TRUE;
      break;
    }
  }
  if (!found) {
    return SH_STATUS_NOT_FOUND;
  }
  sh_uint8 *start_data=(sh_uint8*)partial_slab;
  *out=start_data+found_index*alloc->object_size_bytes;
  *index_in_slab=found_index;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_generic_alloc(struct sh_slab_generic_SLAB_ALLOCATOR* alloc,sh_page_PAGE_TABLE_POOL *ptp,void** out_obj) {
  if (alloc==SH_NULLPTR || ptp==SH_NULLPTR || out_obj==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  sh_uint8* object_ptr;
  sh_slab_generic_OBJECT_INDEX_IN_SLAB index_in_slab;
  SH_STATUS status=sh_slab_generic_find_free_object(alloc,ptp,(void*)&object_ptr,&index_in_slab);
  if (sh_status_error(status)) return status;
  if (object_ptr==SH_NULLPTR) return SH_STATUS_ERROR_NULLPTR_RETURNED;
  for (sh_iter64 i=0;i<alloc->object_size_bytes;i++) object_ptr[i]=0;
  sh_slab_generic_SLAB* slab=(sh_slab_generic_SLAB*)((sh_uint64)object_ptr & ~(0xFFFULL<<alloc->level));
  if (*(sh_uint64*)slab->sig!=SH_SLAB_GENERIC_MAGIC) {
    return SH_STATUS_FOUND_CORRUPTED_SLAB;
  }
  sh_uint16 word_idx=index_in_slab>>6;
  sh_uint16 bit_idx=index_in_slab&63;
  slab->free_bitmap[word_idx]|=(1ULL<<bit_idx);
  slab->used_count++;
  *out_obj=object_ptr;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_slab_generic_dealloc(struct sh_slab_generic_SLAB_ALLOCATOR* alloc,void *object_ptr) {
  if (alloc==SH_NULLPTR || object_ptr==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  if ((sh_uint64)object_ptr%alloc->object_size_bytes!=0) return SH_STATUS_INVALID_PARAMETER;
  sh_slab_generic_SLAB* slab=(sh_slab_generic_SLAB*)((sh_uint64)object_ptr & ~(0xFFFULL<<alloc->level));
  if (*(sh_uint64*)slab->sig!=SH_SLAB_GENERIC_MAGIC) {
    return SH_STATUS_FOUND_CORRUPTED_SLAB;
  }
  sh_uint16 object_idx=(sh_uint16)(((sh_uint64)object_ptr & (0xFFFULL<<alloc->level))/alloc->object_size_bytes);
  sh_uint16 min_idx=(sh_uint16)(alloc->object_per_slab-alloc->actual_object_per_slab);
  if (object_idx<min_idx) return SH_STATUS_INVALID_PARAMETER;
  sh_uint16 word=object_idx>>6;
  sh_uint16 bit=object_idx&63;
  if (!(slab->free_bitmap[word] & (1ULL<<bit))) return SH_STATUS_SLAB_SLOT_ALREADY_FREE;
  slab->free_bitmap[word]&= ~(1ULL<<bit);
  if (slab->used_count==alloc->actual_object_per_slab) {
    slab->prev_partial=SH_NULLPTR;
    slab->next_partial=alloc->partial_head;
    if (alloc->partial_head!=SH_NULLPTR) {
      alloc->partial_head->prev_partial=slab;
    }
    alloc->partial_head=slab;
  }
  slab->used_count--;
  for (sh_iter64 i=0;i<alloc->object_size_bytes;i++) ((sh_uint8*)object_ptr)[i]=0;
  return SH_STATUS_SUCCESS;
}
