#include "arch/kernel/main.hpp"

#include "arch/boot/pointers.hpp"
#include "arch/context_switching/interrupt_descriptor_table/segment_selector.hpp"
#include "arch/context_switching/main.hpp"
#include "arch/kernel/cpu/if.hpp"
#include "arch/kernel/cpu/segment_register.hpp"
#include "arch/memory/heap/bump_allocator.hpp"
#include "arch/memory/heap/global_heap_allocator.hpp"
#include "arch/memory/main.hpp"
#include "arch/memory/multiboot/reader.hpp"
#include "arch/stl/vector.hpp"
#include "arch/video/vga/text.hpp"

namespace teachos::arch::kernel
{
  auto stack_overflow_test(int count) -> int
  {
    int test[5000] = {};
    if (test[0] == 0xFFFF)
    {
      return count;
    }
    count = stack_overflow_test(count);
    return count++;
  }

  auto heap_test() -> void
  {
    auto test2 = new memory::multiboot::memory_information{};
    auto test3 = new memory::multiboot::memory_information{};
    auto test4 = *test2;
    auto test5 = *test3;
    test4.kernel_end = 5000;
    test5.kernel_end = 3000;
    auto test6 = test4.kernel_end;
    auto test7 = test5.kernel_end;
    auto test8 = memory::multiboot::read_multiboot2();
    if (test6 && test7 && test8.kernel_end)
    {
      video::vga::text::write("Heap test successful", video::vga::text::common_attributes::green_on_black);
    }
    test2->kernel_end = 2000;
    test2->kernel_start = 1000;
    test2->multiboot_start = 2000;
    delete test2;
    delete test3;

    auto test9 = new int(50);
    delete test9;
  }

  auto return_function() -> void
  {
    video::vga::text::write("User Mode!!!", video::vga::text::common_attributes::green_on_black);
  }

  /**
   * @brief Switch context into the mode defined in the segment selectors.
   *
   * Setup the stack IRETQ expects to switch the mode:
   * 1. push data selector
   * 2. push current stack pointer
   * 3. push eflags
   * 4. push code segment selector
   * 5. push return address
   *
   * @param data_segment
   * @param code_segment
   * @param address
   */
  [[gnu::naked]]
  auto switch_context(context_switching::interrupt_descriptor_table::segment_selector data_segment,
                      context_switching::interrupt_descriptor_table::segment_selector code_segment, uint64_t address)
      -> void
  {
    asm volatile("mov %[data_segment], %%rax\n"
                 "mov %%rax, %%ds\n"
                 "mov %%rax, %%es\n"
                 "mov %%rax, %%fs\n"
                 "mov %%rax, %%gs\n"
                 "mov %%rsp, %%rax\n"

                 "push %[data_segment]\n"
                 "push %%rax\n"
                 "pushfq\n"
                 "push %[code_segment]\n"
                 "mov %[return_function], %%rax\n"
                 "push %%rax\n"

                 "iretq\n"
                 :
                 : [data_segment] "m"(data_segment), [code_segment] "m"(code_segment), [return_function] "r"(address)
                 : "rax");
  }

  auto main() -> void
  {
    video::vga::text::clear();
    video::vga::text::cursor(false);
    video::vga::text::write("TeachOS is starting up...", video::vga::text::common_attributes::green_on_black);
    video::vga::text::newline();

    memory::initialize_memory_management();
    // stack_overflow_test(0);

    memory::heap::global_heap_allocator::register_heap_allocator(memory::heap::heap_allocator_type::LINKED_LIST);

    heap_test();

    decltype(auto) descriptor_tables = context_switching::initialize_descriptor_tables();

    context_switching::interrupt_descriptor_table::segment_selector user_code_segment_selector{
        3U, context_switching::interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_USER};
    context_switching::interrupt_descriptor_table::segment_selector user_data_segment_selector{
        4U, context_switching::interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_USER};

    cpu::set_segment_registers(user_data_segment_selector);
    switch_context(user_data_segment_selector, user_code_segment_selector,
                   reinterpret_cast<uint64_t>(&return_function));

    (void)descriptor_tables;
  }
}  // namespace teachos::arch::kernel