global _start
extern kmain
extern handle_isr
extern pml4

%define FLAGS       0b10
%define MAGIC       0x1BADB002
%define CHECKSUM    -(MAGIC + FLAGS)

section .multiboot
    align 4
    dd MAGIC
    dd FLAGS
    dd CHECKSUM

section .bss
    align 16
        resb 32768
    stack_top:

section .bss
    align 4096
    pml4:
        resb 4096
    pdpt:
        resb 4096
    pd:
        resb 4096

section .data
    align 8
    gdt64:
        dq 0x0000000000000000
    .code:
        dq 0x00AF9A000000FFFF
    .descriptor:
        dw .end - gdt64 - 1
        dq gdt64
    .end:

section .data
    align 8
    multiboot_info:
        dq 0

section .text
    bits 32
    _start:
        mov esp, stack_top

        ; Multiboot will place a magic value in eax
        ; If this magic value is not present, then we might not have multiboot info in ebx,
        ; therefore we throw an error
        cmp eax, 0x2BADB002
        jne error

        ; Save multiboot info pointer for later
        mov [multiboot_info], ebx

        ; Check if cpuid is available by flipping bit 21
        ; in the eflags register and checking if the cpu flipped it back
        pushfd
        pop eax
        mov ecx, eax
        xor eax, 1 << 21
        push eax
        popfd
        pushfd
        pop eax
        ; If cpuid is available, eax should be different than ecx
        xor eax, ecx
        jz error
        ; Finally restore eflags register to the original value
        push ecx
        popfd

        ; Check if extended cpuid is available by calling cpuid with 0x80000000,
        ; If cpuid is available, eax will be greater than 0x80000000 after the call
        mov eax, 0x80000000
        cpuid
        cmp eax, 0x80000001
        jb error

        ; Check if long mode is available by calling cpuid with 0x80000001
        ; this will place the extended features of the cpu in edx
        ; Bit 29 tells us if long mode is supported or not
        mov eax, 0x80000001
        cpuid
        test edx, 1 << 29
        jz error

        ; Enable PAE by setting bit 5 in cr4 to 1
        mov eax, cr4
        or eax, 1 << 5
        mov cr4, eax

        ; pml4[0] -> pdpt
        mov eax, pdpt
        or eax, 0x03
        mov [pml4], eax
        ; pdpt[0] -> pd
        mov eax, pd
        or eax, 0x03
        mov [pdpt], eax
        ; Map first 32 2mb pages for the kernel for a total of 64mb
        mov edi, pd
        mov eax, 0x83
        mov ecx, 32
    .setup_pd:
        mov [edi], eax
        add eax, 0x200000
        add edi, 8
        loop .setup_pd

        ; Load cr3 with the address of pml4
        mov eax, pml4
        mov cr3, eax

        lgdt [gdt64.descriptor]

        ; Enable long mode by setting bit 8 to 1 in EFER (Extended Feature Enable Register)
        mov ecx, 0xc0000080
        rdmsr
        or eax, 1 << 8
        wrmsr

        ; Enable paging by setting bit 31 in cr0 to 1
        mov eax, cr0
        or eax, 1 << 31
        mov cr0, eax

        jmp 0x8:long_mode

    error:
        cli
        mov byte [0xb8000], 'E'
        mov byte [0xb8002], 'R'
        mov byte [0xb8004], 'R'
    .hang:
        hlt
        jmp .hang

section .bss
    align 4096
    idt64:
        resb 4096

section .data
    align 8
    idt64_descriptor:
        dw 4095
        dq idt64

section .data
    align 8
    isr_stub_table:
        %assign i 0
        %rep 256
            dq isr_stub_%[i]
            %assign i i+1
        %endrep

%macro ISR_NOERR 1
isr_stub_%1:
    push qword 0
    push qword %1
    jmp isr_common
%endmacro

%macro ISR_ERR 1
isr_stub_%1:
    push qword %1
    jmp isr_common
%endmacro

section .text
    bits 64
    long_mode:
        ; Clear segment registers in long mode
        xor ax, ax
        mov ds, ax
        mov es, ax
        mov fs, ax
        mov gs, ax
        mov ss, ax

        ; Fill in the idt table with the stub functions
        mov rdi, idt64
        mov rsi, isr_stub_table
        mov rcx, 256
    .loop:
        mov rax, [rsi]
        mov [rdi], ax
        mov word [rdi + 2], 0x08
        mov word [rdi + 4], 0x8E00
        shr rax, 16
        mov [rdi + 6], ax
        shr rax, 16
        mov [rdi + 8], eax
        mov dword [rdi + 12], 0
        add rdi, 16
        add rsi, 8
        loop .loop
        lidt [idt64_descriptor]

        ; Finally, we call in to c
        mov rdi, [multiboot_info]
        call kmain
    .hang:
        hlt
        jmp .hang

    isr_common:
        push rax
        push rbx
        push rcx
        push rdx
        push rsi
        push rdi
        push rbp
        push r8
        push r9
        push r10
        push r11
        push r12
        push r13
        push r14
        push r15

        mov rdi, rsp
        call handle_isr

        pop r15
        pop r14
        pop r13
        pop r12
        pop r11
        pop r10
        pop r9
        pop r8
        pop rbp
        pop rdi
        pop rsi
        pop rdx
        pop rcx
        pop rbx
        pop rax

        add rsp, 16
        iretq

    ; CPU exceptions 0-31. Some of these contain error codes, so we define them manually
    ISR_NOERR   0
    ISR_NOERR   1
    ISR_NOERR   2
    ISR_NOERR   3
    ISR_NOERR   4
    ISR_NOERR   5
    ISR_NOERR   6
    ISR_NOERR   7
    ISR_ERR     8
    ISR_NOERR   9
    ISR_ERR     10
    ISR_ERR     11
    ISR_ERR     12
    ISR_ERR     13
    ISR_ERR     14
    ISR_NOERR   15
    ISR_NOERR   16
    ISR_ERR     17
    ISR_NOERR   18
    ISR_NOERR   19
    ISR_NOERR   20
    ISR_NOERR   21
    ISR_NOERR   22
    ISR_NOERR   23
    ISR_NOERR   24
    ISR_NOERR   25
    ISR_NOERR   26
    ISR_NOERR   27
    ISR_NOERR   28
    ISR_NOERR   29
    ISR_ERR     30
    ISR_NOERR   31

    ; Hardware interrupts and user-defined interrupts (32-255). These don't have error codes
    %assign i 32
    %rep 224
        ISR_NOERR i
        %assign i i+1
    %endrep