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 is $0
  str screen_name is $18
  int mode is $11
  nat screen_buffer is $21

  # use load immediate because it is 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 is $20
  nat size is $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 is $15
  str mouse_name is $16
  load_address &mouse_namespace -> mouse_name
  syscall OPEN mouse_name mode -> mouse # Mouse mouse = open("/dev/mouse/0", 0);

  byte color is $1
  nat x_pos is $12
  nat y_pos is $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 is $11

  draw_loop:
    # load mouse click data
    syscall REFRESH mouse

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

    jump_eq_nat &draw_loop left_down zero

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

    nat box_size is $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 is $25
    load_absolute_32 &SELECTED_COLOR -> selected_color

    nat brush_size is $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 is $0, int click_y is $1,
  int box_x is $2, int box_y is $3, byte color is $4, int box_size is $5)

  # Compute right
  int right_edge is $6
  add_int box_x box_size -> right_edge
  
  # Compute bottom = box_y + box_size  
  int bottom_edge is $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

  fail:
    return

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

  # Constants
  nat background_color is $5
  load_absolute_32 &GRAY -> background_color
  
  byte selected_color is $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 is $6
  load_immediate 20 -> outline_size

  nat fill_size is $7
  load_immediate 17 -> fill_size

  nat offset is $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 is $0, nat screen_width is $1, 
  byte color is $2, nat x_start is $3, nat y_start is $4, nat width is $5, nat height is $6)

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

  int i is $30
  load_immediate 1 -> i

  int zero is $26
  load_immediate 0 -> zero

  int row_end is $27
  nat pixel_ptr is $29

  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