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
  do set_selected
    load_absolute_32 DARK_GRAY -> background_color
  else 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