#include "arch/memory/higher_half_mapper.hpp"

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

#include "arch/memory/page_table.hpp"
#include "arch/memory/page_utilities.hpp"

#include <algorithm>
#include <array>
#include <cstddef>
#include <memory>
#include <ranges>
#include <utility>

namespace arch::memory
{

  higher_half_mapper::higher_half_mapper(page_table * root)
      : m_root{root}
  {}

  auto higher_half_mapper::map(kapi::memory::page page, kapi::memory::frame frame, flags flags) -> std::byte *
  {
    auto table = get_or_create_page_table(page);
    if (!table)
    {
      return nullptr;
    }

    auto const index = pml_index(1, page);
    auto & entry = (*table)[index];

    if (entry.present())
    {
      kapi::system::panic("[x86_64:MEM] Tried to map a page that is already mapped!");
    }

    entry.frame(frame, to_table_flags(flags) | page_table::entry::flags::present);

    return static_cast<std::byte *>(page.start_address());
  }

  auto higher_half_mapper::unmap(kapi::memory::page page) -> void
  {
    if (!try_unmap(page))
    {
      kapi::system::panic("[x86_64:MEM] Tried to unmap a page that is not mapped!");
    }
  }

  auto higher_half_mapper::try_unmap(kapi::memory::page page) noexcept -> bool
  {
    auto table_path = std::array<std::pair<page_table *, std::size_t>, 4>{};
    table_path[0] = std::pair{m_root, pml_index(4, page)};

    for (auto level = 4uz; level > 1uz; --level)
    {
      auto [table, index] = table_path[4 - level];
      auto & entry = (*table)[index];

      if (!entry.present())
      {
        return false;
      }

      auto next_table = to_higher_half_pointer<page_table>(entry.frame()->start_address());
      auto next_index = pml_index(4 - level - 1, page);
      table_path[4 - level - 1] = std::pair{next_table, next_index};
    }

    std::ranges::for_each(std::views::reverse(table_path), [previous_was_empty = true](auto & step) mutable {
      auto [table, index] = step;
      auto & entry = (*table)[index];
      auto frame = entry.frame();

      if (previous_was_empty)
      {
        entry.clear();
        previous_was_empty = table->empty();
        kapi::memory::get_frame_allocator().release(*frame);
      }
    });

    return true;
  }

  auto higher_half_mapper::get_or_create_page_table(kapi::memory::page page) noexcept -> page_table *
  {
    auto table = m_root;

    for (auto level = 4uz; level > 1uz; --level)
    {
      auto index = pml_index(level, page);
      auto & entry = (*table)[index];

      if (!entry.present())
      {
        auto table_frame = kapi::memory::allocate_frame();
        if (!table_frame)
        {
          return nullptr;
        }

        auto new_table = to_higher_half_pointer<page_table>(table_frame->start_address());
        std::construct_at(new_table);

        auto const flags = page_table::entry::flags::present | page_table::entry::flags::writable |
                           page_table::entry::flags::user_accessible;
        entry.frame(*table_frame, flags);
      }

      table = to_higher_half_pointer<page_table>(entry.frame()->start_address());
    }

    return table;
  }

}  // namespace arch::memory