#include "arch/context_switching/main.hpp"

#include "arch/boot/pointers.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/msr.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
  {
    auto constexpr IA32_STAR_ADDRESS = 0xC0000081;
    auto constexpr IA32_LSTAR_ADDRESS = 0xC0000082;
    auto constexpr IA32_FMASK_ADDRESS = 0xC0000084;

    constexpr interrupt_descriptor_table::segment_selector KERNEL_CODE_SEGMENT_SELECTOR{
        1U, interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_KERNEL};
    constexpr interrupt_descriptor_table::segment_selector KERNEL_DATA_SEGMENT_SELECTOR{
        2U, 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_global_descriptor_table_register() -> void { kernel::cpu::call(KERNEL_CODE_POINTER); }

    auto user_mode_main() -> void
    {
      video::vga::text::write("Successfully entered user mode!", video::vga::text::common_attributes::green_on_black);

      uint64_t new_value = 60U;
      asm volatile("mov %[input], %%rax"
                   : /* no output from call */
                   : [input] "r"(new_value)
                   : "memory");
      asm volatile("syscall");

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

    auto syscall_handler() -> void
    {
      uint64_t syscall_number{};
      asm volatile("mov %%rax, %[output]" : [output] "=r"(syscall_number));

      switch (syscall_number)
      {
        case 1:
          // Write VGA
          break;
        case 60U:
          // Exit
          break;
        default:
          break;
      }

      asm volatile("sysretq");
    }

    auto enable_systemcall() -> void
    {
      uint64_t const syscall_function = reinterpret_cast<uint64_t>(syscall_handler);
      kernel::cpu::write_msr(IA32_LSTAR_ADDRESS, syscall_function);
      kernel::cpu::write_msr(IA32_FMASK_ADDRESS, 1 << 9U);  // Disable interrupt flag during syscall.

      uint64_t kernel_cs = KERNEL_CODE_SEGMENT_SELECTOR;
      // We want to provide the user code segment, but the instruction calculates + 0x10 to fill the
      // cs register (See https://www.felixcloutier.com/x86/sysret).
      uint64_t user_cs = USER_CODE_SEGMENT_SELECTOR - 0x10;
      uint64_t star_value = (kernel_cs << 32) | (user_cs << 48);
      kernel::cpu::write_msr(IA32_STAR_ADDRESS, star_value);

      kernel::cpu::set_efer_bit(kernel::cpu::efer_flags::SCE);
    }

  }  // namespace

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

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

      segment_descriptor_table::update_global_descriptor_table_register();
      interrupt_descriptor_table::update_interrupt_descriptor_table_register();

      reload_global_descriptor_table_register();
      segment_descriptor_table::update_task_state_segment_register();

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

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

  auto switch_to_user_mode() -> void
  {
    enable_systemcall();
    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