// SPDX-License-Identifier: MPL-2.0
#include "memory/pba.h"
#include "memory/pez/pez.h"
#include "kernel/log.h"
SH_STATUS sh_pba_init(sh_pba_PAGE_BLOCK_ALLOCATOR *pba,sh_page_VIRTUAL_ADDRESS start_va,sh_uint64 area_pages_amount,sh_uint64 block_pages) {
  if (pba==SH_NULLPTR) {
    return SH_STATUS_INVALID_PARAMETER;
  }
  if (area_pages_amount%block_pages!=0) {
    return SH_STATUS_INVALID_PARAMETER;
  }
  if (start_va%(SH_PAGE_SIZE*block_pages)!=0) {
    return SH_STATUS_INVALID_PARAMETER;
  }
  pba->start_va=start_va;
  pba->block_pages=block_pages;
  pba->block_count=0;
  pba->total_pages=area_pages_amount;
  pba->max_blocks=area_pages_amount/block_pages;
  return SH_STATUS_SUCCESS;
}
SH_STATUS sh_pba_alloc(sh_pba_PAGE_BLOCK_ALLOCATOR *pba,sh_page_PAGE_TABLE_POOL *ptp,sh_page_VIRTUAL_ADDRESS *ptr) {
  if (pba==SH_NULLPTR || ptr==SH_NULLPTR) return SH_STATUS_INVALID_PARAMETER;
  if (pba->block_count==pba->max_blocks) return SH_STATUS_PBA_FULL;
  if (!sh_pez_is_available()) {
    sh_page_PHYSICAL_ADDRESS block_pa;
    SH_STATUS status=sh_page_search_physical_contiguous_block_na(pba->block_pages,&block_pa);
    if (sh_status_error(status)) return status;
    sh_page_VIRTUAL_ADDRESS temp_ptr=pba->start_va+pba->block_count*pba->block_pages*SH_PAGE_SIZE;
    status=sh_page_is_va_range_mapped_ptp(ptp,temp_ptr,pba->block_pages*SH_PAGE_SIZE);
    if (status==SH_STATUS_VA_FULLY_MAPPED || status==SH_STATUS_VA_PARTIALLY_MAPPED) {
      return status;
    } else if (sh_status_error(status)) {
      return status;
    } else if (status!=SH_STATUS_VA_NOT_MAPPED) {
      return status;
    }
    status=sh_page_set_pages_range_bitmap((sh_uint8*)sh_page_get_physical_bitmap_ptr(),sh_page_get_physical_memory_amount_pages(),block_pa/SH_PAGE_SIZE,pba->block_pages,SH_TRUE);
    if (sh_status_error(status)) return status;
    status=sh_page_map_contiguous_pages_range_ptp(ptp,temp_ptr,block_pa,SH_PAGE_RW | SH_PAGE_NX | SH_PAGE_PRESENT,pba->block_pages*SH_PAGE_SIZE);
    if (sh_status_error(status)) {
      SH_STATUS status2=sh_page_set_pages_range_bitmap((sh_uint8*)sh_page_get_physical_bitmap_ptr(),sh_page_get_physical_memory_amount_pages(),block_pa/SH_PAGE_SIZE,pba->block_pages,SH_FALSE);
      if (sh_status_error(status2)) return status2;
      return status;
    }
    *ptr=temp_ptr;
  } else {
    if (sh_pez_get_reference_phys_plane()==SH_NULLPTR) return SH_STATUS_ERROR_PEZ_INITIALIZED;
    sh_page_PHYSICAL_ADDRESS block_pa;
    SH_STATUS status=sh_pez_alloc_physical_pages(sh_pez_get_reference_phys_plane(),(sh_uint32)pba->block_pages,&block_pa);
    if (sh_status_error(status)) return status;
    sh_page_VIRTUAL_ADDRESS temp_ptr=pba->start_va+pba->block_count*pba->block_pages*SH_PAGE_SIZE;
    status=sh_page_is_va_range_mapped_ptp(ptp,temp_ptr,pba->block_pages*SH_PAGE_SIZE);
    if (status==SH_STATUS_VA_FULLY_MAPPED || status==SH_STATUS_VA_PARTIALLY_MAPPED) {
      return status;
    } else if (sh_status_error(status)) {
      return status;
    } else if (status!=SH_STATUS_VA_NOT_MAPPED) {
      return status;
    }
    status=sh_page_map_contiguous_pages_range_ptp(ptp,temp_ptr,block_pa,SH_PAGE_RW | SH_PAGE_NX | SH_PAGE_PRESENT,pba->block_pages*SH_PAGE_SIZE);
    if (sh_status_error(status)) {
      SH_STATUS status2=sh_page_set_pages_range_bitmap((sh_uint8*)sh_page_get_physical_bitmap_ptr(),sh_page_get_physical_memory_amount_pages(),block_pa/SH_PAGE_SIZE,pba->block_pages,SH_FALSE);
      if (sh_status_error(status2)) return status2;
      return status;
    }
    *ptr=temp_ptr;
  }
  pba->block_count++;
  return SH_STATUS_SUCCESS;
}