global start
extern long_mode_start

%define VRAM 0xb8000

section .bss
    align 4096
    p4_table: resb 4096
    p3_table: resb 4096
    p2_table: resb 4096

    stack_bottom: resb 4096 ;reserve 4096 bytes of stack space
    stack_top:

section .text
bits 32

start:
    mov esp, stack_top
    mov edi, ebx ; move multiboot info pointer to edi

    call test_multiboot
    call test_cpuid
    call use_cpuid
    ; loop a bit
    mov dword [VRAM], 0x2f4b2f4f ;write OK
    call spin_loop

    call setup_page_tables
    call enable_paging
    lgdt [gdt64.pointer]
    ;update selectors
    mov ax, gdt64.data
    mov ss, ax ;stack selector
    mov ds, ax ;data selector
    mov es, ax ;extra selector

    ;long jump to update code selector (cs)
    jmp gdt64.code:long_mode_start

    hlt

error:
    mov dword [VRAM],   0x4f524f45   ;ER
    mov dword [VRAM+4], 0x4f3a4f52   ;R:
    mov dword [VRAM+8], 0x4f204f20   ; <2 spaces>
    mov byte [VRAM+0xa], al
    hlt

spin_loop:
    mov ecx, 0xfffffff
    .loop:
    dec ecx
    cmp ecx, 0
    jne spin_loop.loop
    ret

test_multiboot:
    cmp eax, 0x36d76289
    jne no_multiboot
    ret

    no_multiboot:
        mov al, "0"
        jmp error

test_cpuid:
    pushfd               ; Store the FLAGS-register.
    pop eax              ; Restore the A-register.
    mov ecx, eax         ; Set the C-register to the A-register.
    xor eax, 1 << 21     ; Flip the ID-bit, which is bit 21.
    push eax             ; Store the A-register.
    popfd                ; Restore the FLAGS-register.
    pushfd               ; Store the FLAGS-register.
    pop eax              ; Restore the A-register.
    push ecx             ; Store the C-register.
    popfd                ; Restore the FLAGS-register.
    xor eax, ecx         ; Do a XOR-operation on the A-register and the C-register.
    jz no_cpuid          ; The zero flag is set, no CPUID.
    ret ;CPUID is available

    no_cpuid:
    mov al, "1"
    jmp error

use_cpuid:
    call print_cpu_string
    call test_long_mode

test_long_mode:
    mov eax, 0x8000_0000
    cpuid
    cmp eax, 0x8000_0001
    jb no_long_mode
    mov eax, 0x8000_0001
    cpuid
    test edx, 1 << 29
    jz no_long_mode
    ret
    no_long_mode:
    mov al, "2"
    jmp error


print_cpu_string:
    mov eax, 0x0
    cpuid

    mov eax, ebx
    mov [VRAM+0], al
    mov [VRAM+2], ah
    shr eax, 16
    mov [VRAM+4], al
    mov [VRAM+6], ah
    mov eax, edx

    mov [VRAM+8], al
    mov [VRAM+0xa], ah
    shr eax, 16
    mov [VRAM+0xc], al
    mov [VRAM+0xe], ah
    mov eax, ecx

    mov [VRAM+0x10], al
    mov [VRAM+0x12], ah
    shr eax, 16
    mov [VRAM+0x14], al
    mov [VRAM+0x16], ah
    mov [VRAM+0x18], byte ' '
    ret


; page table stuff
; in long mode x86 paging has 4096 byte pages, 4-level lookup table

setup_page_tables:
    ;first p4 entry maps to p3
    mov eax, p3_table
    or eax, 0b11 ;set present+writeable page table flags
    mov [p4_table], eax

    ; first p3 entry maps to p2
    mov eax, p2_table
    or eax, 0b11 ; present + writeable
    mov [p3_table], eax

    ;loop over p2 entries, setting them to huge pages (2MiB)
    mov ecx, 0
    .map_p2_table:
        ;map ecx'th P2 entry to ecx*2MiB
        mov eax, 0x20_0000 ;2MiB
        mul ecx
        or eax, 0b1000_0011 ;writeable + huge + present
        mov [p2_table + ecx * 8], eax ;write to table
        inc ecx
        cmp ecx, 512 ;we do this 512 times to fill out the P2 table
        jne .map_p2_table

    ret


enable_paging:
    ;load P4 address to cr3 register
    mov eax, p4_table
    mov cr3, eax

    ; enable physical address extension (PAE)
    mov eax, cr4
    or eax, 1 << 5
    mov cr4, eax

    ;set long mode bit in EFER Model Specific Register (MSR)
    mov ecx, 0xc000_0080
    rdmsr ;reads msr[ecx] to edx:eax
    or eax, 1 << 8
    wrmsr

    ;enable paging in cr0
    mov eax, cr0
    or eax, 1 << 31
    or eax, 1 << 16
    mov cr0, eax

    ret

section .rodata
    ; multiboot grub establishes a 32-bit GDT
    ; we need to establish a 64-bit one, which can go in initialized read-only data
    ; dq = define quad, i.e. 64 bits
    gdt64:
        dq 0 ; zero entry
    .code: equ $ - gdt64
        dq (1<<44) | (1<<47) | (1<<41) | (1<<43) | (1<<53) ; code segment
    .data: equ $ - gdt64
        dq (1<<44) | (1<<47) | (1<<41) ; data segment
    .pointer:
        dw .pointer - gdt64
        dq gdt64

