#include "kapi/memory.hpp"

#include <kapi/memory.hpp>

#include "kernel/tests/simulated_memory.hpp"

#include <kstd/print>
#include <kstd/units>

#include <cstddef>
#include <cstdint>
#include <optional>
#include <utility>

namespace kapi::memory
{

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

    struct test_boostrap_frame_allocator : frame_allocator
    {
      auto mark_used(frame) -> void override {}

      auto allocate_many(std::size_t count) noexcept -> std::optional<std::pair<frame, std::size_t>> override
      {
        auto start = next_free_frame;
        next_free_frame += count;
        return std::pair{frame{start}, count};
      }

      auto release_many(std::pair<frame, std::size_t>) -> void override {}

      std::size_t next_free_frame{};
    } boostrap_allocator;

    struct test_page_mapper : page_mapper
    {
      auto map(page page, frame frame, flags flags) -> std::byte * override
      {
        kstd::println("mapping page {} onto frame {} with flags {}", page.number(), frame.number(),
                      static_cast<std::uint64_t>(flags));

        if ((page.start_address() & pmm_metadata_base.raw()) == pmm_metadata_base.raw())
        {
          auto offset = page.start_address() & ~pmm_metadata_base.raw();
          return kernel::tests::simulated_memory::pmm_metadata_base() + offset;
        }

        return nullptr;
      }

      auto unmap(page page) -> void override
      {
        kstd::println("unmapping page {}", page.number());
      }

      auto try_unmap(page page) noexcept -> bool override
      {
        kstd::println("trying to unmap page {}", page.number());
        return false;
      }
    } test_mapper;

    auto handoff_to_kernel_pmm(frame_allocator & new_allocator) -> void
    {
      for (auto i = 0uz; i < boostrap_allocator.next_free_frame; ++i)
      {
        new_allocator.mark_used(frame{i});
      }
    }

  }  // namespace

  auto init() -> void
  {
    kernel::tests::simulated_memory::init(simulate_memory_size);
    set_frame_allocator(boostrap_allocator);
    set_page_mapper(test_mapper);

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

}  // namespace kapi::memory