#include "arch/memory/multiboot/reader.hpp"

#include "arch/boot/pointers.hpp"
#include "arch/exception_handling/assert.hpp"
#include "arch/memory/multiboot/elf_symbols_section.hpp"
#include "arch/memory/multiboot/info.hpp"

#include <algorithm>

namespace teachos::arch::memory::multiboot
{
  namespace
  {
    template<typename T>
      requires std::is_pointer<T>::value
    auto align_to_8_byte_boundary(T ptr, uint32_t size) -> T
    {
      return reinterpret_cast<T>(reinterpret_cast<uint8_t *>(ptr) + ((size + 7) & ~7));
    }

    auto process_memory_map(memory_map_header * mminfo) -> memory_area_container
    {
      auto const expected_entry_size = mminfo->entry_size;
      auto constexpr actual_entry_size = sizeof(memory_area);
      exception_handling::assert(expected_entry_size == actual_entry_size,
                                 "[Multiboot Reader] Unexpected memory area entry size");

      auto const total_size = mminfo->info.size;
      auto const total_entries_size = total_size - sizeof(memory_map_header) + actual_entry_size;
      auto const number_of_entries = total_entries_size / actual_entry_size;

      auto const begin = memory_area_container::iterator{&mminfo->entries};
      auto const end = begin + number_of_entries;
      return memory_area_container{begin, end};
    }

    auto process_elf_sections(elf_symbols_section_header * symbol, std::size_t & kernel_start,
                              std::size_t & kernel_end) -> elf_section_header_container
    {
      auto const expected_entry_size = symbol->entry_size;
      auto constexpr actual_entry_size = sizeof(elf_section_header);
      exception_handling::assert(expected_entry_size == actual_entry_size,
                                 "[Multiboot Reader] Unexpected elf section header entry size");

      auto const expected_total_size = symbol->info.size;
      auto const actual_total_entry_size = actual_entry_size * symbol->number_of_sections;
      auto constexpr actual_total_section_size = sizeof(elf_symbols_section_header) - sizeof(uint32_t);
      auto const actual_total_size = actual_total_entry_size + actual_total_section_size;
      exception_handling::assert(expected_total_size == actual_total_size,
                                 "[Multiboot Reader] Unexpected elf symbols section header total size");

      auto const begin = elf_section_header_container::iterator{reinterpret_cast<elf_section_header *>(&symbol->end)};
      auto const end = begin + symbol->number_of_sections;
      exception_handling::assert(begin->is_null(),
                                 "[Multiboot Reader] Elf symbols section not starting with SHT_NULL section");

      elf_section_header_container sections{begin, end};

      auto const elf_section_with_lowest_physical_address =
          std::ranges::min_element(sections, [](elf_section_header const & a, elf_section_header const & b) {
            return a.physical_address < b.physical_address;
          });

      auto const elf_section_with_highest_physical_address =
          std::ranges::max_element(sections, [](elf_section_header const & a, elf_section_header const & b) {
            auto a_physical_address_end = a.physical_address + a.section_size;
            auto b_physical_address_end = b.physical_address + b.section_size;
            return a_physical_address_end < b_physical_address_end;
          });

      auto const symbol_table_section_count = std::ranges::count_if(sections, [](elf_section_header const & section) {
        return section.type == elf_section_type::DYNAMIC_SYMBOL_TABLE || section.type == elf_section_type::SYMBOL_TABLE;
      });
      auto const dynamic_section_count = std::ranges::count_if(
          sections, [](elf_section_header const & section) { return section.type == elf_section_type::DYNAMIC; });

      exception_handling::assert(
          symbol_table_section_count == 1U,
          "[Multiboot Reader] ELF Specifications allows only (1) symbol table section, but got more");
      exception_handling::assert(
          dynamic_section_count <= 1U,
          "[Multiboot Reader] ELF Specifications allows only (1) or less dynamic sections, but got more");

      auto const lowest_elf_section = *elf_section_with_lowest_physical_address;
      kernel_start = lowest_elf_section.physical_address;

      auto const highest_elf_section = *elf_section_with_highest_physical_address;
      kernel_end = highest_elf_section.physical_address + highest_elf_section.section_size;

      return sections;
    }
  }  // namespace

  auto read_multiboot2() -> memory_information
  {
    memory_information mem_info{UINT64_MAX,
                                0U,
                                elf_section_header_container{},
                                boot::multiboot_information_pointer,
                                0U,
                                memory_area_container{}};

    auto const multiboot_information_pointer = reinterpret_cast<info_header *>(boot::multiboot_information_pointer);
    auto const multiboot_tag = &multiboot_information_pointer->tags;
    mem_info.multiboot_end = mem_info.multiboot_start + multiboot_information_pointer->total_size;

    for (auto tag = multiboot_tag; tag->type != tag_type::END; tag = align_to_8_byte_boundary(tag, tag->size))
    {
      switch (tag->type)
      {
        case tag_type::ELF_SECTIONS: {
          auto const symbol = reinterpret_cast<elf_symbols_section_header *>(tag);
          mem_info.sections = process_elf_sections(symbol, mem_info.kernel_start, mem_info.kernel_end);
          break;
        }
        case tag_type::MEMORY_MAP: {
          auto const mminfo = reinterpret_cast<memory_map_header *>(tag);
          mem_info.areas = process_memory_map(mminfo);
          break;
        }
        default:
          // All other cases are not important and can be ignored.
          break;
      }
    }
    return mem_info;
  }
}  // namespace teachos::arch::memory::multiboot