#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);
    constexpr auto number_of_simulated_frames = simulate_memory_size / frame::size;

    struct test_bootstrap_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{};
    };

    auto constinit bootstrap_allocator = std::optional<test_bootstrap_frame_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;
      }
    };

    auto constinit test_mapper = std::optional<test_page_mapper>{};

    auto handoff_to_kernel_pmm(frame_allocator & new_allocator) -> void
    {
      auto first_free_frame = bootstrap_allocator->next_free_frame;
      auto number_of_free_frames = number_of_simulated_frames - first_free_frame;
      new_allocator.release_many({frame{first_free_frame}, number_of_free_frames});
    }

    auto constinit previous_frame_allocator = std::optional<frame_allocator *>{};
    auto constinit previous_page_mapper = std::optional<page_mapper *>{};

  }  // namespace

  auto init() -> void
  {
    kernel::tests::simulated_memory::init(simulate_memory_size);

    bootstrap_allocator.emplace();
    test_mapper.emplace();

    previous_frame_allocator = set_frame_allocator(*bootstrap_allocator);
    previous_page_mapper = set_page_mapper(*test_mapper);

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

  auto reset() -> void
  {
    kernel::tests::simulated_memory::reset();

    if (previous_frame_allocator && *previous_frame_allocator)
    {
      set_frame_allocator(**previous_frame_allocator);
    }

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

    bootstrap_allocator.reset();
    test_mapper.reset();
  }

}  // namespace kapi::memory