global const str screen_namespace = "/dev/screen/0"
global const str mouse_namespace = "/dev/mouse/0"
global const byte BLACK = 0
global const byte WHITE = 255
global const byte DARK_GRAY = 73
global const byte GRAY = 146
global const byte LIGHT_GRAY = 182
global byte SELECTED_COLOR = 255

function main ()
  # Open screen
  plex screen $0
  str screen_name $18
  int mode $11
  nat screen_buffer $21

  # use load immediate because it a pointer to a string, not a value
  load_address &screen_namespace -> screen_name 
  load_immediate 0 -> mode
  syscall OPEN screen_name mode -> screen # Screen screen = open("/dev/screen/0", 0);

  nat width $20
  nat size $22
  load_offset_32 screen 8 -> width # load width
  load_offset_32 screen 12 -> size # load size
  load_immediate 16 -> $1 # offset for screen buffer
  add_nat screen $1 -> screen_buffer

  # open mouse
  plex mouse $15
  str mouse_name $16
  load_address &mouse_namespace -> mouse_name
  syscall OPEN mouse_name mode -> mouse # Mouse mouse = open("/dev/mouse/0", 0);

  byte color $1
  nat x_pos $12
  nat y_pos $13

  load_absolute_32 &BLACK -> color 
  load_immediate 1 -> x_pos
  load_immediate 1 -> y_pos
  call &draw_outlined_swatch screen_buffer color x_pos y_pos width

  load_absolute_32 &WHITE -> color
  load_immediate 21 -> x_pos
  load_immediate 1 -> y_pos
  call &draw_outlined_swatch screen_buffer color x_pos y_pos width 
  
  # screen.draw#
  syscall WRITE screen screen_buffer size

  nat zero $11

  loop draw_loop
    # load mouse click data
    syscall REFRESH mouse

    byte left_down $9
    load_offset_8 mouse 16 -> left_down # load btn1 pressed

    jump_eq_nat &draw_loop left_down zero

    nat mouse_x $7
    nat mouse_y $8
    load_offset_32 mouse 8 -> mouse_x # load x
    load_offset_32 mouse 12 -> mouse_y # load y

    nat box_size $14
    load_immediate 20 -> box_size

    # first row
    load_absolute_32 &BLACK -> color 
    load_immediate 1 -> x_pos
    load_immediate 1 -> y_pos
    call &draw_outlined_swatch screen_buffer color x_pos y_pos width
    call &set_color_if_clicked mouse_x mouse_y x_pos y_pos color box_size


    load_absolute_32 &WHITE -> color
    load_immediate 21 -> x_pos
    load_immediate 1 -> y_pos
    call &draw_outlined_swatch screen_buffer color x_pos y_pos width 
    call &set_color_if_clicked mouse_x mouse_y x_pos y_pos color box_size

    syscall WRITE screen screen_buffer size

    byte selected_color $25
    load_absolute_32 &SELECTED_COLOR -> selected_color

    nat brush_size $19
    load_immediate 5 -> brush_size 

    call &draw_box screen_buffer width selected_color mouse_x mouse_y brush_size brush_size

    jump &draw_loop

  # Flush and exit
  exit 0

function set_color_if_clicked (int click_x $0, int click_y $1,
  int box_x $2, int box_y $3, byte color $4, int box_size $5)

  # Compute right
  int right_edge $6
  add_int box_x box_size -> right_edge
  
  # Compute bottom = box_y + box_size  
  int bottom_edge $7
  add_int box_y box_size -> bottom_edge

  # Bounds check: x in [box_x, right] and y in [box_y, bottom]
  jump_lt_int &fail click_x box_x
  jump_ge_int &fail click_x right_edge
  jump_lt_int &fail click_y box_y
  jump_ge_int &fail click_y bottom_edge

  store_absolute_8 &SELECTED_COLOR color

  else fail
    return

function draw_outlined_swatch(nat base $0, 
  byte color $1, int x $2, int y $3, int width $4)

  # Constants
  nat background_color $5
  load_absolute_32 &GRAY -> background_color
  
  byte selected_color $10
  load_absolute_32 &SELECTED_COLOR -> selected_color

  jump_eq_int &set_selected selected_color color
  jump &end_set_selected
  set_selected:
    load_absolute_32 &DARK_GRAY -> background_color
  end_set_selected:

  nat outline_size $6
  load_immediate 20 -> outline_size

  nat fill_size $7
  load_immediate 17 -> fill_size

  nat offset $8
  load_immediate 2 -> offset

  call &draw_box base width background_color x y outline_size outline_size

  add_int x offset -> $9    # x + 2
  add_int y offset -> $10   # y + 2

  call &draw_box base width color $9 $10 fill_size fill_size

  return

function draw_box (nat base $0, nat screen_width $1, 
  byte color $2, nat x_start $3, nat y_start $4, nat width $5, nat height $6)

  # Compute start address: base + y*640 + x
  nat offset $15
  mul_int y_start screen_width -> offset
  add_int offset x_start -> offset
  add_nat offset base -> offset
  nat fat_ptr_size $25
  load_immediate 4 -> fat_ptr_size
  add_nat offset fat_ptr_size -> offset      # need to add offset for fat pointer size

  int i $30
  load_immediate 1 -> i

  int zero $26
  load_immediate 0 -> zero

  int row_end $27
  nat pixel_ptr $29

  loop draw_box_outer
    add_int offset width -> row_end        # current + width
    register_move offset -> pixel_ptr      # set pixel point
    memset_8 pixel_ptr color width         # draw row
    add_int offset screen_width -> offset  # next row += 640
    sub_int height i -> height             # decrement row count
    jump_gt_int &draw_box_outer height zero
  return