#include "kapi/memory.hpp"

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

#include "x86_64/boot/boot.hpp"
#include "x86_64/boot/ld.hpp"
#include "x86_64/cpu/registers.hpp"
#include "x86_64/memory/mmu.hpp"
#include "x86_64/memory/page_table.hpp"
#include "x86_64/memory/paging_root.hpp"
#include "x86_64/memory/region_allocator.hpp"
#include "x86_64/memory/scoped_mapping.hpp"

#include <multiboot2/information.hpp>

#include <atomic>
#include <memory>
#include <span>

namespace teachos::memory
{
  std::size_t const PLATFORM_FRAME_SIZE{4096};
  std::size_t const PLATFORM_PAGE_SIZE{PLATFORM_FRAME_SIZE};

  namespace
  {
    // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
    auto constinit allocator = static_cast<frame_allocator *>(nullptr);

    constexpr auto static unused_page_address = 0x0000'7fff'cafe'faceuz;

    auto create_memory_information() -> x86_64::region_allocator::memory_information
    {
      auto const & mbi = boot::bootstrap_information.mbi;
      auto mbi_span = std::span{std::bit_cast<std::byte *>(mbi), mbi->size_bytes()};
      auto image_span = std::span{&boot::x86_64::_start_physical, &boot::x86_64::_end_physical};

      return {.image_range =
                  std::make_pair(physical_address{&image_span.front()}, physical_address{&image_span.back()}),
              .mbi_range = std::make_pair(physical_address{&mbi_span.front()}, physical_address{&mbi_span.back()}),
              .memory_map = mbi->memory_map()};
    };

    auto create_early_frame_allocator()
    {
      auto memory_map = boot::bootstrap_information.mbi->maybe_memory_map();
      if (!memory_map)
      {
        system::panic("[x86_64] Failed to create early allocator, no memory map available.");
      }

      return x86_64::region_allocator{create_memory_information()};
    }

    auto enable_cpu_protections() -> void
    {
      cpu::x86_64::cr0::set(cpu::x86_64::cr0::flags::write_protect);
      cpu::x86_64::i32_efer::set(cpu::x86_64::i32_efer::flags::execute_disable_bit_enable);
    }

    auto inject_faux_pml4(frame_allocator & allocator) -> void
    {
      using entry_flags = x86_64::page_table::entry::flags;
      using page_table = x86_64::page_table;

      auto temporary_page = page::containing(linear_address{unused_page_address});
      auto temporary_page_address = temporary_page.start_address();

      auto & pml4 = x86_64::paging_root::get();

      // NOLINTBEGIN(cppcoreguidelines-avoid-magic-numbers)
      auto faux_pml4_frame =
          allocator.allocate()
              .and_then([&](auto frame) -> auto {
                auto index = temporary_page_address >> 39 & 0x1ffu;
                pml4[index].frame(frame, entry_flags::present | entry_flags::writable);
                return pml4.next(index);
              })
              .and_then([&](auto pml) -> auto {
                std::construct_at(pml);
                auto index = temporary_page_address >> 30 & 0x1ffu;
                (*pml)[index].frame(*allocator.allocate(), entry_flags::present | entry_flags::writable);
                return pml->next(index);
              })
              .and_then([&](auto pml) -> auto {
                std::construct_at(pml);
                auto index = temporary_page_address >> 21 & 0x1ffu;
                (*pml)[index].frame(*allocator.allocate(), entry_flags::present | entry_flags::writable);
                return pml->next(index);
              })
              .transform([&](auto pml) -> auto {
                std::construct_at(pml);
                auto index = temporary_page_address >> 12 & 0x1ffu;
                (*pml)[index].frame(*allocator.allocate(), entry_flags::present | entry_flags::writable);
                return pml;
              })
              .and_then([&](auto pml) -> auto {
                auto faux_pml4_pointer = std::bit_cast<page_table *>(temporary_page_address.raw());
                auto faux_pml4 = std::construct_at<page_table>(faux_pml4_pointer);

                auto index = temporary_page_address >> 12 & 0x1ffu;
                auto frame = (*pml)[index].frame();

                (*faux_pml4)[510].frame(*frame, entry_flags::present | entry_flags::writable);
                return frame;
              });
      // NOLINTEND(cppcoreguidelines-avoid-magic-numbers)

      if (!faux_pml4_frame)
      {
        system::panic("[MEM] Failed to map and construct faux PML4");
      }

      // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers)
      pml4[510].frame(*faux_pml4_frame, entry_flags::present | entry_flags::writable);
      x86_64::tlb_flush_all();

      cio::println("[MEM] Injected faux PML4 as recursive map.");
    }

  }  // namespace

  auto active_allocator() -> frame_allocator &
  {
    if (!allocator)
    {
      system::panic("[x86_64] The frame allocator has not been set yet.");
    }

    return *allocator;
  }

  auto init() -> void
  {
    auto static constinit is_initialized = std::atomic_flag{};
    if (is_initialized.test_and_set())
    {
      system::panic("[x86_64] Memory management has already been initialized.");
    }

    auto allocator = create_early_frame_allocator();
    enable_cpu_protections();

    // TODO: remove
    inject_faux_pml4(allocator);

    // TODO: implement
    auto temporary_mapper = x86_64::scoped_mapping{linear_address{unused_page_address}, allocator};
    auto new_pml4_frame = allocator.allocate();

    auto new_plm4 = temporary_mapper.map_as<x86_64::page_table>(
        *new_pml4_frame, x86_64::page_table::entry::flags::present | x86_64::page_table::entry::flags::writable);
    (*new_plm4)[510].frame(new_pml4_frame.value(),
                           x86_64::page_table::entry::flags::present | x86_64::page_table::entry::flags::writable);

    // paging::kernel_mapper kernel(allocator, memory_information);
    // kernel.remap_kernel();
    // video::vga::text::write("Kernel remapping successful", video::vga::text::common_attributes::green_on_black);
    // video::vga::text::newline();

    // remap_heap(heap::KERNEL_HEAP_START, heap::KERNEL_HEAP_SIZE);
    // video::vga::text::write("Heap remapping successful", video::vga::text::common_attributes::green_on_black);
    // video::vga::text::newline();
  }

}  // namespace teachos::memory