#include "arch/memory/kernel_mapper.hpp"

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

#include "arch/boot/ld.hpp"

#include <kstd/print>
#include <kstd/units>

#include <elf/format.hpp>
#include <elf/section_header.hpp>
#include <multiboot2/information.hpp>

#include <algorithm>
#include <array>
#include <bit>
#include <cstdint>
#include <ranges>
#include <string_view>
#include <utility>

using namespace std::string_view_literals;
using namespace kstd::units_literals;

namespace arch::memory
{

  namespace
  {
    constexpr auto static ignored_section_prefixes = std::array{
      ".boot_"sv,
    };

    constexpr auto static user_accessible_prefixes = std::array{
      ".user"sv,
      ".stl"sv,
    };

  }  // namespace

  kernel_mapper::kernel_mapper(multiboot2::information_view const * mbi)
      : m_mbi{std::move(mbi)}
      , m_kernel_load_base{std::bit_cast<std::uintptr_t>(&arch::boot::TEACHOS_VMA)}
  {}

  auto kernel_mapper::remap_kernel(kapi::memory::page_mapper & mapper) -> void
  {
    auto elf_information = m_mbi->maybe_elf_symbols<elf::format::elf64>();
    if (!elf_information)
    {
      kapi::system::panic("[x86_64:MEM] ELF section information is not available.");
    }

    auto sections = *elf_information;
    auto allocated_sections =
        std::views::all(sections) | std::views::filter(&elf::section_header<elf::format::elf64>::allocated) |
        std::views::filter([&](auto const & section) -> auto {
          auto name = sections.name(section);
          return !std::ranges::any_of(ignored_section_prefixes,
                                      [&](auto const & prefix) -> auto { return name.starts_with(prefix); });
        });

    if (allocated_sections.empty())
    {
      kapi::system::panic("[x86_64:MEM] No allocated ELF sections were found.");
    }

    std::ranges::for_each(allocated_sections,
                          [&](auto const & section) -> auto { map_section(section, sections.name(section), mapper); });
  }

  auto kernel_mapper::map_section(section_header_type const & section, std::string_view name,
                                  kapi::memory::page_mapper & mapper) -> void
  {
    auto number_of_pages =
        (kstd::units::bytes{section.size} + (kapi::memory::page::size - 1_B)) / kapi::memory::page::size;
    auto linear_start_address = kapi::memory::linear_address{section.virtual_load_address};
    auto physical_start_address = kapi::memory::physical_address{section.virtual_load_address & ~m_kernel_load_base};

    kstd::println("[x86_64:MEM] mapping {}"
                  "\n  {} bytes -> page count: {}"
                  "\n  {} @ {}",
                  name, section.size, number_of_pages, linear_start_address, physical_start_address);

    auto first_page = kapi::memory::page::containing(linear_start_address);
    auto first_frame = kapi::memory::frame::containing(physical_start_address);

    auto page_flags = kapi::memory::page_mapper::flags::empty;

    if (section.writable())
    {
      page_flags |= kapi::memory::page_mapper::flags::writable;
    }

    if (section.executable())
    {
      page_flags |= kapi::memory::page_mapper::flags::executable;
    }

    auto is_prefix_of_name = [=](auto prefix) -> bool {
      return name.starts_with(prefix);
    };

    if (!std::ranges::any_of(user_accessible_prefixes, is_prefix_of_name))
    {
      page_flags |= kapi::memory::page_mapper::flags::supervisor_only;
    }

    for (auto i = 0uz; i < number_of_pages; ++i)
    {
      mapper.map(first_page + i, first_frame + i, page_flags);
    }
  }

}  // namespace arch::memory