#include <Arduino.h>
#include "WASM.h"

WASM vm;

/**
;; ------------------------------------------------------------
;;
;; LEDBounce_F32.wat
;; (c) 2025 RIoT Secure AB
;;
;; https://www.hackster.io/Ripred/bouncing-balls-sketch-on-uno-r4-wifi-led-matrix-d0657e
;;
;; This module implements the LED Bounce example on hackster.io
;; for the Arduino UNO R4. The sketch uses WebAssembly host
;; functions. It generates 5 random balls and movement, then
;; proceeds to animate them with a small delay between.
;;
;; F32 - float math (optimal for MCU with FPU)
;;
;; Host functions expected:
;;   (import "brawl:device/matrix" "begin"        (func ...))
;;   (import "brawl:device/matrix" "loadPixels"   (func ...))
;;   (import "brawl:device/matrix" "renderFrame"  (func ...))
;;   (import "brawl:device/sys"    "delay"        (func ...))
;;
;; Entry point:
;;   (start $start) -- calls setup() once, then repeatedly loop()
;;
(module
  ;; ---- imports ----
  (import "brawl:device/matrix" "begin"        (func $begin))
  (import "brawl:device/matrix" "loadPixels"   (func $load (param i32 i32 i32)))
  (import "brawl:device/matrix" "renderFrame"  (func $render))
  (import "brawl:device/sys"    "delay"        (func $delay (param i32)))

  ;; ---- memory ----
  (memory (export "mem") 1)

  ;; ---- constants ----
  (global $MAX_X  i32 (i32.const 12))
  (global $MAX_Y  i32 (i32.const 8))     ;; matrix has 96 LEDs
  (global $POINTS i32 (i32.const 5))

  ;; ---- constants :: memory locations ----
  (global $RNG    i32 (i32.const 0))     ;; [0x00 - 0x03] - random seed            ( 4 bytes)
  (global $FB     i32 (i32.const 4))     ;; [0x04 - 0x63] - LED representation     (96 bytes)
  (global $STATE  i32 (i32.const 100))   ;; [0x64 - 0xB3] - dot struct (x,y,dx,dy) (80 bytes)

  ;; ---- random number generator ----
  (func $rand32 (result i32)
    (local $s i32)
    (local.set $s (i32.load (global.get $RNG)))
    (local.set $s
      (i32.add
        (i32.mul (local.get $s) (i32.const 1664525))
        (i32.const 1013904223)))
    (i32.store (global.get $RNG) (local.get $s))
    (local.get $s)
  )

  ;; ---- start function which includes the setup and loop ----
  (func $start
    (local $i i32) (local $j i32) (local $base i32)
    (local $x f32) (local $y f32) (local $dx f32) (local $dy f32)
    (local $r i32) (local $tmp i32) (local $idx i32)
    (local $MAXXf f32) (local $MAXYf f32) (local $xi i32) (local $yi i32)

    ;; device init
    (call $begin)

    ;; set the initial random seed
    (i32.store (global.get $RNG) (i32.const 0x12345678))

    ;; bounds in f32: [0, MAXX), [0, MAXY)
    (local.set $MAXXf (f32.convert_i32_s (global.get $MAX_X)))
    (local.set $MAXYf (f32.convert_i32_s (global.get $MAX_Y)))

    ;; initialize points
    (local.set $i (i32.const 0))
    (loop $INIT

      ;; initial position for the point in x [0, MAX_X)
      (local.set $r (call $rand32))
      (local.set $r (i32.and (local.get $r) (i32.const 0x7fffffff)))
      (local.set $x (f32.convert_i32_u (i32.rem_u (local.get $r) (global.get $MAX_X))))

      ;; initial position for the point in y [0, MAX_Y)
      (local.set $r (call $rand32))
      (local.set $r (i32.and (local.get $r) (i32.const 0x7fffffff)))
      (local.set $y (f32.convert_i32_u (i32.rem_u (local.get $r) (global.get $MAX_Y))))

      ;; dx = ± (1 / {2,3,4})
      (local.set $r (call $rand32))
      (local.set $r (i32.and (local.get $r) (i32.const 0x7fffffff)))
      (local.set $tmp (i32.add (i32.const 2) (i32.rem_u (local.get $r) (i32.const 3))))
      (local.set $dx (f32.div (f32.const 1) (f32.convert_i32_u (local.get $tmp))))
      (if (i32.and (call $rand32) (i32.const 1))
        (then)
        (else (local.set $dx (f32.neg (local.get $dx)))))

      ;; dy = ± (1 / {2,3,4})
      (local.set $r (call $rand32))
      (local.set $r (i32.and (local.get $r) (i32.const 0x7fffffff)))
      (local.set $tmp (i32.add (i32.const 2) (i32.rem_u (local.get $r) (i32.const 3))))
      (local.set $dy (f32.div (f32.const 1) (f32.convert_i32_u (local.get $tmp))))
      (if (i32.and (call $rand32) (i32.const 1))
        (then)
        (else (local.set $dy (f32.neg (local.get $dy)))))

      ;; store point
      (local.set $base (i32.add (global.get $STATE) (i32.mul (local.get $i) (i32.const 16))))
      (f32.store (i32.add (local.get $base) (i32.const 0))  (local.get $x))
      (f32.store (i32.add (local.get $base) (i32.const 4))  (local.get $y))
      (f32.store (i32.add (local.get $base) (i32.const 8))  (local.get $dx))
      (f32.store (i32.add (local.get $base) (i32.const 12)) (local.get $dy))

      (local.set $i (i32.add (local.get $i) (i32.const 1)))
      (br_if $INIT (i32.lt_u (local.get $i) (global.get $POINTS)))
    )

    ;; animation loop
    (loop $FOREVER

      ;; clear buffer
      (local.set $j (i32.const 0))
      (loop $CLR
        (i32.store8 (i32.add (global.get $FB) (local.get $j)) (i32.const 0))
        (local.set $j (i32.add (local.get $j) (i32.const 1)))
        (br_if $CLR (i32.lt_u (local.get $j) (i32.const 96)))
      )

      ;; draw points
      (local.set $i (i32.const 0))
      (loop $DRAW
        (local.set $base (i32.add (global.get $STATE) (i32.mul (local.get $i) (i32.const 16))))
        (local.set $x (f32.load (i32.add (local.get $base) (i32.const 0))))
        (local.set $y (f32.load (i32.add (local.get $base) (i32.const 4))))

        ;; clamp to boundaries - truncate f32 to i32
        (local.set $xi
          (i32.trunc_f32_s
            (f32.min (f32.max (local.get $x) (f32.const 0))
                     (f32.sub (local.get $MAXXf) (f32.const 1)))))
        (local.set $yi
          (i32.trunc_f32_s
            (f32.min (f32.max (local.get $y) (f32.const 0))
                     (f32.sub (local.get $MAXYf) (f32.const 1)))))
        (local.set $idx
          (i32.add (global.get $FB)
            (i32.add (local.get $xi) (i32.mul (local.get $yi) (global.get $MAX_X)))))
        (i32.store8 (local.get $idx) (i32.const 1))

        ;; turn on the LED
        (local.set $i (i32.add (local.get $i) (i32.const 1)))
        (br_if $DRAW (i32.lt_u (local.get $i) (global.get $POINTS)))
      )

      ;; render to LED display
      (call $load (global.get $FB) (global.get $MAX_X) (global.get $MAX_Y))
      (call $render)
      (call $delay (i32.const 50)) ;; delay between processing

      ;; update with bounce physics
      (local.set $i (i32.const 0))
      (loop $UPD
        (local.set $base (i32.add (global.get $STATE) (i32.mul (local.get $i) (i32.const 16))))
        (local.set $x  (f32.load (i32.add (local.get $base) (i32.const 0))))
        (local.set $y  (f32.load (i32.add (local.get $base) (i32.const 4))))
        (local.set $dx (f32.load (i32.add (local.get $base) (i32.const 8))))
        (local.set $dy (f32.load (i32.add (local.get $base) (i32.const 12))))

        ;; reflect if next or previous step leaves bounds
        (if
          (i32.or
            (i32.or
              (f32.lt (f32.add (local.get $x) (local.get $dx)) (f32.const 0))
              (f32.lt (f32.sub (local.get $x) (local.get $dx)) (f32.const 0)))
            (i32.or
              (f32.ge (f32.add (local.get $x) (local.get $dx)) (local.get $MAXXf))
              (f32.ge (f32.sub (local.get $x) (local.get $dx)) (local.get $MAXXf))))
          (then (local.set $dx (f32.neg (local.get $dx))))
        )
        (if
          (i32.or
            (i32.or
              (f32.lt (f32.add (local.get $y) (local.get $dy)) (f32.const 0))
              (f32.lt (f32.sub (local.get $y) (local.get $dy)) (f32.const 0)))
            (i32.or
              (f32.ge (f32.add (local.get $y) (local.get $dy)) (local.get $MAXYf))
              (f32.ge (f32.sub (local.get $y) (local.get $dy)) (local.get $MAXYf))))
          (then (local.set $dy (f32.neg (local.get $dy))))
        )

        (local.set $x (f32.add (local.get $x) (local.get $dx)))
        (local.set $y (f32.add (local.get $y) (local.get $dy)))

        ;; store the values
        (f32.store (i32.add (local.get $base) (i32.const 0))  (local.get $x))
        (f32.store (i32.add (local.get $base) (i32.const 4))  (local.get $y))
        (f32.store (i32.add (local.get $base) (i32.const 8))  (local.get $dx))
        (f32.store (i32.add (local.get $base) (i32.const 12)) (local.get $dy))

        (local.set $i (i32.add (local.get $i) (i32.const 1)))
        (br_if $UPD (i32.lt_u (local.get $i) (global.get $POINTS)))
      )

      (br $FOREVER)
    )
  )

  (export "start" (func $start))
  (start $start)
)

;; ------------------------------------------------------------
**/

const uint8_t wasm_bytecode[] =
{
  0x00, 0x61, 0x73, 0x6d, 0x01, 0x00, 0x00, 0x00, 0x01, 0x12, 0x04, 0x60, 0x00, 0x00, 0x60, 0x03,
  0x7f, 0x7f, 0x7f, 0x00, 0x60, 0x01, 0x7f, 0x00, 0x60, 0x00, 0x01, 0x7f, 0x02, 0x79, 0x04, 0x13,
  0x62, 0x72, 0x61, 0x77, 0x6c, 0x3a, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x6d, 0x61, 0x74,
  0x72, 0x69, 0x78, 0x05, 0x62, 0x65, 0x67, 0x69, 0x6e, 0x00, 0x00, 0x13, 0x62, 0x72, 0x61, 0x77,
  0x6c, 0x3a, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x6d, 0x61, 0x74, 0x72, 0x69, 0x78, 0x0a,
  0x6c, 0x6f, 0x61, 0x64, 0x50, 0x69, 0x78, 0x65, 0x6c, 0x73, 0x00, 0x01, 0x13, 0x62, 0x72, 0x61,
  0x77, 0x6c, 0x3a, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x6d, 0x61, 0x74, 0x72, 0x69, 0x78,
  0x0b, 0x72, 0x65, 0x6e, 0x64, 0x65, 0x72, 0x46, 0x72, 0x61, 0x6d, 0x65, 0x00, 0x00, 0x10, 0x62,
  0x72, 0x61, 0x77, 0x6c, 0x3a, 0x64, 0x65, 0x76, 0x69, 0x63, 0x65, 0x2f, 0x73, 0x79, 0x73, 0x05,
  0x64, 0x65, 0x6c, 0x61, 0x79, 0x00, 0x02, 0x03, 0x03, 0x02, 0x03, 0x00, 0x05, 0x03, 0x01, 0x00,
  0x01, 0x06, 0x20, 0x06, 0x7f, 0x00, 0x41, 0x0c, 0x0b, 0x7f, 0x00, 0x41, 0x08, 0x0b, 0x7f, 0x00,
  0x41, 0x05, 0x0b, 0x7f, 0x00, 0x41, 0x00, 0x0b, 0x7f, 0x00, 0x41, 0x04, 0x0b, 0x7f, 0x00, 0x41,
  0xe4, 0x00, 0x0b, 0x07, 0x0f, 0x02, 0x03, 0x6d, 0x65, 0x6d, 0x02, 0x00, 0x05, 0x73, 0x74, 0x61,
  0x72, 0x74, 0x00, 0x05, 0x08, 0x01, 0x05, 0x0a, 0xa6, 0x05, 0x02, 0x25, 0x01, 0x01, 0x7f, 0x23,
  0x03, 0x28, 0x02, 0x00, 0x21, 0x00, 0x20, 0x00, 0x41, 0x8d, 0xcc, 0xe5, 0x00, 0x6c, 0x41, 0xdf,
  0xe6, 0xbb, 0xe3, 0x03, 0x6a, 0x21, 0x00, 0x23, 0x03, 0x20, 0x00, 0x36, 0x02, 0x00, 0x20, 0x00,
  0x0b, 0xfd, 0x04, 0x05, 0x03, 0x7f, 0x04, 0x7d, 0x03, 0x7f, 0x02, 0x7d, 0x02, 0x7f, 0x10, 0x00,
  0x23, 0x03, 0x41, 0xf8, 0xac, 0xd1, 0x91, 0x01, 0x36, 0x02, 0x00, 0x23, 0x00, 0xb2, 0x21, 0x0a,
  0x23, 0x01, 0xb2, 0x21, 0x0b, 0x41, 0x00, 0x21, 0x00, 0x03, 0x40, 0x10, 0x04, 0x21, 0x07, 0x20,
  0x07, 0x41, 0xff, 0xff, 0xff, 0xff, 0x07, 0x71, 0x21, 0x07, 0x20, 0x07, 0x23, 0x00, 0x70, 0xb3,
  0x21, 0x03, 0x10, 0x04, 0x21, 0x07, 0x20, 0x07, 0x41, 0xff, 0xff, 0xff, 0xff, 0x07, 0x71, 0x21,
  0x07, 0x20, 0x07, 0x23, 0x01, 0x70, 0xb3, 0x21, 0x04, 0x10, 0x04, 0x21, 0x07, 0x20, 0x07, 0x41,
  0xff, 0xff, 0xff, 0xff, 0x07, 0x71, 0x21, 0x07, 0x41, 0x02, 0x20, 0x07, 0x41, 0x03, 0x70, 0x6a,
  0x21, 0x08, 0x43, 0x00, 0x00, 0x80, 0x3f, 0x20, 0x08, 0xb3, 0x95, 0x21, 0x05, 0x10, 0x04, 0x41,
  0x01, 0x71, 0x04, 0x40, 0x05, 0x20, 0x05, 0x8c, 0x21, 0x05, 0x0b, 0x10, 0x04, 0x21, 0x07, 0x20,
  0x07, 0x41, 0xff, 0xff, 0xff, 0xff, 0x07, 0x71, 0x21, 0x07, 0x41, 0x02, 0x20, 0x07, 0x41, 0x03,
  0x70, 0x6a, 0x21, 0x08, 0x43, 0x00, 0x00, 0x80, 0x3f, 0x20, 0x08, 0xb3, 0x95, 0x21, 0x06, 0x10,
  0x04, 0x41, 0x01, 0x71, 0x04, 0x40, 0x05, 0x20, 0x06, 0x8c, 0x21, 0x06, 0x0b, 0x23, 0x05, 0x20,
  0x00, 0x41, 0x10, 0x6c, 0x6a, 0x21, 0x02, 0x20, 0x02, 0x41, 0x00, 0x6a, 0x20, 0x03, 0x38, 0x02,
  0x00, 0x20, 0x02, 0x41, 0x04, 0x6a, 0x20, 0x04, 0x38, 0x02, 0x00, 0x20, 0x02, 0x41, 0x08, 0x6a,
  0x20, 0x05, 0x38, 0x02, 0x00, 0x20, 0x02, 0x41, 0x0c, 0x6a, 0x20, 0x06, 0x38, 0x02, 0x00, 0x20,
  0x00, 0x41, 0x01, 0x6a, 0x21, 0x00, 0x20, 0x00, 0x23, 0x02, 0x49, 0x0d, 0x00, 0x0b, 0x03, 0x40,
  0x41, 0x00, 0x21, 0x01, 0x03, 0x40, 0x23, 0x04, 0x20, 0x01, 0x6a, 0x41, 0x00, 0x3a, 0x00, 0x00,
  0x20, 0x01, 0x41, 0x01, 0x6a, 0x21, 0x01, 0x20, 0x01, 0x41, 0xe0, 0x00, 0x49, 0x0d, 0x00, 0x0b,
  0x41, 0x00, 0x21, 0x00, 0x03, 0x40, 0x23, 0x05, 0x20, 0x00, 0x41, 0x10, 0x6c, 0x6a, 0x21, 0x02,
  0x20, 0x02, 0x41, 0x00, 0x6a, 0x2a, 0x02, 0x00, 0x21, 0x03, 0x20, 0x02, 0x41, 0x04, 0x6a, 0x2a,
  0x02, 0x00, 0x21, 0x04, 0x20, 0x03, 0x43, 0x00, 0x00, 0x00, 0x00, 0x97, 0x20, 0x0a, 0x43, 0x00,
  0x00, 0x80, 0x3f, 0x93, 0x96, 0xa8, 0x21, 0x0c, 0x20, 0x04, 0x43, 0x00, 0x00, 0x00, 0x00, 0x97,
  0x20, 0x0b, 0x43, 0x00, 0x00, 0x80, 0x3f, 0x93, 0x96, 0xa8, 0x21, 0x0d, 0x23, 0x04, 0x20, 0x0c,
  0x20, 0x0d, 0x23, 0x00, 0x6c, 0x6a, 0x6a, 0x21, 0x09, 0x20, 0x09, 0x41, 0x01, 0x3a, 0x00, 0x00,
  0x20, 0x00, 0x41, 0x01, 0x6a, 0x21, 0x00, 0x20, 0x00, 0x23, 0x02, 0x49, 0x0d, 0x00, 0x0b, 0x23,
  0x04, 0x23, 0x00, 0x23, 0x01, 0x10, 0x01, 0x10, 0x02, 0x41, 0x32, 0x10, 0x03, 0x41, 0x00, 0x21,
  0x00, 0x03, 0x40, 0x23, 0x05, 0x20, 0x00, 0x41, 0x10, 0x6c, 0x6a, 0x21, 0x02, 0x20, 0x02, 0x41,
  0x00, 0x6a, 0x2a, 0x02, 0x00, 0x21, 0x03, 0x20, 0x02, 0x41, 0x04, 0x6a, 0x2a, 0x02, 0x00, 0x21,
  0x04, 0x20, 0x02, 0x41, 0x08, 0x6a, 0x2a, 0x02, 0x00, 0x21, 0x05, 0x20, 0x02, 0x41, 0x0c, 0x6a,
  0x2a, 0x02, 0x00, 0x21, 0x06, 0x20, 0x03, 0x20, 0x05, 0x92, 0x43, 0x00, 0x00, 0x00, 0x00, 0x5d,
  0x20, 0x03, 0x20, 0x05, 0x93, 0x43, 0x00, 0x00, 0x00, 0x00, 0x5d, 0x72, 0x20, 0x03, 0x20, 0x05,
  0x92, 0x20, 0x0a, 0x60, 0x20, 0x03, 0x20, 0x05, 0x93, 0x20, 0x0a, 0x60, 0x72, 0x72, 0x04, 0x40,
  0x20, 0x05, 0x8c, 0x21, 0x05, 0x0b, 0x20, 0x04, 0x20, 0x06, 0x92, 0x43, 0x00, 0x00, 0x00, 0x00,
  0x5d, 0x20, 0x04, 0x20, 0x06, 0x93, 0x43, 0x00, 0x00, 0x00, 0x00, 0x5d, 0x72, 0x20, 0x04, 0x20,
  0x06, 0x92, 0x20, 0x0b, 0x60, 0x20, 0x04, 0x20, 0x06, 0x93, 0x20, 0x0b, 0x60, 0x72, 0x72, 0x04,
  0x40, 0x20, 0x06, 0x8c, 0x21, 0x06, 0x0b, 0x20, 0x03, 0x20, 0x05, 0x92, 0x21, 0x03, 0x20, 0x04,
  0x20, 0x06, 0x92, 0x21, 0x04, 0x20, 0x02, 0x41, 0x00, 0x6a, 0x20, 0x03, 0x38, 0x02, 0x00, 0x20,
  0x02, 0x41, 0x04, 0x6a, 0x20, 0x04, 0x38, 0x02, 0x00, 0x20, 0x02, 0x41, 0x08, 0x6a, 0x20, 0x05,
  0x38, 0x02, 0x00, 0x20, 0x02, 0x41, 0x0c, 0x6a, 0x20, 0x06, 0x38, 0x02, 0x00, 0x20, 0x00, 0x41,
  0x01, 0x6a, 0x21, 0x00, 0x20, 0x00, 0x23, 0x02, 0x49, 0x0d, 0x00, 0x0b, 0x0c, 0x00, 0x0b, 0x0b

};
const size_t  wasm_bytecode_len = 896;

#ifdef __cplusplus
extern "C" {
#endif

// we must define a host printf for WASM_OUT
void hprintf(const char *fmt, ...)
{
    char buf[128];
    va_list args;

    va_start(args, fmt);
    vsnprintf(buf, sizeof(buf), fmt, args);
    va_end(args);

    Serial.print(buf);
    Serial.flush();
}

#ifdef __cplusplus
}
#endif

void setup()
{
    wasm_error_t err;

    Serial.begin(115200);
    while (!Serial) ;

    Serial.println("wasm (runtime)");
    Serial.println("Copyright 2025, RIoT Secure AB");
    Serial.println();

    err = vm.initialize();
    if (err)
    {
        Serial.print("ERROR: vm.initialize() : "); Serial.println(vm.getWASMErrorMessage(err));
        return;
    }

    err = vm.loadModule(wasm_bytecode, wasm_bytecode_len);
    if (err)
    {
        Serial.print("ERROR: - vm.loadModule() : "); Serial.println(vm.getWASMErrorMessage(err));
        return;
    }

    err = vm.validate();
    if (err)
    {
        Serial.print("ERROR: - vm.validate() : "); Serial.println(vm.getWASMErrorMessage(err));
        return;
    }

    err = vm.exec();
    if (err)
    {
        Serial.print("ERROR: - vm.exec() : "); Serial.println(vm.getWASMErrorMessage(err));
        return;
    }

    err = vm.unloadModule();
    if (err)
    {
        Serial.print("ERROR: - vm.unloadModule() : "); Serial.println(vm.getWASMErrorMessage(err));
        return;
    }

    err = vm.terminate();
    if (err)
    {
        Serial.print("ERROR: - vm.terminate() : "); Serial.println(vm.getWASMErrorMessage(err));
        return;
    }

    Serial.println(":: done");
}

void
loop()
{
}

//--------------------------------------------------------------------------