.section .efi.text, "ax"
.global _start
.code64  // UEFI is already in 64-bit protected mode

_start:
    // Save UEFI parameters
    mov %rcx, uefi_image_handle
    mov %rdx, uefi_system_table

    // Get memory map
    call get_memory_map

    // Exit boot services
    mov uefi_image_handle, %rcx
    mov memory_map_key, %rdx
    call exit_boot_services
    
    // Disable interrupts
    cli

    // Load GDT for long mode
    lgdt gdt64_pointer(%rip)

    // Setup page tables (identity map first 2MB + map kernel to higher half)
    call setup_page_tables

    // Enable PAE
    mov %cr4, %rax
    or $(1 << 5), %rax
    mov %rax, %cr4

    // Set long mode bit
    mov $0xC0000080, %ecx
    rdmsr
    or $(1 << 8), %eax
    wrmsr

    // Enable paging
    mov %cr0, %rax
    or $0x80000000, %rax
    mov %rax, %cr0

    // Long jump to higher half kernel
    lea higher_half(%rip), %rax
    jmp *%rax

.align 8
higher_half:
    // Now in long mode at higher half
    // Setup stack and jump to Rust code
    mov $kernel_stack_top, %rsp
    call kmain

// GDT for long mode
.align 16
gdt64:
    .quad 0                    // Null descriptor
    .quad 0x00AF9A000000FFFF  // Code segment
    .quad 0x00CF92000000FFFF  // Data segment

gdt64_pointer:
    .word gdt64_pointer - gdt64 - 1
    .quad gdt64

.section .efi.data, "aw"
uefi_image_handle:  .quad 0
uefi_system_table:  .quad 0
memory_map_key:     .quad 0

.section .bss
.align 4096
kernel_stack_bottom:
    .skip 16384  // 16 KB stack
kernel_stack_top:

// Page tables
.align 4096
pml4_table:
    .skip 4096
pdpt_table:
    .skip 4096
pd_table:
    .skip 4096