bits 16
org 0x8000

xor ax, ax
mov ds, ax
mov es, ax
mov ss, ax
mov sp, 0x7000

FIRST_PARTITION_SECTOR equ 2048 ; TODO Read MBR.
STAGE_BUFFER_SEGMENT   equ 0x0900
STAGE_BUFFER_OFFSET    equ 0x0000
KERNEL_BUFFER_SEGMENT  equ 0x2000

jmp prepare_kernel

dap:
.dap_size:    db 16
.reserved:    db 0
.src_sectors: dw 1
.dst_offset:  dw 0
.dst_segment: dw 0
.src_lba:     dq 0
dap_end:

bpb:
.bytes_per_sector:    dw 0
.sectors_per_cluster: db 0
.reserved_sectors:    dw 0
.fats_count:          db 0
.root_entry_count:    dw 0
.total_sectors_16:    dw 0
.media_type:          db 0
.sectors_per_fat:     dw 0
bpb_end:

fat_start:  dw 0
root_start: dw 0
root_size:  dw 0
data_start: dw 0

kernel_filename:      db 'KERNEL  BIN'
kernel_cluster:       dw 0
kernel_buffer_offset: dw 0

read_sectors:
    xor bx, bx
    mov ds, bx

    mov [dap.src_lba], eax
    mov [dap.src_sectors], cx
    mov [dap.dst_segment], es
    mov [dap.dst_offset], di
    mov ah, 0x42
    mov dl, 0x80
    mov si, dap
    int 0x13
    jc disk_error
    ret

disk_error:
    mov ah, 0x0E
    mov al, 'E'
    int 0x10
    jmp $

prepare_kernel:
load_bpb:
    xor bx, bx
    mov ds, bx

    mov eax, FIRST_PARTITION_SECTOR
    mov cx, 1
    mov bx, STAGE_BUFFER_SEGMENT
    mov es, bx
    mov di, STAGE_BUFFER_OFFSET
    call read_sectors

    mov bx, STAGE_BUFFER_SEGMENT
    mov ds, bx

    mov si, STAGE_BUFFER_OFFSET + 11
    xor bx, bx
    mov es, bx
    mov di, bpb
    mov cx, bpb_end - bpb
    rep movsb

find_data_start:
    xor bx, bx
    mov ds, bx

    mov ax, FIRST_PARTITION_SECTOR
    add ax, [bpb.reserved_sectors]
    mov [fat_start], ax

    xor ax, ax
    mov al, [bpb.fats_count]
    mul word [bpb.sectors_per_fat]
    add ax, [fat_start]
    mov [root_start], ax ; dx should be zero for sane FAT16 values

    mov ax, [bpb.root_entry_count]
    shl ax, 5 ; * 32
    shr ax, 9 ; / 512 ; there should be at least 512 entries, so safe
    mov [root_size], ax

    mov ax, [root_start]
    add ax, [root_size]
    mov [data_start], ax

find_kernel:
    xor bx, bx
    mov ds, bx

    xor eax, eax
    mov ax, [root_start]
    mov cx, [root_size]
    mov bx, STAGE_BUFFER_SEGMENT
    mov es, bx
    mov di, STAGE_BUFFER_OFFSET
    call read_sectors

    mov cx, [bpb.root_entry_count]
    mov bx, STAGE_BUFFER_SEGMENT
    mov ds, bx
    mov si, STAGE_BUFFER_OFFSET

.check_entry:
    mov bx, STAGE_BUFFER_SEGMENT
    mov ds, bx

    cmp byte [si], 0x00
    je .no_more_entries
    cmp byte [si], 0xE5
    je .continue

    push si
    push cx
    xor bx, bx
    mov es, bx
    mov di, kernel_filename
    mov cx, 11
    rep cmpsb
    pop cx
    pop si
    je .found

.continue:
    add si, 32
    loop .check_entry

.no_more_entries:
    mov ah, 0x0E
    mov al, 'K'
    int 0x10
    jmp $

.found:
    mov bx, STAGE_BUFFER_SEGMENT
    mov ds, bx

    mov ax, [si + 26]

    xor bx, bx
    mov ds, bx

    mov [kernel_cluster], ax

load_fat:
    xor bx, bx
    mov ds, bx

    xor eax, eax
    mov ax, [fat_start]
    mov cx, [bpb.sectors_per_fat]
    mov bx, STAGE_BUFFER_SEGMENT
    mov es, bx
    mov di, STAGE_BUFFER_OFFSET
    call read_sectors

load_kernel_cluster:
    xor bx, bx
    mov ds, bx

    xor eax, eax
    mov ax, [kernel_cluster]
    sub ax, 2
    xor bx, bx
    mov bl, [bpb.sectors_per_cluster]
    mul bx
    add ax, [data_start]
    xor cx, cx
    mov cl, [bpb.sectors_per_cluster]
    mov bx, KERNEL_BUFFER_SEGMENT
    mov es, bx
    mov di, [kernel_buffer_offset]
    call read_sectors

    mov ax, [bpb.sectors_per_cluster]
    shl ax, 9 ; * 512
    add ax, [kernel_buffer_offset]
    mov [kernel_buffer_offset], ax

find_next_kernel_cluster:
    xor bx, bx
    mov ds, bx

    mov si, [kernel_cluster]
    shl si, 1 ; * 2
    add si, STAGE_BUFFER_OFFSET
    mov bx, STAGE_BUFFER_SEGMENT
    mov ds, bx
    mov ax, [si]
    xor bx, bx
    mov ds, bx
    mov [kernel_cluster], ax
    cmp ax, 0xFFF8
    jb load_kernel_cluster

; ------------------------------------------------------------------------------

jmp load_gdt

gdt_start:
gdt_0: dq 0x0000000000000000
gdt_32_code:
.limit_low:     dw 0xFFFF
.base_low:      dw 0x0000
.base_middle:   db 0x00
.access:        db 10011010b ; present, ring 0, code, executable, readable
.flags:         db 11001111b ; 32-bit, 4KB granularity
.base_high:     db 0x00
gdt_32_code_end:
gdt_64_code:
.limit_low:     dw 0xFFFF
.base_low:      dw 0x0000
.base_middle:   db 0x00
.access:        db 10011010b ; present, ring 0, code, executable, readable
.flags:         db 10101111b ; long mode (L bit)
.base_high:     db 0x00
gdt_64_code_end:
gdt_data:
.limit_low:     dw 0xFFFF
.base_low:      dw 0x0000
.base_middle:   db 0x00
.access:        db 10010010b ; present, ring 0, data, writable
.flags:         db 00000000b
.base_high:     db 0x00
gdt_data_end:
gdt_end:

gdt_descriptor:
.size:   dw gdt_end - gdt_start - 1
.offset: dd gdt_start
gdt_descriptor_end:

load_gdt:
    cli

    lgdt [gdt_descriptor]

    mov eax, cr0
    or eax, 0x00000001
    mov cr0, eax

    jmp 0x0008:protected_mode

; ------------------------------------------------------------------------------

bits 32
protected_mode:
    mov ax, 0x0018          ; data segment is now third entry = 0x18
    mov ds, ax
    mov es, ax
    mov ss, ax
    mov esp, 0x00090000

    ; zero page table memory
    mov edi, 0x00010000
    mov ecx, 256 * 5
    xor eax, eax
    rep stosd

    ; PML4[0] -> PDPT at 0x00011000
    mov dword [0x00010000], 0x00011003
    ; PML4[511] -> PDPT at 0x00012000
    mov dword [0x00010FF8], 0x00012003
    ; PML4[0].PDPT[0] -> PD at 0x00013000
    mov dword [0x00011000], 0x00013003
    ; PML4[511].PDPT[510] -> PD at 0x00014000
    mov dword [0x00012FF0], 0x00014003
    ; PD: fill 64 entries, each mapping 2MB
    mov edi, 0x00013000
    mov eax, 0x00000083   ; present, writable, huge page, base 0
    mov ecx, 64

.fill_pd:
    mov ebx, edi
    add ebx, 0x1000
    mov dword [edi], eax
    mov dword [ebx], eax
    add edi, 8
    add eax, 0x00200000   ; next 2MB
    loop .fill_pd

    ; load PML4 into CR3
    mov eax, 0x00010000
    mov cr3, eax

    ; enable PAE
    mov eax, cr4
    or eax, 0x00000020
    mov cr4, eax

    ; enable long mode in EFER MSR
    mov ecx, 0xC0000080
    rdmsr
    or eax, 0x00000100
    wrmsr

    ; enable paging (and keep protection enabled)
    mov eax, cr0
    or eax, 0x80000000
    mov cr0, eax

    jmp 0x0010:long_mode    ; 0x10 = second entry = 64-bit code segment

; ------------------------------------------------------------------------------

bits 64
long_mode:
    mov ax, 0x0018          ; data segment
    mov ds, ax
    mov es, ax
    mov ss, ax
    mov rsp, 0xFFFFFFFF80090000

    xor rax, rax
    xor rbx, rbx
    xor rcx, rcx
    xor rdx, rdx

    jmp 0xFFFFFFFF80020000