%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


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

  mov eax, 0xa1b2c3d4
  mov ebx, 0x12345678
  call panic_fancy

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

  ; TODO: enable a20
  hlt

.a20_enabled:
  mov ax, msg_a20_enabled
  call vga_println
  
  hlt


; Wrapper function around panic_fancy which provides default arguments.
; Arguments: none
; Does not return
panic_default:
  push bp
  mov bp, sp
  push word 0
  push word 0
  ; Spoof return address
  push dword [bp]
  jmp panic_fancy


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

  sub sp, 16 ; Buffer: bp - 16

  push eax   ; bp - 20
  push ebx   ; bp - 24
  push ecx   ; bp - 28
  push edx   ; bp - 32
  push esi   ; bp - 36
  push edi   ; bp - 40

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

  mov ax, ss
  mov es, ax
  
  lea ax, [bp - 16]
  push ax
  mov ecx, [bp - 20]
  call dump_reg
  pop cx
  xor ax, ax
  mov dx, 8
  call vga_print_raw

  lea ax, [bp - 16]
  push ax
  mov ecx, [bp - 24]
  call dump_reg
  pop cx
  mov ax, VGA_WIDTH
  mov dx, 8
  call vga_print_raw

.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


enable_a20:
  fnstart
  fnret


; 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


msg_boot1_loaded db "boot1 loaded. hello!", 0
msg_a20_enabled  db "a20 enabled", 0
msg_a20_disabled db "a20 not enabled", 0
msg_panic        db "panic!", 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