#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 kapi::memory
{

  namespace
  {
    //! The size of the simulated RAM.
    constexpr auto memory_size = kstd::units::MiB(32);
    constexpr auto number_of_frames = memory_size / 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<frame_allocator *>{};
    auto constinit old_mapper = std::optional<page_mapper *>{};

    auto handoff_to_kernel_pmm(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({frame{first_free_frame}, number_of_free_frames});
    }

  }  // namespace

  auto init() -> void
  {
    bump_allocator.emplace();
    test_mapper.emplace(memory_size);

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

    init_pmm(memory_size / frame::size, handoff_to_kernel_pmm);
  }

  auto reset() -> 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();
  }

}  // namespace kapi::memory