; MEMORY LAYOUT
; R = reserved, U = usable
; --------------------------------------------------------------------
; R | 0x000000 - 0x000400: real-mode interrupt vector table
; R | 0x000400 - 0x000500: bios data area
; U | 0x000500 - 0x004000: stack
; U | 0x004000 - 0x006a00: globals
; U | 0x006a00 - 0x007c00: memory map
; U | 0x007c00 - 0x007e00: boot sector
; U | 0x007e00 - 0x080000: conventional usable memory
; R | 0x080000 - 0x0a0000: extended bios data area (maximum possible size)
; R | 0x0a0000 - 0x0c0000: video memory
; R | 0x0c0000 - 0x0c8000: video bios
; R | 0x0c8000 - 0x0f0000: bios expansions
; R | 0x0f0000 - 0x100000: motherboard bios

%include "defines.s"

; BIOS puts our boot sector at 0000:7c00
[org BOOT0_LOADPOINT]
; We're (probably) in real mode
[bits 16]

main:
  ; Disable interrupts
  cli

  xor ax, ax

  mov ds, ax
  mov es, ax

  ; Put the stack base at 0x4000.
  ; Stack grows high->low, so we'll grow away from our globals and program text.
  mov ss, ax
  mov bp, STACK_BASE
  mov sp, bp

  ; Segment for VGA (0xb800 * 16 = 0xb8000)
  mov ax, 0xb800
  mov fs, ax

  ; Set VGA mode
  ; https://mendelson.org/wpdos/videomodes.txt
  mov ax, 0x0003
  int 0x10

  ; Store boot drive number
  xor dh, dh
  push dx

  ; Get drive geometry
  mov di, 0x00
  mov ah, 0x08
  int 0x13
  jc panic
  ; Load sectors per track into cx & spill
  and cl, 0x3f
  xor ch, ch
  push cx
  ; Load number of heads into bx & spill
  movzx bx, dh
  inc bx
  push bx

  ; Load LBA 1.
  mov ax, 1
  mov bx, 0x7e00
  call read_lba

  ; Check the GPT header magic "EFI PART"
  mov cx, GPT_MAGIC_LEN
  mov si, gpt_magic
  mov di, 0x7e00
  repe cmpsb
  jne panic

  ; Ensure the 8-byte GPT starting LBA fits in 16 bits
  mov di, 0x7e00 ; The rep increments di so we need to reset it
  mov eax, [di + 0x4c]
  mov bx, [di + 0x4a]
  or ax, bx
  or eax, eax
  jnz panic
  ; Store the first 16 bits of the GPT starting LBA (we have made sure the remaining bits are 0)
  push word [di + 0x48]

  ; Load number of partitions
  mov ax, [di + 0x50]
  mov bx, [di + 0x52]
  or bx, bx
  jz .gpt_n_partitions_loaded
  ; Number of partitions overflows 16 bits, so we just concern ourselves with the first 65535.
  ; That's an awful lot of partitions anyway.
  mov ax, 0xffff
.gpt_n_partitions_loaded:
  push ax

  ; Load GPT entry size
  mov eax, [di + 0x54] ; Operand size override otherwise this is going to be painful
  mov ebx, eax
  ; Assert that the entry size is 128 * 2^n for some integer n>=0. This is required for a valid GPT
  ; and has the nice properties that:
  ; - If each entry is larger than a sector (512 bytes), they'll be sector-aligned.
  ; - If each entry is smaller than a sector, an integer number of them will fit into a sector.
  or eax, eax  ; Test size != 0 because 128 * 2^n != 0
  jz panic
  test eax, 127 ; Test size is a multiple of 128
  jnz panic
  ; Use the (n & (n - 1)) == 0 trick to test if the entry size is a power of 2. Since we already
  ; know it's a nonzero multiple of 128, if size is a power of 2 then size = 128 * 2^n holds.
  ; Therefore we don't need to bother dividing by 128 first (shr 7), which saves a couple of bytes.
  mov ecx, ebx
  dec ecx
  and ecx, eax
  jnz panic

  ; Find the "sector stride", which is the number of sectors we increment by each time we want to
  ; load a new entry.
  shr eax, 9      ; Divide by sector size to get sectors per entry
  cmp eax, 0xffff ; Make sure sectors per entry fits in 16 bits
  ja panic
  or ax, ax
  jnz .gpt_sector_stride_loaded
  ; Sector stride must be at least one or we'll load the same sector each time!
  inc ax
.gpt_sector_stride_loaded:
  push ax

  ; Find the "byte stride", which is the number of bytes we increment by each time we want to load
  ; the next entry in the same sector.
  cmp ebx, 512
  jb .gpt_find_entries_per_sector
  push word 0 ; Arbitrary byte stride since there's only one entry per sector
  push word 1 ; 1 entry per sector, since an entry is larger than a sector
  jmp .gpt_found_entries_per_sector
.gpt_find_entries_per_sector:
  push bx     ; Store byte stride = entry length in this case
  xor dx, dx
  mov ax, 512
  div bx      ; Find entries per sector
  push ax
.gpt_found_entries_per_sector:

  ; Set up stack variables for our second stage search loop.
  xor ax, ax
  push ax                                ; Current entry
  push ax                                ; Current entry within the current sector
  push ax                                ; Number of sectors loaded
  push word [bp - GPT_ENTRIES_START_LBA] ; Current LBA

  ; Search for the partition storing our second stage.
.loop_find_stage2:
  mov dx, [bp - GPT_CURRENT_ENTRY_IDX]
  cmp [bp - GPT_N_ENTRIES_16], dx
  ; Panic if we've run out of partitions and haven't found the second stage yet.
  jbe panic

  ; If we haven't loaded any sectors yet, load the first one.
  cmp word [bp - GPT_SECTORS_LOADED], 0
  je .load_first_lba
  ; If there's still more entries in the current sector, skip loading a new sector
  mov ax, [bp - GPT_SECTOR_ENTRY_IDX]   ; Load current entry index within the current sector
  cmp [bp - GPT_ENTRIES_PER_SECTOR], ax ; Compare to entries per sector
  ja .process_current_entry
  
  mov ax, [bp - GPT_SECTOR_STRIDE]        ; Load sector stride
  add word [bp - GPT_CURRENT_LBA], ax     ; Increment current LBA by sector stride
  mov word [bp - GPT_SECTOR_ENTRY_IDX], 0 ; Reset the current entry index within the current sector
.load_first_lba:
  ; Read the current LBA to 0x8000 (just past the end of the GPT header)
  mov ax, [bp - GPT_CURRENT_LBA]
  mov bx, 0x8000
  call read_lba
  ; Increment number of sectors loaded
  inc word [bp - GPT_SECTORS_LOADED]

.process_current_entry:
  ; Calculate the address of the current GPT entry.
  mov ax, [bp - GPT_SECTOR_ENTRY_IDX] ; Load current entry index within current sector
  xor dx, dx
  mul word [bp - GPT_BYTE_STRIDE]     ; Get the byte offset in the current sector of the current entry
  add ax, 0x8000                      ; Convert offset to address (we loaded the sector at 0x8000)

  ; Compare entry GUID to our stage 2 partition GUID.
  mov cx, GUID_LEN
  mov si, guid_stage2
  mov di, ax
  repe cmpsb
  je .found_stage2

  ; Next iteration
  inc word [bp - GPT_CURRENT_ENTRY_IDX] ; Increment current entry index
  inc word [bp - GPT_SECTOR_ENTRY_IDX]  ; Increment current entry index within the current sector
  jmp .loop_find_stage2

.found_stage2:
  push ax ; Address of the GPT entry for stage 2
  mov si, ax

  ; Load partition LBA start.
  mov eax, [si + 0x20]
  mov ebx, [si + 0x24]
  ; Ensure it fits in 16 bits.
  or ebx, ebx
  jnz panic
  cmp ebx, 0xffff
  ja panic
  ; Load partition LBA end.
  mov ecx, [si + 0x28]
  mov edx, [si + 0x2c]
  ; Assert that the end LBA is greater than or equal to the start LBA, so we have at least one
  ; sector to load (end LBA is inclusive).
  or edx, edx
  jnz .stage2_end_lba_ok
  cmp eax, ecx
  ja panic
.stage2_end_lba_ok:

  mov bx, BOOT1_LOADPOINT
  call read_lba
  jmp bx

  ; Load a single boot disk sector. Panic on failure.
  ; Inputs:
  ; - ax: LBA to load
  ; - bx: address to read sector to
  ; Clobber: ax, cx, dx
read_lba:
  ; sector - 1 = LBA  % sectors_per_track
  ; temp       = LBA  / sectors_per_track
  ; head       = temp % n_heads
  ; cylinder   = temp / n_heads
  xor dx, dx
  ; Divide by sectors per track. dx = mod (sector - 1), ax = div (temp)
  div word [bp - SECTORS_PER_TRACK]
  ; Put the sector into cx (the bios call will use cl)
  mov cx, dx
  inc cx
  xor dx, dx
  ; Divide by number of heads. dx = mod (head), ax = div (cylinder)
  div word [bp - N_HEADS]
  mov dh, dl
  mov ch, al
  mov dl, byte [bp - BOOT_DRIVE]
  mov ah, 0x02
  mov al, 1
  ; Read sector
  int 0x13
  jc panic
  ret

panic:
  mov ax, 0x0003
  int 0x10
  mov word fs:[0x0000], 0x4f21
  hlt

gpt_magic     db "EFI PART"
GPT_MAGIC_LEN equ $ - gpt_magic

; Our stage2 guid: fdffea69-3651-442f-a11d-88a09bf372dd
guid_stage2 db 0x69, 0xea, 0xff, 0xfd, 0x51, 0x36, 0x2f, 0x44, \
               0xa1, 0x1d, 0x88, 0xa0, 0x9b, 0xf3, 0x72, 0xdd
GUID_LEN   equ $ - guid_stage2

; MBR bootstrap field is 440 bytes long
%if ($ - $$) > 440
%error "exceeded mbr bootstrap field size"
%endif