[bits 16]

%include "fn.s"
%include "layout.s"
%include "globals.s"

extern s2_bin_len
extern s2_bin_sectors
extern s2_magic
extern s2_main

%macro copy_stack_var_to_globals 2
  mov %1, [bp - %2]
  mov [REAL_GLOBALS + %2], %1
%endmacro


section .prelude

s2_data:
.s3_bin_offset_sectors:
  dw 0
.s3_bin_len_sectors:
  dw 0
.padding:
  times (S2_DATA_LEN - 4) db 0xf4

global s2_data
global s2_data.s3_bin_offset_sectors
global s2_data.s3_bin_len_sectors

%if ($ - $$) != S2_DATA_LEN
%error "incorrect prelude data size"
%endif

; Load the rest of stage 2 into memory.
prelude:
  ; Now that we're not doing instruction byte golf like we were in stage 1, we can afford to move
  ; the various stage 1 stack variables to the globals section.
  copy_stack_var_to_globals ax, BOOT_DRIVE
  copy_stack_var_to_globals ax, SECTORS_PER_TRACK
  copy_stack_var_to_globals ax, N_HEADS
  copy_stack_var_to_globals ax, GPT_ENTRIES_START_LBA
  copy_stack_var_to_globals ax, GPT_N_ENTRIES_16
  copy_stack_var_to_globals ax, GPT_SECTOR_STRIDE
  copy_stack_var_to_globals ax, GPT_BYTE_STRIDE
  copy_stack_var_to_globals ax, GPT_ENTRIES_PER_SECTOR
  copy_stack_var_to_globals ax, GPT_CURRENT_ENTRY_IDX
  copy_stack_var_to_globals ax, GPT_SECTOR_ENTRY_IDX
  copy_stack_var_to_globals ax, GPT_SECTORS_LOADED
  copy_stack_var_to_globals ax, GPT_CURRENT_LBA
  copy_stack_var_to_globals ax, STAGE_2_GPT_ENTRY_ADDR

  ; Reset the stack, now we've got everything we need from it.
  mov sp, bp

  mov si, [REAL_GLOBALS + STAGE_2_GPT_ENTRY_ADDR]
  mov eax, [si + 0x20] ; Partition / s2 start LBA lower
  mov ebx, [si + 0x24] ; Partition / s2 start LBA upper
  mov ecx, [si + 0x28] ; Partition end LBA lower
  mov edx, [si + 0x32] ; Partition end LBA upper

  ; Panic if the partition / boot1 starting LBA overflows 16 bits.
  or ebx, ebx
  jnz panic_simple
  ror eax, 16
  or ax, ax
  jnz panic_simple
  ror eax, 16

  ; There must be at least one sector to load.
  mov bx, s2_bin_sectors
  or bx, bx
  jz panic_simple

  ; Calculate the s2 end LBA (inclusive) and panic if it overflows 16 bits.
  ; n.b. ebx is zero before this so both bx and ebx can be used as the s2 end LBA.
  dec bx
  add bx, ax
  jc panic_simple

  ; Panic if the s2 end LBA is after the partition end LBA.
  ; If the upper 32 bits of the partition end LBA are nonzero, then it must be greater than our
  ; 16-bit s2 end LBA.
  or edx, edx
  jnz .end_lba_ok
  ; Compare the s2 end LBA to the lower 32 bits of the partition end LBA.
  cmp ebx, ecx
  ja panic_simple

  ; Save partition / s2 extents for later.
  mov [REAL_GLOBALS + LOADER_PART_END_LBA], ecx
  mov [REAL_GLOBALS + STAGE_2_START_LBA], ax
  mov [REAL_GLOBALS + STAGE_2_END_LBA], bx

.end_lba_ok:

  ; The first sector has already been loaded (we're running it right now!) so increment the
  ; current LBA.
  inc ax
  push ax                     ; Current LBA
  push bx                     ; s2 end LBA
  mov ebx, S2_LOAD_ADDR + 512 ; Current sector load address

.load_loop:
  mov ax, [bp - 0x02]      ; Load current LBA
  cmp word [bp - 0x04], ax ; Compare to s2 end LBA
  jb .load_done

  mov ecx, ebx
  call read_sector
  jc panic_simple
  
  add ebx, 512
  inc word [bp - 0x02]
  jmp .load_loop

.load_done:

  ; Check the magic bytes at the end of s2.
  push es
  mov ebx, s2_magic
  call addr32_to_addr16
  cmp dword es:[bx], S2_MAGIC
  pop es
  jne panic_simple
  
  jmp s2_main


; Converts a 32-bit address to a 16-bit sector and offset.
; Arguments:
; - ebx: 32-bit address
; Return:
; - es: 16-bit address segment (unchanged on failure)
; - ebx: 16-bit address offset
; - cf: unset on success, set on failure
; Clobber: none
addr32_to_addr16:
  fnstart
  push es
  push eax

  mov eax, ebx
  ; Divide addr by 16 and saturate to 16 bits to get the segment.
  shr eax, 4
  ror eax, 16
  or ax, ax
  jz .segment_ok
  mov eax, 0xffff0000
.segment_ok:
  ror eax, 16
  mov es, ax

  ; Calculate offset = addr - (16 * segment), failing if the offset doesn't fit in 16 bits.
  shl eax, 4
  sub ebx, eax
  ror ebx, 16
  or bx, bx
  jnz .fail
  ror ebx, 16
  
  pop eax
  add sp, 2 ; Discard the original es from the stack
  pop bp
  clc
  ret

.fail:
  pop eax
  pop es
  stc
  fnret

global addr32_to_addr16


; Reads a single sector at the given LBA into memory.
; Arguments:
; - ax: start LBA
; - ecx: address to read sector to
; Return:
; - cf: unset on success, set on failure
; Clobber: eax, ecx, edx
read_sector:
  ; sector - 1 = LBA  % sectors_per_track
  ; temp       = LBA  / sectors_per_track
  ; head       = temp % n_heads
  ; cylinder   = temp / n_heads

  fnstart
  push es
  push ebx
  
  mov ebx, ecx
  call addr32_to_addr16
  jc .return
  
  ; Calculate sector and temp
  xor dx, dx
  ; Divide by sectors per track. dx = mod (sector - 1), ax = div (temp)
  div word [REAL_GLOBALS + SECTORS_PER_TRACK]
  ; Put the sector into cx (the bios call will use cl)
  mov cx, dx
  inc cx
  
  ; Calculate head and cylinder
  xor dx, dx
  ; Divide by number of heads. dx = mod (head), ax = div (cylinder)
  div word [REAL_GLOBALS + N_HEADS]
  mov dh, dl
  mov ch, al

  mov dl, byte [REAL_GLOBALS + BOOT_DRIVE]
  mov ah, 0x02
  mov al, 1
  ; Read sector
  int 0x13

.return:
  pop ebx
  pop es
  fnret

global read_sector


panic_simple:
  mov ax, 0x0003
  int 0x10
  mov word fs:[0x0000], 0x4f21
.halt:
  hlt
  jmp .halt

global panic_simple


%if ($ - $$) > 512
%error "stage 2 exceeded sector size"
%endif