#include "kapi/memory.hpp"

#include <kapi/memory.hpp>

#include "kernel/test_support/bump_frame_allocator.hpp"
#include "kernel/test_support/page_mapper.hpp"

#include <kstd/units>

#include <optional>

namespace
{
  //! The size of the simulated RAM.
  constexpr auto physical_size = kstd::units::MiB(32);
  constexpr auto virtual_size = kstd::units::GiB(1);

  constexpr auto number_of_frames = physical_size / kapi::memory::frame::size;

  auto constinit bump_allocator = std::optional<kernel::tests::bump_frame_allocator>{};
  auto constinit test_mapper = std::optional<kernel::tests::page_mapper>{};

  auto constinit old_allocator = std::optional<kapi::memory::frame_allocator *>{};
  auto constinit old_mapper = std::optional<kapi::memory::page_mapper *>{};

  auto handoff_to_kernel_pmm(kapi::memory::frame_allocator & new_allocator) -> void
  {
    auto first_free_frame = bump_allocator->next_free_frame;
    auto number_of_free_frames = number_of_frames - first_free_frame;
    new_allocator.release_many({kapi::memory::frame{first_free_frame}, number_of_free_frames});
  }

}  // namespace

namespace kapi::memory
{

  auto init() -> void
  {
    bump_allocator.emplace();
    test_mapper.emplace(physical_size, virtual_size);

    old_allocator = set_frame_allocator(*bump_allocator);
    old_mapper = set_page_mapper(*test_mapper);

    init_pmm(physical_size / frame::size, handoff_to_kernel_pmm);
  }
}  // namespace kapi::memory

namespace kernel::tests::memory
{

  auto deinit() -> void
  {
    if (old_allocator && *old_allocator)
    {
      set_frame_allocator(**old_allocator);
    }

    if (old_mapper && *old_mapper)
    {
      set_page_mapper(**old_mapper);
    }

    bump_allocator.reset();
    test_mapper.reset();
  }

  auto virtual_base() -> kapi::memory::linear_address
  {
    return test_mapper->memory.virtual_base();
  }

}  // namespace kernel::tests::memory