%include "defines.s"
%include "ps2.s"

[bits 16]

extern boot1_bin_len
extern boot1_bin_sectors
extern boot1_magic
extern _start

extern gdt_flat_slice
extern gdt_flat
extern GDT_FLAT_IDX_CODE_32
extern GDT_FLAT_IDX_CODE_16
extern GDT_FLAT_IDX_DATA

section .prelude

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

; boot0 loads only our first sector into memory. We must load the rest.
self_load:
  ; Now that we're not doing instruction byte golf like we were in boot0, we can afford to move
  ; the various boot0 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, BOOT1_GPT_ENTRY_ADDR

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

  mov si, [GLOBALS + BOOT1_GPT_ENTRY_ADDR]
  mov eax, [si + 0x20] ; Partition / boot1 start LBA lower
  mov ebx, [si + 0x24] ; Partition / boot1 start LBA upper
  mov ecx, [si + 0x28] ; Partition end LBA lower
  mov edx, [si + 0x32] ; Partition 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

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

  ; Panic if the boot1 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 boot1 end LBA.
  or edx, edx
  jnz .end_lba_ok
  ; Compare the boot1 end LBA to the lower 32 bits of the partition end LBA.
  cmp ebx, ecx
  ja panic_simple

.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                        ; boot1 end LBA
  mov ebx, BOOT1_LOADPOINT + 512 ; Current sector load address

.self_load_loop:
  mov ax, [bp - 0x02]      ; Load current LBA
  cmp word [bp - 0x04], ax ; Compare to boot1 end LBA
  jb .self_load_done

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

.self_load_done:

  ; Check the magic bytes at the end of boot1.
  push es
  mov ebx, boot1_magic
  call addr32_to_addr16
  cmp dword es:[bx], BOOT1_MAGIC
  pop es
  jne panic_simple
  
  jmp prelude_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
  

; 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 [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 [GLOBALS + N_HEADS]
  mov dh, dl
  mov ch, al

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

.return:
  pop ebx
  pop es
  fnret


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


%if ($ - $$) > 512
%error "boot1 self-loader exceeded sector size"
%endif


; Check whether the A20 line is enabled. Writes to the boot sector identifier.
; Arguments: none
; Return:
; - ax: 0 if A20 disabled, nonzero if A20 enabled
; Clobber: none
test_a20:
  push bp
  mov bp, sp
  push gs
  
  ; Restore the boot sector identifier in case it was overwritten by anything.
  mov word [0x7dfe], 0xaa55

  mov ax, 0xffff
  mov gs, ax
  xor ax, ax

  ; If the word at 0x107dfe (1 MiB after the boot sector identifier) is different to the boot
  ; sector identifier, than A20 must be enabled.
  cmp word gs:[0x7e0e], 0xaa55
  setne al
  jne .return

  ; Even if A20 was enabled, the two words may have been equal by chance, so we temporarily swap
  ; the boot sector identifier bytes and test again.
  ror word [0x7dfe], 8
  cmp word gs:[0x7e0e], 0x55aa
  setne al
  ror word [0x7dfe], 8
  jmp .return

.return:
  pop gs
  pop bp
  ret


; Wait for the Intel 8042 input buffer to become empty, so we can write.
; Arguments: none
; Return: none
; Clobber: al
intel_8042_wait_write:
.loop:
  ; Read the 8042 status register.
  in al, INTEL_8042_IN_STATUS
  ; Input buffer status flag set means the input buffer is full, so loop in this case.
  test al, INTEL_8042_STATUS_MASK_IBUF
  jnz .loop
  ret


; Wait for the Intel 8042 output buffer to become filled, so we can read.
; Arguments: none
; Return: none
; Clobber: al
intel_8042_wait_read:
.loop:
  ; Read the 8042 status register.
  in al, INTEL_8042_IN_STATUS
  ; Output buffer status flag unset means output buffer is empty, so loop in this case.
  test al, INTEL_8042_STATUS_MASK_OBUF
  jz .loop
  ret


; Try to enable A20 using the Intel 8042 PS/2 keyboard controller.
; Arguments: none
; Return: none
; Clobber: ax, cx, dx
enable_a20_intel_8042:
  ; Temporarily disable the keyboard.
  call intel_8042_wait_write
  mov al, INTEL_8042_CMD_PS2_1_DISABLE
  out INTEL_8042_OUT_CMD, al

  ; Read the controller output port.
  call intel_8042_wait_write
  mov al, INTEL_8042_CMD_CONTROLLER_OUT_PORT_READ
  out INTEL_8042_OUT_CMD, al
  call intel_8042_wait_read
  in al, INTEL_8042_IO_DATA

  ; The second bit is "A20 enabled", so set it.
  mov cl, al
  or cl, 2

  ; Write the modified byte back to the controller output port.
  call intel_8042_wait_write
  mov al, INTEL_8042_CMD_CONTROLLER_OUT_PORT_WRITE
  out INTEL_8042_OUT_CMD, al
  call intel_8042_wait_write
  mov al, cl
  out INTEL_8042_IO_DATA, al

  ; Re-enable the keyboard.
  call intel_8042_wait_write
  mov al, INTEL_8042_CMD_PS2_1_ENABLE
  out INTEL_8042_OUT_CMD, al

  ; Wait for writes to finish.
  call intel_8042_wait_write

  ret


prelude_main:
  call test_a20
  test al, al
  jnz .a20_enabled

  ; Try to enable A20 using the Intel 8042 PS/2 keyboard controller.
  call enable_a20_intel_8042
  call test_a20
  test al, al
  jnz .a20_enabled

  ; TODO: try other methods first before we panic:
  ; - [ ] BIOS interrupt
  ; - [ ] Fast A20 enable
  jmp panic_simple

.a20_enabled:
  mov ax, 0x0003
  int 0x10

  ; Ensure interrupts are definitely disabled.
  cli

  ; Load our flat-address-space GDT.
  lgdt [gdt_flat_slice]

  ; Set the protected-mode bit in cr0.
  mov eax, cr0
  or al, 0x01
  mov cr0, eax

  ; Long jump to set the code segment to gdt_flat.segment_code, and to clear the instruction
  ; pipeline.
  jmp GDT_FLAT_IDX_CODE_32:.protected_mode_32

[bits 32]
.protected_mode_32:

  ; Set the data segments to gdt_flat.segment_data.
  mov eax, GDT_FLAT_IDX_DATA
  mov ds, eax
  mov es, eax
  mov fs, eax
  mov gs, eax
  mov ss, eax

  ; Reset the stack.
  mov ebp, STACK_BASE
  mov esp, ebp

  jmp _start
  jmp panic_simple