[BITS 16]
[ORG 0x7C00]

start:
    ; Set up segments
    xor ax, ax
    mov ds, ax
    mov es, ax
    mov ss, ax
    mov sp, 0x7C00

    ; Save boot drive passed in dl by the BIOS
    mov [boot_drive], dl

    ; Clear screen
    mov ah, 0x00
    mov al, 0x03
    int 0x10

    ; Print boot drive number in dl
    mov al, dl
    call print_hex_byte
    mov ah, 0x0E
    mov al, 13
    int 0x10
    mov al, 10
    int 0x10

    ; Print startup messages
    mov si, msg_start
    call print_string

    mov si, msg_key
    call print_string

    mov si, msg_drm
    call print_string

    mov si, msg_kernel
    call print_string

    ; Check if Extensions are present
    mov dl, [boot_drive]
    mov ah, 0x41
    mov bx, 0x55AA
    int 0x13
    jc load_chs
    cmp bx, 0xAA55
    jne load_chs

    ; Load Stage 2 from disk using LBA read
    mov dl, [boot_drive]
    mov ah, 0x42
    mov si, dap
    int 0x13
    jc disk_error
    jmp jump_to_stage2

load_chs:
    mov al, 7           ; Read 7 sectors
    mov bx, 0x7E00      ; Destination offset
    call read_sectors_chs
    jc disk_error

jump_to_stage2:
    mov dl, [boot_drive]
    jmp 0x0000:0x7E00



disk_error:
    mov al, ah
    call print_hex_byte
    mov ah, 0x0E
    mov al, 13
    int 0x10
    mov al, 10
    int 0x10
    mov si, msg_err
    call print_string
halt:
    hlt
    jmp halt

print_hex_byte:
    push ax
    push cx
    mov cl, 4
    shr al, cl
    call print_nibble
    pop ax
    push ax
    call print_nibble
    pop cx
    pop ax
    ret

print_nibble:
    and al, 0x0F
    add al, '0'
    cmp al, '9'
    jbe .print
    add al, 7
.print:
    mov ah, 0x0E
    int 0x10
    ret

; --- Functions ---
print_string:
    mov ah, 0x0E ; Teletype output
.loop:
    lodsb
    or al, al
    jz .done
    int 0x10
    jmp .loop
.done:
    ret

read_sectors_chs:
    push es
    push ax             ; Save number of sectors to read
    push bx             ; Save destination offset

    ; 1. Get drive parameters
    mov ah, 0x08
    mov dl, [boot_drive]
    xor di, di          ; ES:DI = 0:0 to prevent BIOS bugs
    mov es, di
    int 0x13
    jc .use_default

    ; Save H_count (dh + 1) in bl, and S_track (cl & 0x3F) in bh
    movzx bx, dh
    inc bx              ; bx = H_count
    and cx, 0x003F      ; cx = S_track
    
    jcxz .use_default
    test bx, bx
    jz .use_default
    jmp .calc_chs

.use_default:
    mov cx, 63          ; default sectors per track
    mov bx, 16          ; default heads

.calc_chs:
    ; 2. Perform LBA to CHS calculation for LBA 34
    mov ax, 34          ; ax = LBA
    xor dx, dx
    div cx              ; ax = LBA / S_track, dx = LBA % S_track

    ; Sector S = (LBA % S_track) + 1
    mov cl, dl
    inc cl              ; cl = Sector S (bits 0-5)

    ; Head H = (LBA / S_track) % H_count
    xor dx, dx
    div bx              ; ax = Cylinder C (should be 0), dx = Head H
    mov dh, dl          ; dh = Head H

    ; Cylinder C is 0, so ch = 0, and cl bits 6-7 are 0
    mov ch, 0

    ; Restore destination offset to bx, and pop ax
    pop bx              ; es:bx = destination buffer
    pop ax              ; al = number of sectors to read (7)
    pop es              ; Restore original ES

    ; Load boot drive to dl
    mov dl, [boot_drive]

    ; Call CHS Read
    mov ah, 0x02
    int 0x13
    ret

chs_fail:
    pop bx
    pop ax
    pop es
    stc
    ret

; --- Data ---
msg_start:  db "Welcome to ArkOS Bootloader", 13, 10, 0
msg_key:    db "Checking Hardware Secure Key... OK", 13, 10, 0
msg_drm:    db "Loading DRM Module... OK", 13, 10, 0
msg_kernel: db "Loading Kernel & Launching OS...", 13, 10, 0
msg_err:    db "Disk Error loading Stage 2!", 13, 10, 0


boot_drive: db 0

align 4
dap:
    db 0x10 ; size of DAP
    db 0    ; zero
    dw 7    ; number of sectors to read
    dw 0x7E00 ; offset
    dw 0x0000 ; segment
    dq 34   ; LBA 34 (Stage 2 starts at sector 34 in hybrid GPT layout)

times 446-($-$$) db 0
; Partition table placeholder (4 partitions * 16 bytes = 64 bytes)
times 64 db 0
; Boot signature
dw 0xAA55