#include "x86_64/boot/boot.hpp"

/**
 * @brief Uninitialized data for the bootstrapping process.
 */
.section .boot_bss, "aw", @nobits

/**
 * @brief Storage for the multiboot2 information pointer.
 */
.global multiboot_information_pointer
multiboot_information_pointer: .skip 8

.align 4096

page_maps_start:
page_map_level_4:      .skip 512 * 8
page_map_level_3_high: .skip 512 * 8
page_map_level_3_low:  .skip 512 * 8
page_map_level_2:      .skip 512 * 8
page_maps_end = .
page_maps_size = page_maps_end - page_maps_start

/**
 * @brief Storage for the bootstrap stack.
 */
.section .boot_stack, "aw", @nobits
.align 16

early_stack_bottom: .skip 1 << 8
early_stack_top:
early_stack_size = early_stack_top - early_stack_bottom

/**
 * @brief Constants for the bootstrapping process.
 */
.section .boot_rodata, "a", @progbits

.global global_descriptor_table_data

/**
 * @brief A basic GDT for long mode.
 */
global_descriptor_table:
global_descriptor_table_null = . - global_descriptor_table
.quad 0
global_descriptor_table_code = . - global_descriptor_table
.quad GDT_READ_WRITE | GDT_EXECUTABLE | GDT_DESCRIPTOR_TYPE | GDT_PRESENT | GDT_LONG_MODE | (1 << 55)
global_descriptor_table_data = . - global_descriptor_table
.quad GDT_READ_WRITE | GDT_DESCRIPTOR_TYPE | GDT_PRESENT | (1 << 54) | (1 << 55)
global_descriptor_table_end:

message_prefix_panic:                    .string  "Panic: "
message_not_loaded_by_multiboot2:        .string  "The operating system was not loaded by a Multiboot 2 loader."
message_cpuid_instruction_no_supported:  .string  "The 'cpuid' instruction is not supported on this platform."
message_long_mode_not_supported:         .string  "Long mode is not supported by this platform."
message_return_from_kernel_main:         .string  "Execution returned from kernel main."

/**
 * @brief Initialized data for the bootstrapping process.
 */
.section .boot_data, "aw", @progbits

/**
 * @brief A pointer to the current position within the VGA text buffer.
 */
.global vga_buffer_pointer
vga_buffer_pointer: .quad 0xb8000

/**
 * @brief Code for the bootstrapping process.
 */
.section .boot_text, "ax", @progbits
.align 16
.code32

.macro pie_base
    push %esi
    call 0f
    0:
    pop %esi
.endm

.macro function_start
    push %ebp
    mov  %esp, %ebp
.endm

.macro function_end
    leave
    ret
.endm

.macro pie_function_start
    function_start
    pie_base
.endm

.macro pie_function_end
    pop %esi
    function_end
.endm

/**
 * @brief Prepare the environment and start the kernel.
 *
 * This function performs all necessary checks to ensure the system was loaded
 * by the expected loader and supports all features required to run the kernel.
 * If successful, it prepares the system by setting up memory virtualization
 * and then start the kernel proper.
 *
 * @param %eax The Multiboot 2 magic marker.
 * @param %ebx The Multiboot 2 information pointer.
 * @return void This function does not return.
 */
.global _start
_start:
    call 0f
0:
    pop %esi

    lea (early_stack_top - 0b)(%esi), %ecx

    mov %ecx, %esp
    mov %esp, %ebp

    call _assert_loaded_by_multiboot2_loader
    call _save_multiboot_information_pointer

    call _assert_cpuid_instruction_is_supported
    call _assert_cpu_supports_long_mode

    push $HUGE_PAGES_TO_MAP
    call _prepare_page_maps
    add $4, %esp

    call _enable_paging
    call _enable_sse
    call _reload_gdt

    lea (_entry64 - 0b)(%esi), %eax
    pushl $global_descriptor_table_code
    pushl %eax
    lret

/**
 * @brief Halt the system.
 *
 * This function will instruct the CPU to halt. It will try to keep the CPU
 * halted, even if interrupts occur.
 *
 * @return This function never returns.
 */
_halt:
    function_start

1:
    hlt
    jmp 1b

    function_end

/**
 * @brief Print a message via the VGA text buffer.
 *
 * @param ebp+12 The message to print.
 * @param ebp+8  The color to print the message in.
 */
_print:
    pie_function_start

    push %edi
    push %ebx

    mov 8(%ebp), %al
    mov 12(%ebp), %edx
    mov $0, %ecx
    lea (vga_buffer_pointer - 0b)(%esi), %edi
    mov (%edi), %edi

1:
    mov (%edx, %ecx), %bl
    test %bl, %bl
    je 2f
    mov %bl, (%edi, %ecx, 2)
    mov %al, 1(%edi, %ecx, 2)
    inc %ecx
    jmp 1b

2:
    shl $1, %ecx
    add %ecx, %edi
    lea (vga_buffer_pointer - 0b)(%esi), %ecx
    mov %edi, (%ecx)

    pop %ebx
    pop %edi

    pie_function_end

/**
 * @brief Print a given panic message and then halt the machine as if by calling ::halt()
 *
 * @param ebp+4 A message to print.
 * @return This function does not return.
 */
_panic:
    pie_function_start

    lea (message_prefix_panic - 0b)(%esi), %eax

    push %eax
    push $0x4f
    call _print

    mov  16(%ebp), %eax
    mov  %eax,     8(%ebp)
    call _print
    add $8, %esp

    call _halt

    pie_function_end

/**
 * Assert that we were loaded by a Multiboot 2 compliant bootloader.
 *
 * This assertion will panic the system if the magic signature was not found.
 * If we were loaded my an appropriate bootloader, this function also saves
 * the provided MBI pointer to `multiboot_information_pointer`.
 */
_assert_loaded_by_multiboot2_loader:
    pie_function_start

    cmp $MULTIBOOT2_MAGIC, %eax
    je 1f
    lea (message_not_loaded_by_multiboot2 - 0b)(%esi), %eax
    push %eax
    call _panic
1:
    pie_function_end

/**
 * @brief Store the multiboot 2 information pointer in the global memory.
 *
 * @return void
 */
_save_multiboot_information_pointer:
    pie_function_start

    lea (multiboot_information_pointer - 0b)(%esi), %eax
    mov %ebx, (%eax)

    pie_function_end

/**
 * @brief Assert that the CPU supports the CPUID instruction.
 *
 * The primary way to check for support of the instruction is to flip the ID
 * bin in EFLAGS and then check if this changed was accepted. If so, the CPU
 * supports the CPUID instruction, otherwise it most-likely doesn't.
 */
_assert_cpuid_instruction_is_supported:
    pie_function_start

    pushfl
    pop %eax
    mov %eax, %ecx

    xor $(1 << 21), %eax  /* Flip the ID bit */
    push %eax             /* Move the new bitset on the stack for loading */
    popfl                 /* Load the flags with ID set back into EFLAGS */
    pushfl                /* Copy the flags back onto the stack */
    pop %eax              /* Load the flags for further checking */

    push %ecx
    popfl

    cmp %ecx, %eax
    jne 1f
    lea (message_cpuid_instruction_no_supported - 0b)(%esi), %eax
    push %eax
    call _panic

1:
    pie_function_end

/**
 * @brief Assert that the CPU supports going into long mode.
 */
_assert_cpu_supports_long_mode:
    pie_function_start

    mov $0x80000000, %eax
    cpuid
    cmp $0x80000001, %eax
    jb 1f

    mov $0x80000001, %eax
    cpuid
    test $(1 << 29), %edx
    jnz 2f
1:
    lea (message_long_mode_not_supported - 0b)(%esi), %eax
    push %eax
    call _panic
2:
    pie_function_end

/**
 * @brief Prepare a recursive page map hierarchy
 *
 * @param ebp+8 The number of huge pages to map
 * @return void
 */
_prepare_page_maps:
    pie_function_start

    push %edi

    call _clear_page_map_memory

    lea (page_map_level_4 - 0b)(%esi), %edi
    mov %edi, %eax
    or $0b11, %eax
    mov %eax, (510 * 8)(%edi)

    lea (page_map_level_3_low - 0b)(%esi), %eax
    or $0b11, %eax
    mov %eax, (%edi)

    lea (page_map_level_3_high - 0b)(%esi), %eax
    or $0b11, %eax
    mov %eax, (511 * 8)(%edi)

    lea (page_map_level_3_low - 0b)(%esi), %edi
    lea (page_map_level_2 - 0b)(%esi), %eax
    or $0b11, %eax
    mov %eax, (%edi)

    lea (page_map_level_3_high - 0b)(%esi), %edi
    lea (page_map_level_2 - 0b)(%esi), %eax
    or $0b11, %eax
    mov %eax, (510 * 8)(%edi)

    lea (page_map_level_2 - 0b)(%esi), %edi
    mov 8(%ebp), %ecx

1:
    dec %ecx
    mov $(1 << 21), %eax
    mul %ecx
    or $0b10000011, %eax
    mov %eax, (%edi, %ecx, 8)

    test %ecx, %ecx
    jnz 1b

    pop %edi

    pie_function_end

/**
 * @brief Clear all page map memory by filling it with 0s.
 *
 * @return void
 */
_clear_page_map_memory:
    pie_function_start

    push %edi

    xor %eax, %eax
    mov $page_maps_size, %ecx
    shr $2, %ecx
    lea (page_maps_start - 0b)(%esi), %edi
    rep stosl

    pop %edi

    pie_function_end

/**
 * @p Enable memory virtualization via paging.
 *
 * Note: This routine expects for there to be a valid set of page maps already
 * set up for use.
 *
 * @return void
 */
_enable_paging:
    pie_function_start

    lea (page_map_level_4 - 0b)(%esi), %eax
    mov %eax, %cr3

    /* Enable Physical Address Extension */
    mov %cr4, %eax
    or $(1 << 5), %eax
    mov %eax, %cr4

    /* Enable long mode support */
    mov $0xC0000080, %ecx
    rdmsr
    or $(1 << 8), %eax
    wrmsr

    /* Enable paging */
    mov %cr0, %eax
    or $(1 << 31), %eax
    mov %eax, %cr0

    pie_function_end

/**
 * @brief Enable use of SSE instructions.
 */
_enable_sse:
    function_start

    mov %cr0, %eax
    and $0xfffffffb, %eax
    or $0x00000002, %eax
    mov %eax, %cr0

    mov %cr4, %eax
    or $(3 << 9), %eax
    mov %eax, %cr4

    function_end

/**
 * @brief Prepare a new GTD and load make it active.
 *
 * @return void
 */
_reload_gdt:
    pie_function_start

    sub $10, %esp
    lea (global_descriptor_table - 0b)(%esi), %eax
    movw $(global_descriptor_table_end - global_descriptor_table -1), (%esp)
    mov %eax, 2(%esp)
    movl $0, 6(%esp)

    lgdt (%esp)
    add $10, %esp

    pie_function_end