[bits 16]

%include "fn.s"
%include "ps2.s"

%macro mov_out 3
  mov %1, %3
  out %2, %1
%endmacro

; 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

global test_a20


; 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_out al, INTEL_8042_OUT_CMD, INTEL_8042_CMD_PS2_1_DISABLE

  ; Read the controller output port.
  call intel_8042_wait_write
  mov_out al, INTEL_8042_OUT_CMD, INTEL_8042_CMD_CONTROLLER_OUT_PORT_READ
  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_out al, INTEL_8042_OUT_CMD, INTEL_8042_CMD_CONTROLLER_OUT_PORT_WRITE
  call intel_8042_wait_write
  mov_out al, INTEL_8042_IO_DATA, cl

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

  ; Wait for writes to finish.
  call intel_8042_wait_write

  ret

global enable_a20_intel_8042


; 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