#include "arch/context_switching/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
  {

    /**
     * @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 far_return(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");
    }

    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};

  }  // namespace

  auto initialize_descriptor_tables() -> descriptor_tables
  {
    kernel::cpu::clear_interrupt_flag();

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

    interrupt_descriptor_table::segment_selector segment_selector{
        1U, interrupt_descriptor_table::segment_selector::REQUEST_LEVEL_KERNEL};
    kernel::cpu::far_pointer pointer{&kernel::cpu::reload_segment_registers, segment_selector};
    kernel::cpu::call(pointer);

    segment_descriptor_table::update_task_state_segment_register();

    kernel::cpu::set_interrupt_flag();

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

  auto user_mode_main() -> void
  {
    auto current_segment = kernel::cpu::read_code_segment_register();
    exception_handling::assert(USER_CODE_SEGMENT_SELECTOR == current_segment,
                               "[Context Switching] Context switch into user mode not successful");
    exception_handling::assert(USER_DATA_SEGMENT_SELECTOR == kernel::cpu::validate_data_segment_registers(),
                               "[Context Switching] Context switch into user mode not successful");

    video::vga::text::write("Successfully entered user mode!", video::vga::text::common_attributes::green_on_black);
  }

  auto switch_to_user_mode() -> void
  {
    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
  {
    kernel::cpu::set_segment_registers(data_segment);
    far_return(data_segment, code_segment, reinterpret_cast<uint64_t>(return_function));
  }

}  // namespace teachos::arch::context_switching