.extern _end_physical
.extern _init
.extern kernel_main


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

/**
 * Reserve space for the page maps we are going to used during startup.
 *
 * Note: We are going to use large pages to make the initial mapping code
 * simpler.
 *
 * We need:
 *   - A single PML 4 (since we will only use 4-level paging)
 *   - 2 PML 3s (since we need to map low (1+MiB) memory)
 *   - 2 PML 2s (since we need to map low (1+MiB) memory)
 */

.global page_map_level_4
page_map_level_4: .skip 512 * 8

.global page_map_level_3_low
page_map_level_3_low: .skip 512 * 8

.global page_map_level_2_low
page_map_level_2_low: .skip 512 * 8

/**
 * Reserve some space for the Multiboot 2 information pointer.
 */
.global multiboot_information_pointer
multiboot_information_pointer: .skip 4

/**
 * Stack space for the bootstrapping process.
 *
 * Note: We are going to reserve 1 MiB for now. If/when the kernel requires
 * more space to run, it will have to relocate the stack.  
 */
.section .boot_stack, "aw", @nobits
.align 16

stack_bottom: .skip 1 << 20
stack_top:

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

/**
 * A valid Global Descriptor Table is still required in long mode. However, we
 * only need a single entry for the "code segment", so we will setup a single
 * segment table below.
 *
 * Bit 43: "E" in the access byte => mark the segment as executable.
 * Bit 44: "S" in the access byte => mark the segment as code or data.
 * Bit 47: "P" in the access byte => mark the segment as being present.
 * Bit 53: "L" in the flags byte => mark the segment as being for long mode
 */

global_descriptor_table: .quad 0
global_descriptor_table_code = . - global_descriptor_table
.quad (1<<43) | (1<<44) | (1<<47) | (1<<53)

/**
 * We also need a pointer that we can load into the GDTR.
 *
 * The pointer consists of a word describing the size of the table minus 1 and
 * the pointer to the actual table.
 */
global_descriptor_table_pointer:
.word . - global_descriptor_table - 1
.quad global_descriptor_table

/**
 * We are going to print some messages in case we panic during boot, so we are
 * going to store them here as well
 */
message_prefix_panic:
.string "TeachOS 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."
mesage_long_mode_not_supported:
.string "Long mode is not supported by this platform."

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

/**
 * We need a pointer to our current position in the VGA text buffer.
 */
.global vga_buffer_pointer
vga_buffer_pointer: .long 0xb8000

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

/**
 * Print a given panic message and then halt the machine.
 *
 * Parameters:
 *   - [stack - 0] message: the message to print
 */
_panic:
    push %ebp
    mov %esp, %ebp

    push message_prefix_panic
    push $0x4f
    call _print
    add $8, %esp

    push 8(%ebp)
    push 0x4f
    call _print
    add $8, %esp

    hlt

/**
 * Print a message via the VGA buffer.
 *
 * Parameters:
 *   - [stack - 4] message: the message to print
 *   - [stack - 0] color: the color of the message
 */
_print:
    push %ebp
    mov %esp, %ebp

    push %ebx
    push %esi
    mov 8(%ebp), %eax
    mov 12(%ebp), %ebx
    mov $0, %ecx
    mov (vga_buffer_pointer), %esi

.Lprint_loop:
    mov (%ebx, %ecx), %dl
    test %dl, %dl
    je .Lupdate_vga_buffer_address
    mov %dl, (%esi, %ecx, 2)
    mov %al, 1(%esi, %ecx, 2)
    inc %ecx
    jmp .Lprint_loop

.Lupdate_vga_buffer_address:
    shl $1, %ecx
    add %ecx, (vga_buffer_pointer)

.Lprint_end:
    pop %esi
    pop %ebx
    mov %ebp, %esp
    pop %ebp
    ret

/**
 * This is our entry point after being loaded by the bootloader.
 *
 * Having this in assembly makes it easier for us to keep things together.
 */
.global _start
_start:
    mov $stack_top, %esp
    mov %esp, %ebp

    call assert_loaded_by_multiboot2_loader
    call assert_cpuid_instruction_is_supported
    call assert_cpu_supports_long_mode
    call prepare_page_maps
    call enable_paging
    call enable_sse

    lgdt (global_descriptor_table_pointer)
    jmp $global_descriptor_table_code,$_transition_to_long_mode

    hlt

/**
 * Assert that the CPU supports going into long mode.
 */
assert_cpu_supports_long_mode:
    mov $0x80000000, %eax
    cpuid
    cmp $0x80000001, %eax
    jb .Llong_mode_assertion_failed

    mov $0x80000001, %eax
    cpuid
    test $(1 << 29), %edx
    jz .Llong_mode_assertion_failed
    ret
.Llong_mode_assertion_failed:
    push $mesage_long_mode_not_supported
    call _panic

/**
 * 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:
    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
    je .Lcpuid_assertion_fail
    ret
.Lcpuid_assertion_fail:
    push $message_cpuid_instruction_no_supported
    call _panic

/**
 * 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:
    cmp $0x36d76289, %eax                   /* Check if we received the
                                               expected magic */
    jne .Lmultiboot2_assertion_fail         /* Panic otherwise */
    mov %ebx, multiboot_information_pointer /* Store the MBI pointer */
    ret
.Lmultiboot2_assertion_fail:
    push $message_not_loaded_by_multiboot2
    call _panic

/**
 * Enable paging.
 *
 * Note: This routine expects for there to be a valid set of page maps already
 * set up for use.
 */
enable_paging:
    mov $page_map_level_4, %eax
    mov %eax, %cr3

    /* Map the P4 table recursively */
    mov $page_map_level_4, %eax
    or 0b11, %eax
    // TODO: WHY THIS THROW ERROR?
    mov %eax, [$page_map_level_4 + 511 * 8]

    /* 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

    ret

/**
 * Enable use of SSE instructions.
 */
enable_sse:
    mov %cr0, %eax
    and $0xfffffffb, %eax
    or $0x00000002, %eax
    mov %eax, %cr0

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

    ret

/**
 * Prepare the page maps.
 *
 * We map all physical memory we were loaded in plus one additional page. The
 * mapping is done in terms of huge pages (2 MiB per page) to save on required
 * page map entries.
 */
prepare_page_maps:
    /* Add an entry to the PML4, pointing to the low PML3 */
    mov $page_map_level_3_low, %eax
    or $0x3, %eax
    mov %eax, (page_map_level_4 + ((0x0000000000100000 >> 39) & 0x1ff) * 8)

    /* Add an entry to the low PML3, pointing to the low PML2 */
    mov $page_map_level_2_low, %eax
    or $0x3, %eax
    mov %eax, (page_map_level_3_low + ((0x0000000000100000 >> 30) & 0x1ff) * 8)

    xor %ecx, %ecx

    mov $_end_linear, %esi
    shr $21, %esi
    add $2, %esi

.Lmap_pages:
    mov $(1 << 21), %eax
    mul %ecx
    or $((1 << 0) | (1 << 1) | (1 << 7)), %eax
    mov %eax, page_map_level_2_low(,%ecx,8)

    inc %ecx
    cmp %esi, %ecx
    jne .Lmap_pages

    ret

.section .boot_text, "ax", @progbits
.code64

_transition_to_long_mode:
    xor %rax, %rax
    mov %rax, %ss
    mov %rax, %ds
    mov %rax, %es
    mov %rax, %fs
    mov %rax, %gs

    movl $0xb8000, (vga_buffer_pointer)

    call _init

    call kernel_main
    hlt