#include "arch/kernel/cpu/segment_register.hpp"

#include "arch/context_switching/interrupt_descriptor_table/segment_selector.hpp"
#include "arch/exception_handling/assert.hpp"

namespace teachos::arch::kernel::cpu
{
  [[gnu::naked]]
  auto reload_segment_registers() -> void
  {
    asm volatile("xor %rax, %rax\n"
                 "mov %rax, %ss\n"
                 "mov %rax, %ds\n"
                 "mov %rax, %es\n"
                 "mov %rax, %fs\n"
                 "mov %rax, %gs\n"
                 "ret");
  }

  [[gnu::naked]]
  auto set_segment_registers(context_switching::interrupt_descriptor_table::segment_selector segment_selector) -> void
  {
    asm volatile("xor %%rax, %%rax\n"
                 "mov %[input], %%ax\n"
                 "mov %%rax, %%ds\n"
                 "mov %%rax, %%es\n"
                 "mov %%rax, %%fs\n"
                 "mov %%rax, %%gs\n"
                 "ret"
                 : /* No output from call */
                 : [input] "m"(segment_selector));
  }

  auto read_code_segment_register() -> context_switching::interrupt_descriptor_table::segment_selector
  {
    context_switching::interrupt_descriptor_table::segment_selector current_value;
    asm volatile("mov %%cs, %[output]" : [output] "=r"(current_value));
    return current_value;
  }

  auto validate_data_segment_registers() -> context_switching::interrupt_descriptor_table::segment_selector
  {
    context_switching::interrupt_descriptor_table::segment_selector ss;
    context_switching::interrupt_descriptor_table::segment_selector ds;
    context_switching::interrupt_descriptor_table::segment_selector es;
    context_switching::interrupt_descriptor_table::segment_selector fs;
    context_switching::interrupt_descriptor_table::segment_selector gs;

    asm volatile(
        "mov %%ss, %[ss_output]\n"
        "mov %%ds, %[ds_output]\n"
        "mov %%es, %[es_output]\n"
        "mov %%fs, %[fs_output]\n"
        "mov %%gs, %[gs_output]\n"
        : [ss_output] "=r"(ss), [ds_output] "=r"(ds), [es_output] "=r"(es), [fs_output] "=r"(fs), [gs_output] "=r"(gs));

    auto result = ss == ds && ss == es && ss == fs && ss == gs;
    exception_handling::assert(result, "[Segment Register] Values in data register are not the same.");

    return ss;
  }

}  // namespace teachos::arch::kernel::cpu