#include "arch/context_switching/main.hpp"

#include "arch/boot/pointers.hpp"
#include "arch/context_switching/syscall/main.hpp"
#include "arch/exception_handling/assert.hpp"
#include "arch/kernel/cpu/call.hpp"
#include "arch/kernel/cpu/control_register.hpp"
#include "arch/kernel/cpu/if.hpp"
#include "arch/kernel/cpu/segment_register.hpp"
#include "arch/kernel/cpu/tr.hpp"
#include "arch/video/vga/text.hpp"

namespace teachos::arch::context_switching
{
  namespace
  {
    constexpr interrupt_descriptor_table::segment_selector KERNEL_CODE_SEGMENT_SELECTOR{
        1U, interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_KERNEL};
    constexpr kernel::cpu::far_pointer KERNEL_CODE_POINTER{&kernel::cpu::reload_data_segment_registers,
                                                           KERNEL_CODE_SEGMENT_SELECTOR};
    constexpr context_switching::interrupt_descriptor_table::segment_selector USER_CODE_SEGMENT_SELECTOR{
        3U, context_switching::interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_USER};
    constexpr context_switching::interrupt_descriptor_table::segment_selector USER_DATA_SEGMENT_SELECTOR{
        4U, context_switching::interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_USER};

    auto reload_gdtr() -> void { kernel::cpu::call(KERNEL_CODE_POINTER); }

    auto user_mode_main() -> void
    {
      // TODO: Remove or replace this wall of text
      // RFLAGS is saved into R11, RIP of the next instruction into RCX
      // Required for SYSRETURN to know where to return too.
      // Additional state needs to be saved by calling convention:

      // Syscall Number: RAX, Return Value: RAX (0 indicating no error, and -1 indicating an error, use as a boolean)
      // Argument in this order (max 6. no argument on stack): RDI, RSI, RDX, R10, R8, R9
      // Not used registers: RBX, RSP, R12, R13, R14

      // Actual Source: https://man7.org/linux/man-pages/man2/syscall.2.html More cleare documentation:
      // https://sys.readthedocs.io/en/latest/doc/05_calling_system_calls.html

      const char syscall_message[68] = "Successfully entered user mode and wrote to VGA buffer via syscall!";
      auto error = syscall::syscall(syscall::WRITE, {reinterpret_cast<uint64_t>(&syscall_message)});

      if (!error)
      {
        video::vga::text::write("Successfully made a SYSCALL and returned with SYSRETQ!",
                                video::vga::text::common_attributes::green_on_black);
      }
    }
  }  // namespace

  auto initialize_descriptor_tables() -> descriptor_tables
  {
    static bool initalized = false;

    if (!initalized)
    {
      kernel::cpu::clear_interrupt_flag();

      segment_descriptor_table::update_gdtr();
      interrupt_descriptor_table::update_interrupt_descriptor_table_register();

      reload_gdtr();
      segment_descriptor_table::update_tss_register();

      kernel::cpu::set_interrupt_flag();
      initalized = true;
    }

    descriptor_tables tables = {segment_descriptor_table::get_or_create_gdt(),
                                interrupt_descriptor_table::get_or_create_interrupt_descriptor_table()};
    return tables;
  }

  auto switch_to_user_mode() -> void
  {
    syscall::enable_syscall();
    switch_context(USER_DATA_SEGMENT_SELECTOR, USER_CODE_SEGMENT_SELECTOR, user_mode_main);
  }

  auto switch_context(interrupt_descriptor_table::segment_selector data_segment,
                      interrupt_descriptor_table::segment_selector code_segment, void (*return_function)()) -> void
  {
    (void)initialize_descriptor_tables();
    kernel::cpu::set_data_segment_registers(data_segment);
    kernel::cpu::set_code_segment_register(data_segment, code_segment, reinterpret_cast<uint64_t>(return_function));
  }
}  // namespace teachos::arch::context_switching