#include "arch/context_switching/main.hpp"

#include "arch/boot/pointers.hpp"
#include "arch/exception_handling/assert.hpp"
#include "arch/kernel/cpu/if.hpp"
#include "arch/kernel/cpu/jmp.hpp"
#include "arch/kernel/cpu/tr.hpp"

namespace teachos::arch::context_switching
{
  auto initialize_descriptor_tables() -> descriptor_tables
  {
    decltype(auto) global_descriptor_table = segment_descriptor_table::initialize_global_descriptor_table();
    decltype(auto) interrupt_descriptor_table = interrupt_descriptor_table::initialize_interrupt_descriptor_table();

    // Execute trampoline function for the GDT loading long jump
    boot::reload_segment_register_trampoline();

    // Load task state segment descriptor from the last element in the global descriptor table, done by calculating
    // offset in bytes to the start of the segment descriptor (5 * 16) = 80
    uint16_t const tss_selector =
        (global_descriptor_table.size() - 1) * sizeof(segment_descriptor_table::segment_descriptor);
    kernel::cpu::load_task_register(tss_selector);

    auto const stored_task_register = kernel::cpu::store_task_register();
    arch::exception_handling::assert(tss_selector == stored_task_register,
                                     "[Global Descriptor Table] Loaded TR value is not the same as the stored value.");

    kernel::cpu::set_interrupt_flag();

    descriptor_tables tables = {global_descriptor_table, interrupt_descriptor_table};
    return tables;
  }
}  // namespace teachos::arch::context_switching