#include "arch/memory/region_allocator.hpp"

#include "kapi/memory/address.hpp"
#include "kapi/memory/frame.hpp"

#include <multiboot2/information.hpp>

#include <algorithm>
#include <cstddef>
#include <optional>
#include <ranges>
#include <utility>

namespace arch::memory
{
  namespace
  {
    constexpr auto last_frame(multiboot2::memory_map::region const & region)
    {
      return kapi::memory::frame::containing(kapi::memory::physical_address{region.base + region.size_in_B - 1});
    }

    constexpr auto falls_within(kapi::memory::frame const & candidate, kapi::memory::frame const & start,
                                kapi::memory::frame const & end)
    {
      return candidate >= start && candidate <= end;
    }
  }  // namespace

  region_allocator::region_allocator(memory_information const & mem_info)
      : m_next_frame{}
      , m_current_region{}
      , m_memory_map{mem_info.memory_map}
      , m_kernel_start{kapi::memory::frame::containing(mem_info.image_range.first)}
      , m_kernel_end{kapi::memory::frame::containing(mem_info.image_range.second)}
      , m_multiboot_start{kapi::memory::frame::containing(mem_info.mbi_range.first)}
      , m_multiboot_end{kapi::memory::frame::containing(mem_info.mbi_range.second)}
      , m_multiboot_information{mem_info.mbi}
  {
    choose_next_region();
  }

  auto region_allocator::choose_next_region() -> void
  {
    m_current_region.reset();

    auto remaining_regions =
        m_memory_map |  //
        std::views::filter(&multiboot2::memory_map::region::available) |
        std::views::filter([this](auto const & region) -> bool { return last_frame(region) >= m_next_frame; });

    auto lowest_region =
        std::ranges::min_element(remaining_regions, [](auto lhs, auto rhs) -> bool { return lhs.base < rhs.base; });

    if (lowest_region == remaining_regions.end())
    {
      return;
    }

    m_current_region = *lowest_region;
    if (auto start_of_region = kapi::memory::frame::containing(kapi::memory::physical_address{m_current_region->base});
        start_of_region > m_next_frame)
    {
      m_next_frame = start_of_region;
    }
  }

  auto region_allocator::find_next_frame() -> std::optional<kapi::memory::frame>
  {
    if (!m_current_region || m_next_frame > last_frame(*m_current_region))
    {
      choose_next_region();
      if (!m_current_region)
      {
        return std::nullopt;
      }
    }

    auto advanced = true;
    while (advanced)
    {
      advanced = false;

      if (falls_within(m_next_frame, m_kernel_start, m_kernel_end))
      {
        m_next_frame = m_kernel_end + 1;
        advanced = true;
        continue;
      }

      if (falls_within(m_next_frame, m_multiboot_start, m_multiboot_end))
      {
        m_next_frame = m_multiboot_end + 1;
        advanced = true;
        continue;
      }

      if (m_multiboot_information)
      {
        for (auto const & module : m_multiboot_information->modules())
        {
          auto module_start = kapi::memory::frame::containing(kapi::memory::physical_address{module.start_address});
          auto module_end = kapi::memory::frame::containing(kapi::memory::physical_address{module.end_address});

          if (falls_within(m_next_frame, module_start, module_end))
          {
            m_next_frame = module_end + 1;
            advanced = true;
            break;
          }
        }
      }
    }

    if (m_next_frame > last_frame(*m_current_region))
    {
      choose_next_region();
    }

    return m_current_region.transform([this](auto) -> auto { return m_next_frame; });
  }

  auto region_allocator::allocate_many(std::size_t count) noexcept
      -> std::optional<std::pair<kapi::memory::frame, std::size_t>>
  {
    while (m_current_region)
    {
      auto result = find_next_frame();

      if (result)
      {
        auto region_end = last_frame(*m_current_region);
        if ((region_end.number() - result->number()) >= count)
        {
          m_next_frame = m_next_frame + count;
          return std::make_pair(*result, count);
        }
        else
        {
          m_next_frame = region_end + 1;
          choose_next_region();
        }
      }
    }

    return std::nullopt;
  }

  auto region_allocator::mark_used(kapi::memory::frame frame) -> void
  {
    if (frame < m_next_frame)
    {
      m_next_frame = frame;
      find_next_frame();
    }
  }

  auto region_allocator::release_many(std::pair<kapi::memory::frame, std::size_t>) -> void {}

  auto region_allocator::next_free_frame() noexcept -> std::optional<kapi::memory::frame>
  {
    return find_next_frame();
  }

}  // namespace arch::memory