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

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

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 *& memory_areas, uint8_t & area_count) -> void
    {
      auto expected_entry_size = mminfo->entry_size;
      constexpr auto actual_entry_size = sizeof(memory_area);
      exception_handling::assert(expected_entry_size == actual_entry_size,
                                 "[Multiboot Reader] Unexpected memory area entry size");

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

      memory_areas = &mminfo->entries;
      area_count = number_of_entries;
    }

    auto process_elf_sections(elf_symbols_section_header * symbol, uint64_t & kernel_start,
                              uint64_t & kernel_end) -> void
    {
      auto expected_entry_size = symbol->entry_size;
      constexpr auto 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 expected_total_size = symbol->info.size;
      auto actual_total_entry_size = actual_entry_size * symbol->number_of_sections;
      constexpr auto actual_total_section_size = sizeof(elf_symbols_section_header) - sizeof(uint32_t);
      auto 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 begin = reinterpret_cast<elf_section_header *>(&symbol->end);
      auto end = begin + symbol->number_of_sections;
      exception_handling::assert(begin->is_null(),
                                 "[Multiboot Reader] Elf symbols section not starting with SHT_NULL section");

      std::size_t symbol_table_section_count = 0U;
      std::size_t dynamic_section_count = 0U;

      for (auto section = begin; section != end; ++section)
      {
        switch (section->type)
        {
          case elf_section_type::PROGRAMM: {
            if (section->virtual_address < kernel_start)
            {
              kernel_start = section->virtual_address;
            }
            auto virtual_address_end = section->virtual_address + section->section_size;
            if (virtual_address_end > kernel_end)
            {
              kernel_end = virtual_address_end;
            }
            break;
          }
          case elf_section_type::DYNAMIC_SYMBOL_TABLE:
          case elf_section_type::SYMBOL_TABLE:
            symbol_table_section_count++;
            break;
          case elf_section_type::DYNAMIC:
            dynamic_section_count++;
            break;
          default:
            // All other cases are not important and can be ignored
            break;
        }
      }

      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");
    }
  }  // namespace

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

    auto * multiboot_information_pointer = reinterpret_cast<info_header *>(boot::multiboot_information_pointer);
    auto 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 symbol = reinterpret_cast<elf_symbols_section_header *>(tag);
          process_elf_sections(symbol, mem_info.kernel_start, mem_info.kernel_end);
          break;
        }
        case tag_type::MEMORY_MAP: {
          auto mminfo = reinterpret_cast<memory_map_header *>(tag);
          process_memory_map(mminfo, mem_info.memory_areas, mem_info.area_count);
          break;
        }
        default:
          // All other cases are not important and can be ignored
          break;
      }
    }

    return mem_info;
  }
}  // namespace teachos::arch::memory::multiboot