%include "defines.s"

[org BOOT1_LOADPOINT]
[bits 16]

%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_TOTAL_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 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
  hlt


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


%macro panic 1
  push word %1
  call panic_fancy
%endmacro


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

  ; Disable the cursor (don't want to look at the blink blink blink)
  mov ax, 0x0100
  mov cx, 0x3f00
  int 0x10

  mov word [GLOBALS + VGA_COL], 0x1f00

  call vga_clear

  mov ax, msg_boot1_loaded
  call vga_println

  call test_a20
  test al, al
  jnz .a20_enabled
  
  mov ax, msg_a20_disabled
  call vga_println

  mov ax, msg_a20_8042
  call vga_println

  ; 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
  panic PANIC_TYPE_A20

.a20_enabled:
  mov ax, msg_a20_enabled
  call vga_println
  
  hlt


; Print a panic message then terminate.
; Arguments:
; - word [sp]: panic type
; Does not return
panic_fancy:
  push bp
  mov bp, sp

  ; Flags first so we don't cobber them when we sub (uwu)
  pushfd       ; Temp flags: bp - 0x04
  sub sp, 16   ; Buffer:     bp - 0x14
  push dword 0 ; Registers:  bp - 0x18
  push eax
  push ebx
  push ecx
  push edx
  push esi
  push edi
  push esp
  push ebp
  xor eax, eax
  mov ax, cs
  push eax
  mov ax, ds
  push eax
  mov ax, es
  push eax
  mov ax, fs
  push eax
  mov ax, gs
  push eax
  mov ax, ss
  push eax
  mov ax, word [bp - 0x04]
  push eax

  mov ax, [bp + 0x02]
  mov [bp - 0x18], eax

  mov word [GLOBALS + VGA_COL], 0x4f00
  call vga_clear

  mov ax, VGA_WIDTH + 1
  mov cx, msg_panic
  mov dx, VGA_WIDTH - 1
  call vga_print_raw

  xor bx, bx
  mov di, VGA_WIDTH * 4
.loop_dump_regs:
  cmp bx, 16
  jae .loop_dump_regs_done

  mov si, bx
  shl si, 1
  add si, table_reg_msgs
  mov cx, [si]
  mov ax, di
  add ax, 1
  mov dx, 3
  call vga_print_raw

  push es
  mov ax, ss
  mov es, ax

  ; Format the current saved register value as hex
  lea si, [bp - 0x18]
  mov ax, bx
  shl ax, 2
  sub si, ax
  mov ecx, ss:[si]
  lea ax, [bp - 0x14]
  call dump_reg

  mov ax, di
  add ax, 5
  lea cx, [bp - 0x14]
  mov dx, 8
  call vga_print_raw

  pop es
  inc bx
  add di, VGA_WIDTH / 5
  jmp .loop_dump_regs
  
.loop_dump_regs_done:

  mov bx, [bp + 0x04]
  cmp bx, PANIC_TYPE_MAX
  jae .print_panic_type_done
  shl bx, 1
  add bx, panic_type_msgs
  mov cx, [bx]
  mov ax, VGA_WIDTH * 2 + 1
  mov dx, VGA_WIDTH - 1
  call vga_print_raw
  
.print_panic_type_done:

  ; TODO: unwind stack
  ; - Load saved bp and ip from [bp] and [bp + 2], respectively
  ; - Load bp and ip before that from [prev_bp], [prev_bp + 2]
  ; - Repeat

.halt:
  hlt
  ; Handle non-maskable interrupts
  jmp .halt


; Clear the VGA text buffer.
; Arguments: none
; Return: none
; Clobber: none
vga_clear:
  fnstart
  push di
  push es
  push ax
  push cx

  mov ax, 0xb800
  mov es, ax
  mov ax, [GLOBALS + VGA_COL]
  mov cx, VGA_WIDTH * VGA_HEIGHT
  xor di, di
  rep stosw
  mov word [GLOBALS + TEXTBUF_LINE], 0
  
  pop cx
  pop ax
  pop es
  pop di
  fnret


; Scroll the VGA text buffer up one line.
; Arguments: none
; Return: none
; Clobber: none
vga_scroll:
  fnstart
  push si
  push di
  push ax
  push cx
  push es
  push ds

  mov ax, 0xb800
  mov ds, ax
  mov es, ax

  ; Copy everything up one line.
  mov cx, VGA_WIDTH * (VGA_HEIGHT - 1)
  mov si, VGA_WIDTH * 2
  mov di, 0
  rep movsw

  ; Clear the last line.
  mov ax, [GLOBALS + VGA_COL]
  mov cx, VGA_WIDTH
  mov di, VGA_WIDTH * (VGA_HEIGHT - 1) * 2
  rep stosw

  pop ds
  pop es

  ; Decrement the current textbuf line if it's greater than 0.
  mov cx, [GLOBALS + TEXTBUF_LINE]
  xor ax, ax
  sub cx, 1
  cmovae ax, cx
  mov [GLOBALS + TEXTBUF_LINE], ax

  pop cx
  pop ax
  pop di
  pop si
  fnret


; Write a null-terminated string to the given position in the VGA text buffer.
; Arguments:
; - es: output string segment
; - ax: vga buffer index
; - cx: output string offset
; - dx: maximum length of string to print
; Return:
; - ax: vga buffer index after last character written
; Clobber: none
vga_print_raw:
  fnstart
  push fs
  push si
  push di
  push cx
  push dx

  mov si, cx
  xchg ax, cx
  
  ; Find the distance between the starting index and the end of the buffer.
  mov ax, (VGA_WIDTH * VGA_HEIGHT)
  sub ax, cx
  ; If the starting index is past the end of the buffer, return early.
  jc .done
  ; Clamp the maximum length to the distance between the starting index and the end of the buffer.
  cmp ax, dx
  cmovb dx, ax

  mov di, cx
  shl di, 1
  
  mov ax, 0xb800
  mov fs, ax
  mov ah, [GLOBALS + VGA_COL + 1]

.loop:
  test dx, dx
  jz .done
  dec dx
  mov al, es:[si]
  test al, al
  jz .done
  mov fs:[di], ax
  add di, 2
  inc si
  inc cx
  jmp .loop

.done:
  xchg ax, cx

  pop dx
  pop cx
  pop di
  pop si
  pop fs
  fnret


; Write one line to the VGA text buffer. The string should be null-terminated; we embrace the evil
; of null-termination so this function only takes one argument, so it's slightly less of a faff to
; call in most cases.
; Arguments:
; - es: output string segment
; - ax: output string offset
; Return: none
; Clobber: none
vga_println:
  fnstart
  push ax
  push bx
  push cx
  push dx

  cmp word [GLOBALS + TEXTBUF_LINE], VGA_HEIGHT
  jb .scroll_done
  call vga_scroll
.scroll_done:

  mov bx, ax
  
  xor dx, dx
  mov ax, [GLOBALS + TEXTBUF_LINE]
  mov cx, VGA_WIDTH
  mul cx

  mov dx, VGA_WIDTH
  mov cx, bx

  call vga_print_raw

  inc word [GLOBALS + TEXTBUF_LINE]
  
  pop dx
  pop cx
  pop bx
  pop ax
  fnret


; Convert the value in ecx to hex and write it to the buffer at es:ax. The buffer should be at
; least 8 bytes long.
; - es: output buffer segment
; - ax: output buffer offset
; - ecx: value to convert to hex and print
; Return:
; - ax: the address one after the last byte that was written
; Clobber: none
dump_reg:
  fnstart
  push bx
  push dx
  push ecx
  
  mov bx, ax
  mov dx, 4

.loop:
  test dx, dx
  jz .done
  dec dx

  rol ecx, 8

  mov al, cl
  shr al, 4
  call nybble_to_hex_char
  mov es:[bx], al
  inc bx
  
  mov al, cl
  call nybble_to_hex_char
  mov es:[bx], al
  inc bx

  jmp .loop

.done:
  mov ax, bx

  pop ecx
  pop dx
  pop bx
  fnret


; Convert nybble to lowercase ascii hex char.
; Arguments:
; - al: value to convert
; Return:
; - al: converted ascii hex value
; Clobber: none
nybble_to_hex_char:
  ; We don't use the stack, so no need to change bp.
  and al, 0x0f
  cmp al, 9
  jbe .0_to_9
  add al, (0x61 - 0x0a)
  jmp .done
.0_to_9:
  add al, 0x30
.done:
  ret


; 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, (1 << INTEL_8042_STATUS_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, (1 << INTEL_8042_STATUS_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


msg_boot1_loaded db "boot1 loaded. hello!", 0
msg_a20_enabled  db "a20 enabled", 0
msg_a20_disabled db "a20 not enabled", 0
msg_a20_8042     db "trying 8042", 0
msg_panic        db "panic!", 0

msg_reg_eip   db "eip", 0
msg_reg_eax   db "eax", 0
msg_reg_ebx   db "ebx", 0
msg_reg_ecx   db "ecx", 0
msg_reg_edx   db "edx", 0
msg_reg_esi   db "esi", 0
msg_reg_edi   db "edi", 0
msg_reg_esp   db "esp", 0
msg_reg_ebp   db "ebp", 0
msg_reg_cs    db "cs", 0
msg_reg_ds    db "ds", 0
msg_reg_es    db "es", 0
msg_reg_fs    db "fs", 0
msg_reg_gs    db "gs", 0
msg_reg_ss    db "ss", 0
msg_reg_flags db "flg", 0

table_reg_msgs:
  dw msg_reg_eip
  dw msg_reg_eax
  dw msg_reg_ebx
  dw msg_reg_ecx
  dw msg_reg_edx
  dw msg_reg_esi
  dw msg_reg_edi
  dw msg_reg_esp
  dw msg_reg_ebp
  dw msg_reg_cs
  dw msg_reg_ds
  dw msg_reg_es
  dw msg_reg_fs
  dw msg_reg_gs
  dw msg_reg_ss
  dw msg_reg_flags
  dw 0

msg_panic_generic db "generic panic", 0
msg_panic_a20     db "failed to enable a20 line", 0

panic_type_msgs:
  PANIC_TYPE_GENERIC equ ($ - panic_type_msgs) / 2
  dw msg_panic_generic
  PANIC_TYPE_A20     equ ($ - panic_type_msgs) / 2
  dw msg_panic_a20
  PANIC_TYPE_MAX     equ ($ - panic_type_msgs) / 2
  dw 0

boot1_magic dd BOOT1_MAGIC

BOOT1_TOTAL_LEN     equ $ - $$
BOOT1_TOTAL_SECTORS equ (BOOT1_TOTAL_LEN + 511) / 512

%if (BOOT1_LOADPOINT + BOOT1_TOTAL_LEN) > EBDA_START
%error "boot1 too large to be loaded"
%endif