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

WASM vm;

/**
;; ------------------------------------------------------------
;;
;; LEDBounce_I32.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.
;;
;; I32 - Q16:16 fixed point math (optimal for MCU without 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))
  (global $ONE_Q16_16 i32 (i32.const 65536)) ;; 1 in Q16:16 (fixed point math)

  ;; ---- 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 i32) (local $y i32) (local $dx i32) (local $dy i32)
    (local $r i32) (local $tmp i32) (local $idx i32)
    (local $MAXX i32) (local $MAXY i32)
    (local $nx i32) (local $px i32) (local $ny i32) (local $py i32)

    ;; device init
    (call $begin)

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

    ;; bounds in i32: [0, MAXX), [0, MAXY)
    (local.set $MAXX (i32.shl (global.get $MAX_X) (i32.const 16)))
    (local.set $MAXY (i32.shl (global.get $MAX_Y) (i32.const 16)))

    ;; 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.xor (local.get $r) (i32.shr_u (local.get $r) (i32.const 31))))
      (local.set $x (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.xor (local.get $r) (i32.shr_u (local.get $r) (i32.const 31))))
      (local.set $y (i32.rem_u (local.get $r) (global.get $MAX_Y)))

      ;; dx = ± (1 / {2,3,4})
      (local.set $r (call $rand32))
      (local.set $r (i32.xor (local.get $r) (i32.shr_u (local.get $r) (i32.const 31))))
      (local.set $tmp (i32.add (i32.const 2) (i32.rem_u (local.get $r) (i32.const 3))))
      (local.set $dx (i32.div_s (global.get $ONE_Q16_16) (local.get $tmp)))
      (if (i32.and (call $rand32) (i32.const 1)) (then) (else (local.set $dx (i32.sub (i32.const 0) (local.get $dx)))))

      ;; dy = ± (1 / {2,3,4})
      (local.set $r (call $rand32))
      (local.set $r (i32.xor (local.get $r) (i32.shr_u (local.get $r) (i32.const 31))))
      (local.set $tmp (i32.add (i32.const 2) (i32.rem_u (local.get $r) (i32.const 3))))
      (local.set $dy (i32.div_s (global.get $ONE_Q16_16) (local.get $tmp)))
      (if (i32.and (call $rand32) (i32.const 1)) (then) (else (local.set $dy (i32.sub (i32.const 0) (local.get $dy)))))

      ;; store point
      (local.set $base (i32.add (global.get $STATE) (i32.mul (local.get $i) (i32.const 16))))
      (i32.store (i32.add (local.get $base) (i32.const 0)) (i32.shl (local.get $x) (i32.const 16)))
      (i32.store (i32.add (local.get $base) (i32.const 4)) (i32.shl (local.get $y) (i32.const 16)))
      (i32.store (i32.add (local.get $base) (i32.const 8)) (local.get $dx))
      (i32.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 (i32.shr_s (i32.load (i32.add (local.get $base) (i32.const 0))) (i32.const 16)))
        (local.set $y (i32.shr_s (i32.load (i32.add (local.get $base) (i32.const 4))) (i32.const 16)))

        ;; clamp to boundaries
        (if
          (i32.and
            (i32.lt_u (local.get $x) (global.get $MAX_X))
            (i32.lt_u (local.get $y) (global.get $MAX_Y)))
          (then
            (local.set $idx
              (i32.add (global.get $FB)
                (i32.add (local.get $x) (i32.mul (local.get $y) (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) (i32.const 12) (i32.const 8))
      (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  (i32.load (i32.add (local.get $base) (i32.const 0))))
        (local.set $y  (i32.load (i32.add (local.get $base) (i32.const 4))))
        (local.set $dx (i32.load (i32.add (local.get $base) (i32.const 8))))
        (local.set $dy (i32.load (i32.add (local.get $base) (i32.const 12))))

        ;; reflect if next or previous step leaves bounds
        (local.set $nx (i32.add (local.get $x) (local.get $dx)))
        (local.set $px (i32.sub (local.get $x) (local.get $dx)))
        (if (i32.or
              (i32.or (i32.lt_s (local.get $nx) (i32.const 0))
                      (i32.lt_s (local.get $px) (i32.const 0)))
              (i32.or (i32.ge_s (local.get $nx) (local.get $MAXX))
                      (i32.ge_s (local.get $px) (local.get $MAXX))))
          (then (local.set $dx (i32.sub (i32.const 0) (local.get $dx))))
        )
        (local.set $ny (i32.add (local.get $y) (local.get $dy)))
        (local.set $py (i32.sub (local.get $y) (local.get $dy)))
        (if (i32.or
              (i32.or (i32.lt_s (local.get $ny) (i32.const 0))
                      (i32.lt_s (local.get $py) (i32.const 0)))
              (i32.or (i32.ge_s (local.get $ny) (local.get $MAXY))
                      (i32.ge_s (local.get $py) (local.get $MAXY))))
          (then (local.set $dy (i32.sub (i32.const 0) (local.get $dy))))
        )

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

        ;; store the values
        (i32.store (i32.add (local.get $base) (i32.const 0))  (local.get $x))
        (i32.store (i32.add (local.get $base) (i32.const 4))  (local.get $y))
        (i32.store (i32.add (local.get $base) (i32.const 8))  (local.get $dx))
        (i32.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, 0x27, 0x07, 0x7f, 0x00, 0x41, 0x0c, 0x0b, 0x7f, 0x00, 0x41, 0x08, 0x0b, 0x7f, 0x00,
  0x41, 0x05, 0x0b, 0x7f, 0x00, 0x41, 0x80, 0x80, 0x04, 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, 0x86,
  0x05, 0x02, 0x25, 0x01, 0x01, 0x7f, 0x23, 0x04, 0x28, 0x02, 0x00, 0x21, 0x00, 0x20, 0x00, 0x41,
  0x8d, 0xcc, 0xe5, 0x00, 0x6c, 0x41, 0xdf, 0xe6, 0xbb, 0xe3, 0x03, 0x6a, 0x21, 0x00, 0x23, 0x04,
  0x20, 0x00, 0x36, 0x02, 0x00, 0x20, 0x00, 0x0b, 0xdd, 0x04, 0x01, 0x10, 0x7f, 0x10, 0x00, 0x23,
  0x04, 0x41, 0xf8, 0xac, 0xd1, 0x91, 0x01, 0x36, 0x02, 0x00, 0x23, 0x00, 0x41, 0x10, 0x74, 0x21,
  0x0a, 0x23, 0x01, 0x41, 0x10, 0x74, 0x21, 0x0b, 0x41, 0x00, 0x21, 0x00, 0x03, 0x40, 0x10, 0x04,
  0x21, 0x07, 0x20, 0x07, 0x20, 0x07, 0x41, 0x1f, 0x76, 0x73, 0x21, 0x07, 0x20, 0x07, 0x23, 0x00,
  0x70, 0x21, 0x03, 0x10, 0x04, 0x21, 0x07, 0x20, 0x07, 0x20, 0x07, 0x41, 0x1f, 0x76, 0x73, 0x21,
  0x07, 0x20, 0x07, 0x23, 0x01, 0x70, 0x21, 0x04, 0x10, 0x04, 0x21, 0x07, 0x20, 0x07, 0x20, 0x07,
  0x41, 0x1f, 0x76, 0x73, 0x21, 0x07, 0x41, 0x02, 0x20, 0x07, 0x41, 0x03, 0x70, 0x6a, 0x21, 0x08,
  0x23, 0x03, 0x20, 0x08, 0x6d, 0x21, 0x05, 0x10, 0x04, 0x41, 0x01, 0x71, 0x04, 0x40, 0x05, 0x41,
  0x00, 0x20, 0x05, 0x6b, 0x21, 0x05, 0x0b, 0x10, 0x04, 0x21, 0x07, 0x20, 0x07, 0x20, 0x07, 0x41,
  0x1f, 0x76, 0x73, 0x21, 0x07, 0x41, 0x02, 0x20, 0x07, 0x41, 0x03, 0x70, 0x6a, 0x21, 0x08, 0x23,
  0x03, 0x20, 0x08, 0x6d, 0x21, 0x06, 0x10, 0x04, 0x41, 0x01, 0x71, 0x04, 0x40, 0x05, 0x41, 0x00,
  0x20, 0x06, 0x6b, 0x21, 0x06, 0x0b, 0x23, 0x06, 0x20, 0x00, 0x41, 0x10, 0x6c, 0x6a, 0x21, 0x02,
  0x20, 0x02, 0x41, 0x00, 0x6a, 0x20, 0x03, 0x41, 0x10, 0x74, 0x36, 0x02, 0x00, 0x20, 0x02, 0x41,
  0x04, 0x6a, 0x20, 0x04, 0x41, 0x10, 0x74, 0x36, 0x02, 0x00, 0x20, 0x02, 0x41, 0x08, 0x6a, 0x20,
  0x05, 0x36, 0x02, 0x00, 0x20, 0x02, 0x41, 0x0c, 0x6a, 0x20, 0x06, 0x36, 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, 0x05, 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, 0x06, 0x20, 0x00, 0x41, 0x10, 0x6c, 0x6a, 0x21, 0x02, 0x20,
  0x02, 0x41, 0x00, 0x6a, 0x28, 0x02, 0x00, 0x41, 0x10, 0x75, 0x21, 0x03, 0x20, 0x02, 0x41, 0x04,
  0x6a, 0x28, 0x02, 0x00, 0x41, 0x10, 0x75, 0x21, 0x04, 0x20, 0x03, 0x23, 0x00, 0x49, 0x20, 0x04,
  0x23, 0x01, 0x49, 0x71, 0x04, 0x40, 0x23, 0x05, 0x20, 0x03, 0x20, 0x04, 0x23, 0x00, 0x6c, 0x6a,
  0x6a, 0x21, 0x09, 0x20, 0x09, 0x41, 0x01, 0x3a, 0x00, 0x00, 0x0b, 0x20, 0x00, 0x41, 0x01, 0x6a,
  0x21, 0x00, 0x20, 0x00, 0x23, 0x02, 0x49, 0x0d, 0x00, 0x0b, 0x23, 0x05, 0x41, 0x0c, 0x41, 0x08,
  0x10, 0x01, 0x10, 0x02, 0x41, 0x32, 0x10, 0x03, 0x41, 0x00, 0x21, 0x00, 0x03, 0x40, 0x23, 0x06,
  0x20, 0x00, 0x41, 0x10, 0x6c, 0x6a, 0x21, 0x02, 0x20, 0x02, 0x41, 0x00, 0x6a, 0x28, 0x02, 0x00,
  0x21, 0x03, 0x20, 0x02, 0x41, 0x04, 0x6a, 0x28, 0x02, 0x00, 0x21, 0x04, 0x20, 0x02, 0x41, 0x08,
  0x6a, 0x28, 0x02, 0x00, 0x21, 0x05, 0x20, 0x02, 0x41, 0x0c, 0x6a, 0x28, 0x02, 0x00, 0x21, 0x06,
  0x20, 0x03, 0x20, 0x05, 0x6a, 0x21, 0x0c, 0x20, 0x03, 0x20, 0x05, 0x6b, 0x21, 0x0d, 0x20, 0x0c,
  0x41, 0x00, 0x48, 0x20, 0x0d, 0x41, 0x00, 0x48, 0x72, 0x20, 0x0c, 0x20, 0x0a, 0x4e, 0x20, 0x0d,
  0x20, 0x0a, 0x4e, 0x72, 0x72, 0x04, 0x40, 0x41, 0x00, 0x20, 0x05, 0x6b, 0x21, 0x05, 0x0b, 0x20,
  0x04, 0x20, 0x06, 0x6a, 0x21, 0x0e, 0x20, 0x04, 0x20, 0x06, 0x6b, 0x21, 0x0f, 0x20, 0x0e, 0x41,
  0x00, 0x48, 0x20, 0x0f, 0x41, 0x00, 0x48, 0x72, 0x20, 0x0e, 0x20, 0x0b, 0x4e, 0x20, 0x0f, 0x20,
  0x0b, 0x4e, 0x72, 0x72, 0x04, 0x40, 0x41, 0x00, 0x20, 0x06, 0x6b, 0x21, 0x06, 0x0b, 0x20, 0x03,
  0x20, 0x05, 0x6a, 0x21, 0x03, 0x20, 0x04, 0x20, 0x06, 0x6a, 0x21, 0x04, 0x20, 0x02, 0x41, 0x00,
  0x6a, 0x20, 0x03, 0x36, 0x02, 0x00, 0x20, 0x02, 0x41, 0x04, 0x6a, 0x20, 0x04, 0x36, 0x02, 0x00,
  0x20, 0x02, 0x41, 0x08, 0x6a, 0x20, 0x05, 0x36, 0x02, 0x00, 0x20, 0x02, 0x41, 0x0c, 0x6a, 0x20,
  0x06, 0x36, 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 = 871;

#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()
{
}

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