; File: elevator.asm
; Stage 2 Bootloader - No size limit!
; Job: Load kernel, set up GDT, switch to 64-bit mode, jump to kernel


[org 0x7e00]
[bits 16]

KERNEL_OFFSET equ 0x10000   ; Load kernel at 64KB (safer than 0x1000)
BOOT_DRIVE: db 0            ; Store boot drive number

stage2_start:
    ; Print stage 2 message
    mov si, msg_stage2
    call print_string

    ; Save boot drive number (passed from stage1 in DL)
    mov [BOOT_DRIVE], dl
    
    ; Load kernel from disk
    call load_kernel

    ; Verify kernel was loaded - check for signature
    ; Check if there's actual code at 0x10000
    mov ax, 0x1000
    mov es, ax
    xor bx, bx
    mov ax, [es:bx]          ; Read first 2 bytes of kernel
    
    ; Show what we read (for debugging)
    push ax
    xor ax, ax
    mov es, ax
    
    mov si, msg_kernel_check
    call print_string
    
    pop ax
    push ax
    call print_hex_word      ; Print the word we read
    
    ; Also check bytes 2-3
    mov ax, 0x1000
    mov es, ax
    mov bx, 2
    mov ax, [es:bx]
    xor bx, bx
    mov es, bx
    
    mov si, msg_space
    call print_string
    call print_hex_word
    
    ; Check if first word is non-zero
    pop ax
    test ax, ax
    jnz .kernel_found
    
    ; Kernel is zeros - try reading from different LBA offsets
    mov si, msg_searching
    call print_string
    jmp kernel_load_failed

.kernel_found:
    mov si, msg_kernel_ok
    call print_string

    ; Disable interrupts
    cli

    ; Load GDT
    lgdt [gdt_descriptor]

    ; Enable Protected Mode
    mov eax, cr0
    or eax, 0x1
    mov cr0, eax

    ; Far jump to 32-bit code
    jmp 0x08:protected_mode

; ============================================================
; LOAD KERNEL FROM DISK
; ============================================================
load_kernel:
    push ax
    push bx
    push cx
    push dx
    push si
    
    ; Show what we're about to do
    mov si, msg_loading_kernel
    call print_string
    
    ; Check if LBA is available
    mov ah, 0x41
    mov bx, 0x55AA
    mov dl, [BOOT_DRIVE]
    int 0x13
    jc .use_chs            ; If carry set, LBA not supported
    
    ; Use LBA mode
    mov si, msg_using_lba
    call print_string
    
    ; Set up DAP (Disk Address Packet)
    mov word [dap_size], 0x10
    mov word [dap_sectors], 10
    mov word [dap_offset], 0x0000
    mov word [dap_segment], 0x1000
    mov dword [dap_lba_low], 6      ; LBA 6 = bytes 3072+ (sector 7 in CHS)
    mov dword [dap_lba_high], 0
    
    ; Read using LBA
    mov ah, 0x42
    mov dl, [BOOT_DRIVE]
    mov si, dap
    int 0x13
    jc kernel_error
    
    jmp .read_done

.use_chs:
    ; Fall back to CHS mode
    mov si, msg_using_chs
    call print_string
    
    ; Reset disk system first
    mov ah, 0x00
    mov dl, [BOOT_DRIVE]
    int 0x13
    
    ; Set up ES:BX for the destination
    mov ax, 0x1000
    mov es, ax
    xor bx, bx
    
    ; Read using CHS
    mov ah, 0x02           ; Read function
    mov al, 10             ; Number of sectors
    mov ch, 0              ; Cylinder 0
    mov cl, 7              ; Sector 7
    mov dh, 0              ; Head 0
    mov dl, [BOOT_DRIVE]   ; Drive number
    
    int 0x13
    jc kernel_error

.read_done:
    ; Print success
    mov si, msg_sectors_read
    call print_string
    
    mov si, msg_newline
    call print_string
    
    ; Restore ES to 0
    xor ax, ax
    mov es, ax
    
    pop si
    pop dx
    pop cx
    pop bx
    pop ax
    ret

; Disk Address Packet for LBA
align 4
dap:
dap_size:       db 0x10
dap_reserved:   db 0
dap_sectors:    dw 0
dap_offset:     dw 0
dap_segment:    dw 0
dap_lba_low:    dd 0
dap_lba_high:   dd 0

kernel_error:
    ; Restore ES
    xor ax, ax
    mov es, ax
    
    ; Save error code
    mov byte [ERROR_CODE], ah
    
    mov si, msg_kernel_error
    call print_string
    
    ; Show error code
    mov si, msg_error_code
    call print_string
    
    ; Print error code as hex
    mov al, [ERROR_CODE]
    call print_hex
    
    ; Show drive number we tried to use
    mov si, msg_drive
    call print_string
    mov al, [BOOT_DRIVE]
    call print_hex
    
    jmp $

kernel_load_failed:
    xor ax, ax
    mov es, ax
    
    ; Kernel is zeros at expected location
    ; Let's try reading from different LBA sectors to find it
    mov si, msg_verify_failed
    call print_string
    
    ; Try LBA 5
    mov si, msg_try_lba5
    call print_string
    mov word [dap_lba_low], 5
    call try_read_kernel
    
    ; Try LBA 7
    mov si, msg_try_lba7
    call print_string
    mov word [dap_lba_low], 7
    call try_read_kernel
    
    ; Try LBA 8
    mov si, msg_try_lba8
    call print_string
    mov word [dap_lba_low], 8
    call try_read_kernel
    
    jmp $

try_read_kernel:
    push si
    ; Read to 0x10000
    mov ax, 0x1000
    mov es, ax
    mov word [dap_segment], ax
    
    mov ah, 0x42
    mov dl, [BOOT_DRIVE]
    mov si, dap
    int 0x13
    
    ; Check what we got
    xor bx, bx
    mov ax, [es:bx]
    
    xor cx, cx
    mov es, cx
    
    call print_hex_word
    mov si, msg_newline
    call print_string
    
    ; If non-zero, we found it!
    test ax, ax
    jz .not_found
    
    mov si, msg_found
    call print_string
    
    ; Jump directly to protected mode!
    cli
    lgdt [gdt_descriptor]
    mov eax, cr0
    or eax, 0x1
    mov cr0, eax
    jmp 0x08:protected_mode
    
.not_found:
    pop si
    ret

; ============================================================
; PRINT STRING (16-bit)
; ============================================================
print_string:
    pusha
    mov ah, 0x0e
.loop:
    lodsb
    test al, al
    jz .done
    int 0x10
    jmp .loop
.done:
    popa
    ret

; ============================================================
; PRINT HEX BYTE (for debugging)
; ============================================================
print_hex:
    push ax
    push bx
    
    mov bl, al
    shr al, 4           ; Get high nibble
    call print_nibble
    
    mov al, bl
    and al, 0x0F        ; Get low nibble
    call print_nibble
    
    pop bx
    pop ax
    ret

print_nibble:
    cmp al, 9
    jle .digit
    add al, 7           ; Convert to A-F
.digit:
    add al, '0'
    mov ah, 0x0e
    int 0x10
    ret

; Print a 16-bit word as hex
print_hex_word:
    push ax
    mov al, ah          ; Print high byte first
    call print_hex
    pop ax
    call print_hex      ; Print low byte
    ret

; ============================================================
; MESSAGES
; ============================================================
msg_stage2:        db '[STAGE2] Extended bootloader loaded', 0x0D, 0x0A, 0
msg_kernel_loaded: db '[STAGE2] Kernel loaded, switching to Protected Mode...', 0x0D, 0x0A, 0
msg_kernel_error:  db '[ERROR] Kernel load failed!', 0x0D, 0x0A, 0
msg_error_code:    db ' Error code: ', 0 
msg_drive:         db ' Drive: ', 0 
msg_loading_kernel: db '[STAGE2] Loading kernel from disk...', 0x0D, 0x0A, 0 
msg_using_lba:     db '[STAGE2] Using LBA mode for disk access', 0x0D, 0x0A, 0 
msg_using_chs:     db '[STAGE2] Using CHS mode for disk access', 0x0D, 0x0A, 0 
msg_sectors_read:  db '[STAGE2] Sectors read successfully', 0 
msg_newline:       db 0x0D, 0x0A, 0 
msg_kernel_check:  db '[STAGE2] Verifying kernel load: Read words: ', 0 
msg_space:         db ' ', 0 
msg_searching:     db '[STAGE2] Kernel not found, searching alternative sectors...', 0x0D, 0x0A, 0 
msg_kernel_ok:     db '[STAGE2] Kernel verification successful!', 0x0D, 0x0A, 0 
msg_verify_failed: db '[STAGE2] Kernel verification failed.', 0x0D, 0x0A, 0 
msg_try_lba5:      db '[STAGE2] Trying LBA 5...', 0x0D, 0x0A, 0 
msg_try_lba7:      db '[STAGE2] Trying LBA 7...', 0x0D, 0x0A, 0 
msg_try_lba8:      db '[STAGE2] Trying LBA 8...', 0x0D, 0x0A, 0 
msg_found:         db '[STAGE2] Kernel found!', 0x0D, 0x0A, 0 
ERROR_CODE: db 0 

; ============================================================
; GDT - Global Descriptor Table
; ============================================================
align 8
gdt_start:

gdt_null:
    dq 0x0

gdt_code32:     ; Selector 0x08
    dw 0xffff
    dw 0x0000
    db 0x00
    db 10011010b
    db 11001111b
    db 0x00

gdt_data32:     ; Selector 0x10
    dw 0xffff
    dw 0x0000
    db 0x00
    db 10010010b
    db 11001111b
    db 0x00

gdt_code64:     ; Selector 0x18
    dw 0xffff
    dw 0x0000
    db 0x00
    db 10011010b
    db 10101111b
    db 0x00

gdt_data64:     ; Selector 0x20
    dw 0xffff
    dw 0x0000
    db 0x00
    db 10010010b
    db 10101111b
    db 0x00

gdt_end:

gdt_descriptor:
    dw gdt_end - gdt_start - 1  ; Limit (size - 1)
    dd gdt_start                ; Base address (org already adjusts this)

; ============================================================
; 32-BIT PROTECTED MODE
; ============================================================
[bits 32]
protected_mode:
    ; Set up segments
    mov ax, 0x10
    mov ds, ax
    mov es, ax
    mov fs, ax
    mov gs, ax
    mov ss, ax
    mov esp, 0x90000

    ; Print to VGA
    mov edi, 0xb8000
    mov esi, pm_msg
    call print_pm

    ; Clear page table area (12KB starting at 0x1000)
    mov edi, 0x1000
    mov ecx, 3072       ; 12KB / 4 bytes = 3072 dwords
    xor eax, eax
    rep stosd

    ; Set up page tables
    ; PML4[0] -> PDPT at 0x2000
    mov dword [0x1000], 0x2003

    ; PDPT[0] -> PD at 0x3000
    mov dword [0x2000], 0x3003

    ; Identity map first 4MB (two 2MB pages)
    ; PD[0] -> First 2MB
    mov dword [0x3000], 0x00000083
    mov dword [0x3004], 0x00000000
    
    ; PD[1] -> Second 2MB (covers 0x200000 - 0x3FFFFF)
    mov dword [0x3008], 0x00200083
    mov dword [0x300C], 0x00000000

    ; Load page table
    mov eax, 0x1000
    mov cr3, eax

    ; Enable PAE
    mov eax, cr4
    or eax, 0x20
    mov cr4, eax

    ; Enable Long Mode
    mov ecx, 0xC0000080
    rdmsr
    or eax, 0x100
    wrmsr

    ; Enable Paging
    mov eax, cr0
    or eax, 0x80000000
    mov cr0, eax

    ; Jump to 64-bit mode
    jmp 0x18:long_mode

; Print in 32-bit protected mode
print_pm:
    mov ah, 0x0f
.loop:
    lodsb
    test al, al
    jz .done
    stosw
    jmp .loop
.done:
    ret

pm_msg: db '[PM32] Setting up Long Mode...', 0

; ============================================================
; 64-BIT LONG MODE
; ============================================================
[bits 64]
long_mode:
    ; Set up segments
    mov ax, 0x20
    mov ds, ax
    mov es, ax
    mov fs, ax
    mov gs, ax
    mov ss, ax
    mov rsp, 0x90000

    ; Set up a minimal IDT to prevent triple faults
    ; Point all interrupts to a dummy handler
    call setup_idt

    ; Print message on line 3
    mov rdi, 0xb8000 + (160 * 2)
    mov rsi, lm_msg
    call print_64

    ; Write a marker before jumping to kernel (for debugging)
    mov rdi, 0xb8000 + (160 * 3)
    mov rsi, before_kernel_msg
    call print_64

    ; JUMP TO KERNEL!
    call KERNEL_OFFSET

    ; If kernel returns, show error
    mov rdi, 0xb8000 + (160 * 4)
    mov rsi, kernel_return_msg
    call print_64
    jmp $

; Set up a minimal IDT
setup_idt:
    ; Fill IDT with different handlers for different exceptions
    mov rdi, 0x4000         ; IDT location
    
    ; Set up first 32 entries (CPU exceptions) with specific handlers
    mov rcx, 32
    mov rax, exception_handler
    
.fill_exceptions:
    mov word [rdi], ax              ; Offset low
    mov word [rdi + 2], 0x18        ; Code segment
    mov byte [rdi + 4], 0           ; IST
    mov byte [rdi + 5], 0x8E        ; Present, interrupt gate
    shr rax, 16
    mov word [rdi + 6], ax          ; Offset middle
    shr rax, 16
    mov dword [rdi + 8], eax        ; Offset high
    mov dword [rdi + 12], 0         ; Reserved
    
    mov rax, exception_handler      ; Restore for next iteration
    add rdi, 16
    loop .fill_exceptions
    
    ; Fill remaining entries (32-255) with dummy handler
    mov rcx, 224
    mov rax, dummy_handler
    
.fill_rest:
    mov word [rdi], ax
    mov word [rdi + 2], 0x18
    mov byte [rdi + 4], 0
    mov byte [rdi + 5], 0x8E
    shr rax, 16
    mov word [rdi + 6], ax
    shr rax, 16
    mov dword [rdi + 8], eax
    mov dword [rdi + 12], 0
    
    mov rax, dummy_handler
    add rdi, 16
    loop .fill_rest
    
    ; Load IDT
    lidt [idt_descriptor]
    ret

; Exception handler - shows what went wrong
exception_handler:
    ; Save registers
    push rax
    push rbx
    
    ; Get exception number from stack
    ; (not perfectly accurate but good enough for debugging)
    
    ; Write "EXCEPTION!" in red at line 5
    mov rbx, 0xb8000 + (160 * 5)
    mov rax, 0x0C450C580C430C45  ; "ECXE" in red
    mov [rbx], rax
    mov rax, 0x0C4F0C490C540C50  ; "PTIO" in red  
    mov [rbx + 8], rax
    mov word [rbx + 16], 0x0C21  ; "!" in red
    
    ; Show the instruction pointer where it failed
    mov rax, [rsp + 16]  ; Get RIP from stack (after our 2 pushes)
    mov rbx, 0xb8000 + (160 * 6)
    
    ; Write "RIP: " 
    mov word [rbx], 0x0F52  ; 'R'
    mov word [rbx + 2], 0x0F49  ; 'I'
    mov word [rbx + 4], 0x0F50  ; 'P'
    mov word [rbx + 6], 0x0F3A  ; ':'
    mov word [rbx + 8], 0x0F20  ; ' '
    
    ; Print RIP value as hex
    mov rcx, 16
    mov rbx, 0xb8000 + (160 * 6) + 10
.print_rip:
    rol rax, 4
    mov dl, al
    and dl, 0x0F
    cmp dl, 9
    jle .digit
    add dl, 7
.digit:
    add dl, '0'
    mov byte [rbx], dl
    mov byte [rbx + 1], 0x0F
    add rbx, 2
    dec rcx
    jnz .print_rip
    
    ; Restore and halt
    pop rbx
    pop rax
    
    ; Infinite loop instead of iretq
    cli
    hlt
    jmp $

; Dummy interrupt handler for IRQs (not CPU exceptions)
dummy_handler:
    iretq

idt_descriptor:
    dw (256 * 16) - 1       ; Limit (256 entries * 16 bytes)
    dq 0x4000               ; Base address

; Print in 64-bit mode
print_64:
    mov ah, 0x0f
.loop:
    lodsb
    test al, al
    jz .done
    stosw
    jmp .loop
.done:
    ret

lm_msg: db '[64-BIT] Long mode active...', 0
before_kernel_msg: db '[64-BIT] Calling kernel...', 0
kernel_return_msg: db '[KERNEL] Returned unexpectedly!', 0

; Add pading to make the elevator 2560 bytes. 
times 2560 - ($ - $$) db 0