#include "x86_64/memory/paging_root.hpp"

#include "kapi/memory.hpp"

#include "x86_64/memory/page_utilities.hpp"

#include <cstdint>
#include <optional>

namespace teachos::memory::x86_64
{

  namespace
  {
    constexpr auto PML_RECURSIVE_BASE = std::uintptr_t{0177777'776'776'776'776'0000uz};
  }

  auto paging_root::get() -> paging_root &
  {
    auto pml4_address = std::bit_cast<paging_root *>(PML_RECURSIVE_BASE);
    return *pml4_address;
  }

  auto paging_root::translate(linear_address address) const -> std::optional<physical_address>
  {
    auto offset = address.raw() % PLATFORM_PAGE_SIZE;
    return translate(page::containing(address)).transform([offset](auto frame) -> auto {
      return physical_address{frame.start_address().raw() + offset};
    });
  }

  auto paging_root::translate(page page) const -> std::optional<frame>
  {
    auto pml3 = next(pml_index<4>(page));

    if (!pml3)
    {
      return std::nullopt;
    }

    auto handle_huge_page = [&] -> std::optional<frame> {
      auto pml3_entry = pml3.transform([&](auto pml3) -> auto { return (*pml3)[pml_index<3>(page)]; });
      if (pml3_entry && pml3_entry->huge())
      {
        auto pml3_entry_frame = *pml3_entry->frame();
        return frame{pml3_entry_frame.number() + pml_index<2>(page) * entry_count + pml_index<1>(page)};
      }

      // auto pml3_entry = (**pml3)[page.start_address().raw() >> 39 & 0x1ff];

      return std::nullopt;
    };

    return pml3.and_then([&](auto pml3) -> auto { return pml3->next(pml_index<3>(page)); })
        .and_then([&](auto pml2) -> auto { return pml2->next(pml_index<2>(page)); })
        .and_then([&](auto pml1) -> auto { return (*pml1)[pml_index<1>(page)].frame(); })
        .or_else(handle_huge_page);
  }

}  // namespace teachos::memory::x86_64