#include "x86_64/memory/scoped_mapping.hpp"

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

#include "x86_64/memory/mmu.hpp"
#include "x86_64/memory/page_table.hpp"
#include "x86_64/memory/page_utilities.hpp"
#include "x86_64/memory/paging_root.hpp"

#include <memory>
#include <utility>

namespace teachos::memory::x86_64
{

  scoped_mapping::scoped_mapping(scoped_mapping && other) noexcept
      : m_page{std::exchange(other.m_page, page{})}
      , m_allocator{std::exchange(other.m_allocator, nullptr)}
      , m_mapped{std::exchange(other.m_mapped, false)}
      , m_allocated{std::exchange(other.m_allocated, 0)}
  {}

  scoped_mapping::scoped_mapping(page page, frame_allocator & allocator)
      : m_page{page}
      , m_allocator{&allocator}
      , m_mapped{false}
      , m_allocated{}
  {
    if (paging_root::get().translate(page))
    {
      system::panic("[MEM] Tried to map a page that is already mapped!");
    }
  }

  scoped_mapping::~scoped_mapping()
  {
    if (m_mapped)
    {
      unmap();
      x86_64::tlb_flush(m_page.start_address());
    }
  }

  auto scoped_mapping::operator=(scoped_mapping && other) noexcept -> scoped_mapping &
  {
    if (&other == this)
    {
      return *this;
    }

    using std::swap;

    swap(m_page, other.m_page);
    swap(m_allocator, other.m_allocator);
    swap(m_mapped, other.m_mapped);
    swap(m_allocated, other.m_allocated);

    return *this;
  }

  auto scoped_mapping::map(frame frame, page_table::entry::flags flags) -> std::byte *
  {
    auto & pml4 = paging_root::get();
    auto pml4_index = pml_index<4>(m_page);
    if (!pml4[pml4_index].present())
    {
      auto new_frame = m_allocator->allocate();
      pml4[pml4_index].frame(*new_frame, page_table::entry::flags::present | flags);
      std::construct_at(pml4.next(pml4_index).value());
      m_allocated |= 1uz << 2;
    }

    auto pml3 = pml4.next(pml4_index).value();
    auto pml3_index = pml_index<3>(m_page);
    if (!(*pml3)[pml3_index].present())
    {
      auto new_frame = m_allocator->allocate();
      (*pml3)[pml3_index].frame(*new_frame, page_table::entry::flags::present | flags);
      std::construct_at((*pml3).next(pml3_index).value());
      m_allocated |= 1uz << 1;
    }

    auto pml2 = (*pml3).next(pml3_index).value();
    auto pml2_index = pml_index<2>(m_page);
    if (!(*pml2)[pml2_index].present())
    {
      auto new_frame = m_allocator->allocate();
      (*pml2)[pml2_index].frame(*new_frame, page_table::entry::flags::present | flags);
      std::construct_at((*pml2).next(pml2_index).value());
      m_allocated |= 1uz << 0;
    }

    auto pml1 = (*pml2).next(pml2_index).value();
    auto pml1_index = pml_index<1>(m_page);
    (*pml1)[pml1_index].frame(frame, page_table::entry::flags::present | flags);

    m_mapped = true;

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

  auto scoped_mapping::unmap() -> void
  {
    if (!m_mapped)
    {
      system::panic("[MEM] Tried to release an unmapped temporary mapping!");
    }

    auto pml3 = paging_root::get().next(pml_index<4>(m_page)).value();
    auto pml2 = pml3->next(pml_index<3>(m_page)).value();
    auto pml1 = pml2->next(pml_index<2>(m_page)).value();

    if (m_allocated & 1uz << 0)
    {
      auto pml1_entry = (*pml1)[pml_index<1>(m_page)];
      (*pml1)[pml_index<1>(m_page)].clear();
      if (pml1->empty())
      {
        m_allocator->release(pml1_entry.frame().value());
      }
    }

    if (m_allocated & 1uz << 1)
    {
      auto pml2_entry = (*pml2)[pml_index<2>(m_page)];
      (*pml2)[pml_index<2>(m_page)].clear();
      if (pml2->empty())
      {
        m_allocator->release(pml2_entry.frame().value());
      }
    }

    if (m_allocated & 1uz << 2)
    {
      auto pml3_entry = (*pml3)[pml_index<3>(m_page)];
      (*pml3)[pml_index<3>(m_page)].clear();
      if (pml3->empty())
      {
        m_allocator->release(pml3_entry.frame().value());
      }
    }

    m_mapped = false;
  }

}  // namespace teachos::memory::x86_64