.section .efi.text

// Function to get UEFI memory map
get_memory_map:
    push %rbp
    mov %rsp, %rbp

    // Local variables for memory map
    sub $48, %rsp
    mov $0, -8(%rbp)   // memory_map_size
    mov $0, -16(%rbp)  // memory_map_buffer
    mov $0, -24(%rbp)  // map_key
    mov $0, -32(%rbp)  // descriptor_size
    mov $0, -40(%rbp)  // descriptor_version

    // First call to get size
    lea -8(%rbp), %rcx    // memory_map_size
    mov $0, %rdx          // memory_map
    lea -24(%rbp), %r8    // map_key
    lea -32(%rbp), %r9    // descriptor_size
    push $0               // descriptor_version
    call boot_services_get_memory_map
    
    // Allocate buffer
    mov -8(%rbp), %rcx    // size
    add $1000, %rcx       // Add some extra space
    mov $8, %rdx          // alignment
    lea -16(%rbp), %r8    // buffer pointer
    call boot_services_allocate_pool

    // Get actual map
    mov -8(%rbp), %rcx
    mov -16(%rbp), %rdx
    lea -24(%rbp), %r8
    lea -32(%rbp), %r9
    push -40(%rbp)
    call boot_services_get_memory_map

    // Save map key
    mov -24(%rbp), %rax
    mov %rax, memory_map_key

    leave
    ret

// Function to exit boot services
exit_boot_services:
    push %rbp
    mov %rsp, %rbp
    
    // Parameters already in rcx (image_handle) and rdx (map_key)
    call boot_services_exit

    leave
    ret

// Function to set up page tables
setup_page_tables:
    push %rbp
    mov %rsp, %rbp

    // Clear tables
    mov $pml4_table, %rdi
    mov $12288, %ecx      // 3 pages (PML4, PDPT, PD)
    xor %eax, %eax
    rep stosb

    // Set up identity mapping for first 2MB
    // PML4[0] -> PDPT
    mov $pdpt_table, %eax
    or $3, %eax           // Present + Write
    mov %eax, pml4_table

    // PDPT[0] -> PD
    mov $pd_table, %eax
    or $3, %eax
    mov %eax, pdpt_table

    // PD[0] -> 2MB page
    mov $0x83, %eax       // Present + Write + Huge
    mov %eax, pd_table

    // Map kernel to higher half
    // PML4[511] -> PDPT
    mov $pdpt_table, %eax
    or $3, %eax
    mov %eax, pml4_table + 4088

    // Load CR3
    mov $pml4_table, %rax
    mov %rax, %cr3

    leave
    ret