; Paging
; Base addresses
%define PAGING_TABLE_SIZE         0x1000
%define PAGING_TABLES_BASE        0x8000

; Entry Flags
%define PT_FLAG_P                 (1 << 0)                        ; Present
%define PT_FLAG_W                 (1 << 1)                        ; Writable
%define PT_FLAGS_P_W              (PT_FLAG_P | PT_FLAG_W)         ; Present & Writable

; Base constants [naming convention: P**T_PML4Tindex_PDPTindex_PDTindex_BASE]
%define PML4T_BASE                (PAGING_TABLES_BASE)
%define PDPT_0_BASE               (PAGING_TABLES_BASE + 1 * PAGING_TABLE_SIZE)
%define PDPT_510_BASE             (PAGING_TABLES_BASE + 2 * PAGING_TABLE_SIZE)
%define PDT_0_0_BASE              (PAGING_TABLES_BASE + 3 * PAGING_TABLE_SIZE)
%define PDT_510_0_BASE            (PAGING_TABLES_BASE + 4 * PAGING_TABLE_SIZE)
%define PT_0_0_0_BASE             (PAGING_TABLES_BASE + 5 * PAGING_TABLE_SIZE)
%define PT_510_0_0_BASE           (PAGING_TABLES_BASE + 6 * PAGING_TABLE_SIZE)

; CR4 bit masks
%define CR4_PAE_MASK              1 << 5


; Parameter 1: Base address of the higher level PT
; Parameter 2: Base address of the lower level PT
; Parameter 3: Attributes for the higher level PT entry
%macro SetEntryInHighLevelPT 3
    mov edi, %1                 ; address of the higher level table
    mov dword [edi], %2         ; set one level lower table to point to
    or dword [edi], %3          ; set attributes
%endmacro

; Parameter 1: Base address of the PT
; Parameter 2: Start of the 2MB area we want to map
; Parameter 3: Attributes for the PT's entries
%macro MapTwoMegsToPT 3
    mov edi, %1     ; PT_BASE
    mov edx, %2     ; MEM_BASE
    mov eax, %3     ; Attributes

    mov ecx, 512    ; we want to populate all entries
    %%SetPTEntry:
    mov dword [edi], edx
    or dword [edi], eax
    add edi, 8      ; go to next entry
    add edx, 0x1000 ; go to next memory base
    loop %%SetPTEntry
%endmacro

; we'll set up the page tables at 0x8000
; we'll identity map the first 2MB of memory
; and then map the next physical 2MB at virtual address 255TB
SetUpLongModePaging:
    ; We need one PML4T table, two PDPTs, two PDTs and two PTs
    ; Each has TABLE_SIZE (0x1000) bytes, so we need to clear 7 * TABLE_SIZE bytes
    mov edi, PAGING_TABLES_BASE
    xor eax, eax
    mov ecx, 7 * PAGING_TABLE_SIZE
    rep stosb

    ; PML4T
    ; we need to map this entire table (well ,without the 511th entry, but why not go for that as well)
    ; this requires 8 (bytes per entry) * 512 entries = 4KB

    ; entry 0
    SetEntryInHighLevelPT   PML4T_BASE,                              PDPT_0_BASE,        PT_FLAGS_P_W

    ; entry 510
    SetEntryInHighLevelPT   (PML4T_BASE + 510 * 8),                  PDPT_510_BASE,      PT_FLAGS_P_W
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

    ; PDPTs
    ;;;;;;;;;;
    ; PDPT_0
    SetEntryInHighLevelPT   PDPT_0_BASE,                             PDT_0_0_BASE,       PT_FLAGS_P_W

    ; PDPT_510
    SetEntryInHighLevelPT   PDPT_510_BASE,                           PDT_510_0_BASE,     PT_FLAGS_P_W
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

    ; PDTs
    ;;;;;;;;;;
    ; PDT_0_0
    SetEntryInHighLevelPT   PDT_0_0_BASE,                            PT_0_0_0_BASE,      PT_FLAGS_P_W

    ; PDT_512_0
    SetEntryInHighLevelPT   PDT_510_0_BASE,                          PT_510_0_0_BASE,    PT_FLAGS_P_W
    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;


    ; Now the harder part :), the PTs
    ; PT_0_0_0
    ; Since we are identity mapping the first 2MB, we have to set the first entry to point to address 0,
    ;   the next to 0x1000, the next to 0x2000 etc.
    MapTwoMegsToPT          PT_510_0_0_BASE,        0x200000,   PT_FLAGS_P_W
    MapTwoMegsToPT          PT_0_0_0_BASE,          0,          PT_FLAGS_P_W

    ; PT_510_0_0
    ; We're mapping these addresses to the next 2MB of memory (2MB-4MB)

    ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;


    ; Finally, we'll give the CPU the address to the PML4T
    mov eax, PML4T_BASE
    mov cr3, eax

    ; Now, let's enable PAE paging
    mov eax, cr4
    or eax, CR4_PAE_MASK
    mov cr4, eax

    ; set LM bit
    mov ecx, 0xC0000080
    rdmsr
    or eax, 1 << 8
    wrmsr

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

    ; go to long mode..
    lgdt [gdt64_pointer]
    call 0x8:main64bit

    ret