#include "x86_64/memory/paging_root.hpp"

#include "kapi/memory.hpp"
#include "kapi/system.hpp"

#include "x86_64/memory/page_table.hpp"
#include "x86_64/memory/page_utilities.hpp"
#include "x86_64/memory/scoped_mapping.hpp"

#include <cstddef>
#include <cstdint>
#include <memory>
#include <optional>

namespace teachos::memory::x86_64
{

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

    //! Perform the actual mapping of the page, via the recursive page map.
    //!
    //! On any level above PML1, the entries need to not be no_execute, because the image is densely packed. The entries
    //! also need to be writable, since the mapping is being performed through the recursive page map hierarchy. When
    //! setting the final entry in the PML1, the actually desired flags are set as is, with the present bit added, thus
    //! still enforcing non-writability and non-execution of the affected page.
    template<std::size_t Level>
      requires(Level > 1uz && Level < 5uz)
    auto do_map(recursive_page_table<Level> * pml, page page, page_table::entry::flags flags,
                frame_allocator & allocator)
    {
      auto index = pml_index<Level>(page);
      flags = flags & ~page_table::entry::flags::no_execute;
      flags = flags | page_table::entry::flags::writable;
      if (!(*pml)[index].present())
      {
        auto new_table_frame = allocator.allocate();
        auto mapping = scoped_mapping{page, allocator};
        (*pml)[index].frame(new_table_frame.value(), page_table::entry::flags::present | flags);
        auto new_table = std::optional{std::construct_at(*pml->next(index))};
        return new_table;
      }
      (*pml)[index] |= flags;
      return pml->next(index);
    }

    //! Perform the actual PML1 update.
    auto do_map(page_table * pml, page page, frame frame, page_table::entry::flags flags) -> std::optional<std::byte *>
    {
      auto index = pml_index<1>(page);
      if ((*pml)[index].present())
      {
        system::panic("[x86_64:MEM] Tried to map a page that is already mapped");
      }
      (*pml)[index].frame(frame, page_table::entry::flags::present | flags);
      return std::optional{static_cast<std::byte *>(page.start_address())};
    }
  }  // namespace

  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() % 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)
      {
        return std::nullopt;
      }
      else if (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 pml2 = (*pml3)->next(pml_index<3>(page));
      auto pml2_entry = pml2.transform([&](auto pml2) -> auto { return (*pml2)[pml_index<2>(page)]; });
      if (!pml2_entry)
      {
        return std::nullopt;
      }
      else if (pml2_entry->huge())
      {
        auto pml2_entry_frame = *pml2_entry->frame();
        return frame{pml2_entry_frame.number() + pml_index<1>(page)};
      }

      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);
  }

  auto paging_root::map(page page, frame frame, page_table::entry::flags flags, frame_allocator & allocator)
      -> std::optional<std::byte *>
  {
    return std::optional{this}
        .and_then([&](auto pml) -> auto { return do_map(pml, page, flags, allocator); })
        .and_then([&](auto pml) -> auto { return do_map(pml, page, flags, allocator); })
        .and_then([&](auto pml) -> auto { return do_map(pml, page, flags, allocator); })
        .and_then([&](auto pml) -> auto { return do_map(pml, page, frame, flags); });
  }

}  // namespace teachos::memory::x86_64