From 6203b0a934b696b4e765e7bd3ec52fc9356e553f Mon Sep 17 00:00:00 2001
From: Fernando Chorney <github@djsbx.com>
Date: Sun, 7 Apr 2024 12:51:12 -0500
Subject: [PATCH] Add more SDK Commands for basic functionality

---
 sdk/commands/config.ts                |   4 +-
 sdk/commands/sensor_test.ts           |   2 +-
 sdk/smx.ts                            | 141 ++++++++++++++++++++------
 sdk/state-machines/collate-packets.ts |  21 +++-
 sdk/utils.ts                          |  20 ++++
 ui/common/typed-select.tsx            |   4 +-
 ui/pad-coms.ts                        |   3 +
 ui/stage/stage-test.tsx               |  13 ++-
 ui/state.ts                           |   4 +-
 ui/ui.tsx                             |   5 +-
 10 files changed, 175 insertions(+), 42 deletions(-)

diff --git a/sdk/commands/config.ts b/sdk/commands/config.ts
index 1a2c566..ee6dffe 100644
--- a/sdk/commands/config.ts
+++ b/sdk/commands/config.ts
@@ -205,7 +205,7 @@ export class SMXConfig {
    * TODO: Make this private again later, and maybe make a function called
    * "write_to_stage" or something? Depends on how we want to handle writing/reading
    */
-  encode(): DataView {
-    return smx_config_t.encode(this.config, true);
+  encode(): Array<number> {
+    return Array.from(new Uint8Array(smx_config_t.encode(this.config, true).buffer));
   }
 }
diff --git a/sdk/commands/sensor_test.ts b/sdk/commands/sensor_test.ts
index fdaa829..fa858cc 100644
--- a/sdk/commands/sensor_test.ts
+++ b/sdk/commands/sensor_test.ts
@@ -146,7 +146,7 @@ export class SMXPanelTestData {
        * (square it to convert back), but without this we display a bunch
        * of four and five digit numbers that are too hard to read.
        */
-      return Math.sqrt(value);
+      return Math.round((Math.sqrt(value) + Number.EPSILON) * 100) / 100;
     }
 
     // TODO: This may be necessary for defining sensor value vertial bars in the UI
diff --git a/sdk/smx.ts b/sdk/smx.ts
index 87e0ce5..4e74288 100644
--- a/sdk/smx.ts
+++ b/sdk/smx.ts
@@ -1,11 +1,12 @@
 import * as Bacon from "baconjs";
-import { collatePackets, type DataPacket } from "./state-machines/collate-packets";
+import { collatePackets, type AckPacket, type DataPacket } from "./state-machines/collate-packets";
 import { API_COMMAND, char2byte } from "./api";
 import { SMXConfig, type Decoded } from "./commands/config";
 import { SMXDeviceInfo } from "./commands/data_info";
 import { StageInputs } from "./commands/inputs";
 import { HID_REPORT_INPUT, HID_REPORT_INPUT_STATE, send_data } from "./packet";
 import { SMXSensorTestData, SensorTestMode } from "./commands/sensor_test";
+import { RGB } from "./utils";
 
 /**
  * Class purely to set up in/out event stream "pipes" to properly throttle and sync input/output from a stage
@@ -18,6 +19,8 @@ class SMXEvents {
   public readonly inputState$;
   /** read from this to see all reports sent by the stage in response to a command */
   public readonly otherReports$: Bacon.EventStream<Uint8Array>;
+  /** read from this to see if we've received an ACK */
+  public readonly ackReports$: Bacon.EventStream<AckPacket>;
   /** push to this to write commands to the stage */
   public readonly output$: Bacon.Bus<number[]>;
   /** this is everything pushed to `output$` but properly throttled/timed to the device's ack responses */
@@ -42,12 +45,13 @@ class SMXEvents {
     this.otherReports$ = (report$.filter((e) => e.type === "data") as Bacon.EventStream<DataPacket>).map(
       (e) => e.payload,
     );
+    this.ackReports$ = report$.filter((e) => e.type === "ack") as Bacon.EventStream<AckPacket>;
 
     const finishedCommand$ = report$
       .filter((e) => e.type === "host_cmd_finished")
       .map((e) => e.type === "host_cmd_finished");
 
-    // finishedCommand$.log("Cmd Finished");
+    finishedCommand$.log("Cmd Finished");
 
     const okSend$ = finishedCommand$.startWith(true);
 
@@ -69,8 +73,8 @@ class SMXEvents {
 export class SMXStage {
   private dev: HIDDevice;
   private readonly events: SMXEvents;
-  private test_mode: SensorTestMode = SensorTestMode.CalibratedValues; // TODO: Maybe we just let this be public
-  private debug = false;
+  private test_mode: SensorTestMode = SensorTestMode.CalibratedValues;
+  private debug = true;
 
   info: SMXDeviceInfo | null = null;
   config: SMXConfig | null = null;
@@ -79,8 +83,8 @@ export class SMXStage {
 
   public readonly inputState$: Bacon.EventStream<boolean[]>;
   public readonly deviceInfo$: Bacon.EventStream<SMXDeviceInfo>;
-  public readonly configResponse$: Bacon.EventStream<SMXConfig>;
-  public readonly testDataResponse$: Bacon.EventStream<SMXSensorTestData>;
+  public readonly configResponse$: Bacon.EventStream<SMXConfig | undefined>;
+  public readonly testDataResponse$: Bacon.EventStream<SMXSensorTestData | undefined>;
 
   constructor(dev: HIDDevice) {
     this.dev = dev;
@@ -115,68 +119,144 @@ export class SMXStage {
     this.inputState$ = this.events.inputState$.map((value) => this.handleInputs(value));
   }
 
-  async init(): Promise<SMXConfig> {
+  async init(): Promise<SMXSensorTestData | undefined> {
     // Request the device information
-    this.updateDeviceInfo();
+    return this.updateDeviceInfo().then(() => {
+      // Request the config for this stage
+      return this.updateConfig()?.then(() => {
+        return this.updateTestData();
+      });
+    });
+  }
+
+  /**
+   * TODO: To Implement:
+
+   * Stretch Goal:
+   * SET_PANEL_TEST_MODE
+   *
+   * Double Stretch Goal:
+   * GET_CONFIG (old firmware)
+   * WRITE_CONFIG (old_firmware)
+   * SET_SERIAL_NUMBERS
+   */
 
-    // Request some initial test data
-    this.updateTestData();
+  writeConfig(): Promise<AckPacket | undefined> {
+    if (!this.info || !this.config) return Promise.resolve(undefined);
 
-    // Request the config for this stage
-    return this.updateConfig();
+    const command = this.info.firmware_version < 5 ? API_COMMAND.WRITE_CONFIG : API_COMMAND.WRITE_CONFIG_V5;
+    const encoded_config = this.config.encode();
+    this.events.output$.push([command, encoded_config.length].concat(encoded_config));
+
+    return this.events.ackReports$.firstToPromise();
   }
 
-  updateDeviceInfo() {
+  setLightStrip(color: RGB): Promise<AckPacket | undefined> {
+    if (!this.info) return Promise.resolve(undefined);
+
+    const led_strip_index = 0; // Always 0
+    const number_of_leds = 44; // Always 44 (Unless some older or newer versions have more/less?)
+    const rgb = color.toArray();
+    const light_command = [API_COMMAND.SET_LIGHT_STRIP, led_strip_index, number_of_leds];
+
+    for (let i = 0; i < number_of_leds; i++) {
+      light_command.push(...rgb);
+    }
+
+    console.log(light_command);
+    this.events.output$.push(light_command);
+
+    return this.events.ackReports$.firstToPromise();
+  }
+
+  factoryReset(): Promise<AckPacket | undefined> {
+    if (!this.info || !this.config) return Promise.resolve(undefined);
+
+    /**
+     * Factor reset resets the platform strip color saved to the
+     * configuration, but it doesn't actually apply it to the lights.
+     *
+     * Do this for firmware v5 and up.
+     */
+    if (this.info.firmware_version >= 5) {
+      const color = this.config.config.platformStripColor;
+      this.setLightStrip(new RGB(color.r, color.g, color.b));
+    }
+
+    this.events.output$.push([API_COMMAND.FACTORY_RESET]);
+    return this.events.ackReports$.firstToPromise();
+  }
+
+  forceRecalibration(): Promise<AckPacket> {
+    this.events.output$.push([API_COMMAND.FORCE_RECALIBRATION]);
+    return this.events.ackReports$.firstToPromise();
+  }
+
+  updateDeviceInfo(): Promise<SMXDeviceInfo> {
     this.events.output$.push([API_COMMAND.GET_DEVICE_INFO]);
     return this.deviceInfo$.firstToPromise();
   }
 
-  updateConfig(): Promise<SMXConfig> {
-    this.events.output$.push([API_COMMAND.GET_CONFIG_V5]);
+  updateConfig(): Promise<SMXConfig | undefined> {
+    if (!this.info) return Promise.resolve(undefined);
+
+    const command = this.info.firmware_version < 5 ? API_COMMAND.GET_CONFIG : API_COMMAND.GET_CONFIG_V5;
+    this.events.output$.push([command]);
     return this.configResponse$.firstToPromise();
   }
 
-  updateTestData(mode: SensorTestMode | null = null): Promise<SMXSensorTestData> {
-    if (mode) {
-      this.test_mode = mode;
-    }
+  updateTestData(mode: SensorTestMode | null = null): Promise<SMXSensorTestData | undefined> {
+    if (!this.config) return Promise.resolve(undefined);
+    if (mode) this.test_mode = mode;
+
     this.events.output$.push([API_COMMAND.GET_SENSOR_TEST_DATA, this.test_mode]);
     return this.testDataResponse$.firstToPromise();
   }
 
-  private handleConfig(data: Uint8Array) {
+  private handleConfig(data: Uint8Array): SMXConfig | undefined {
+    if (!this.info) return;
+
+    /*
+    // TODO: Figure out how we want to handle this? I think we can actually convert to/from the new config
+    // from the old config
+    if (this.info.firmware_version < 5) {
+      this.config = new SMXConfigOld(Array.from(data));
+    } else {
+      this.config = new SMXConfig(Array.from(data));
+    }
+    */
     this.config = new SMXConfig(Array.from(data));
 
     // Right now I just want to confirm that decoding and encoding gives us back the same data
     const encoded_config = this.config.encode();
     if (encoded_config) {
-      const buf = new Uint8Array(encoded_config.buffer);
-      this.debug && console.log("Config Encodes Correctly: ", data.slice(2, -1).toString() === buf.toString());
+      this.debug &&
+        console.log("Config Encodes Correctly: ", data.slice(2, -1).toString() === encoded_config.toString());
     }
     this.debug && console.log("Got Config: ", this.config);
+
     return this.config;
   }
 
-  private handleTestData(data: Uint8Array) {
-    this.test = new SMXSensorTestData(
-      Array.from(data),
-      this.test_mode,
-      this.config?.config?.flags?.PlatformFlags_FSR || true,
-    );
+  private handleTestData(data: Uint8Array): SMXSensorTestData | undefined {
+    if (!this.config) return;
+
+    this.test = new SMXSensorTestData(Array.from(data), this.test_mode, this.config.config.flags.PlatformFlags_FSR);
 
     this.debug && console.log("Got Test: ", this.test);
 
     return this.test;
   }
 
-  private handleDeviceInfo(data: Uint8Array) {
+  private handleDeviceInfo(data: Uint8Array): SMXDeviceInfo {
     this.info = new SMXDeviceInfo(Array.from(data));
 
     this.debug && console.log("Got Info: ", this.info);
+
     return this.info;
   }
 
-  private handleInputs(data: Decoded<typeof StageInputs>) {
+  private handleInputs(data: Decoded<typeof StageInputs>): Array<boolean> {
     this.inputs = [
       data.up_left,
       data.up,
@@ -188,6 +268,7 @@ export class SMXStage {
       data.down,
       data.down_right,
     ];
+
     return this.inputs;
   }
 }
diff --git a/sdk/state-machines/collate-packets.ts b/sdk/state-machines/collate-packets.ts
index a7660ca..070e818 100644
--- a/sdk/state-machines/collate-packets.ts
+++ b/sdk/state-machines/collate-packets.ts
@@ -8,12 +8,15 @@ import {
   PACKET_PREAMBLE_SIZE,
 } from "../packet";
 
+const ACK_PACKET = 0x7;
+
 interface PacketHandlingState {
   currentPacket: Uint8Array;
 }
 
 export type DataPacket = { type: "data"; payload: Uint8Array };
-export type Packet = { type: "host_cmd_finished" } | DataPacket;
+export type AckPacket = { type: "ack" };
+export type Packet = { type: "host_cmd_finished" } | DataPacket | AckPacket;
 
 /**
  * Gets called when a packet is received, returns a tuple of new state and an array of
@@ -53,9 +56,15 @@ export const collatePackets: StateF<DataView, PacketHandlingState, Packet> = (st
   }
 
   // The data exists after the first 2 bytes
-  const dataBody = data.slice(2, 2 + byte_len);
+  const dataBody = data.slice(PACKET_PREAMBLE_SIZE, PACKET_PREAMBLE_SIZE + byte_len);
+
+  /**
+   * If packet starts with `ACK_PACKET`, the `byte_len` is 0, and the whole array is just 0,
+   * then this is an ack packet.
+   */
+  const isAck = (cmd & ACK_PACKET) === ACK_PACKET && byte_len === 0 && dataBody.byteLength === 0;
 
-  if ((cmd & PACKET_FLAG_START_OF_COMMAND) === PACKET_FLAG_START_OF_COMMAND && state.currentPacket.length > 0) {
+  if (cmd & PACKET_FLAG_START_OF_COMMAND && state.currentPacket.length > 0) {
     /**
      * When we get a start packet, the read buffer should already be empty. If it isn't,
      * we got a command that didn't end with an END_OF_COMMAND packet and something is wrong.
@@ -85,7 +94,11 @@ export const collatePackets: StateF<DataView, PacketHandlingState, Packet> = (st
 
   if ((cmd & PACKET_FLAG_END_OF_COMMAND) === PACKET_FLAG_END_OF_COMMAND) {
     newState = { currentPacket: new Uint8Array(0) };
-    eventsToPass.push({ type: "data", payload: nextPacket });
+    if (isAck) {
+      eventsToPass.push({ type: "ack" });
+    } else {
+      eventsToPass.push({ type: "data", payload: nextPacket });
+    }
   }
 
   return [newState, eventsToPass.map((e) => new Bacon.Next(e))];
diff --git a/sdk/utils.ts b/sdk/utils.ts
index ab24e8d..8ea5a6a 100644
--- a/sdk/utils.ts
+++ b/sdk/utils.ts
@@ -7,3 +7,23 @@ import { MAX_PACKET_SIZE } from "./packet.ts";
 export function pad_packet(packet: Array<number>, padding = 0): Uint8Array {
   return Uint8Array.from({ length: MAX_PACKET_SIZE }, (_, i) => packet[i] ?? 0);
 }
+
+export class RGB {
+  public r: number;
+  public g: number;
+  public b: number;
+
+  constructor(r: number, g: number, b: number) {
+    this.r = this.clamp(r);
+    this.g = this.clamp(g);
+    this.b = this.clamp(b);
+  }
+
+  private clamp(value: number) {
+    return value < 0 ? 0 : value > 255 ? 255 : value;
+  }
+
+  toArray(): Array<number> {
+    return [this.r, this.g, this.b];
+  }
+}
diff --git a/ui/common/typed-select.tsx b/ui/common/typed-select.tsx
index 4557e4d..d07b036 100644
--- a/ui/common/typed-select.tsx
+++ b/ui/common/typed-select.tsx
@@ -13,7 +13,9 @@ export function TypedSelect<Opts extends string>(props: Props<Opts>) {
   return (
     <select {...selectProps} onChange={(e) => onOptSelected?.(e.currentTarget.value as Opts)}>
       {options.map(([key, label]) => (
-        <option value={key}>{label}</option>
+        <option key={key} value={key}>
+          {label}
+        </option>
       ))}
     </select>
   );
diff --git a/ui/pad-coms.ts b/ui/pad-coms.ts
index df940ef..02aebe5 100644
--- a/ui/pad-coms.ts
+++ b/ui/pad-coms.ts
@@ -45,5 +45,8 @@ type FunctionKeys<T extends object> = keyof {
 /** anything here will appear in the debug UI to dispatch at will */
 export const DEBUG_COMMANDS: Record<string, FunctionKeys<SMXStage>> = {
   requestConfig: "updateConfig",
+  writeConfig: "writeConfig",
   requestTestData: "updateTestData",
+  forceRecalibration: "forceRecalibration",
+  factoryReset: "factoryReset",
 };
diff --git a/ui/stage/stage-test.tsx b/ui/stage/stage-test.tsx
index 209473a..32e4df2 100644
--- a/ui/stage/stage-test.tsx
+++ b/ui/stage/stage-test.tsx
@@ -9,6 +9,7 @@ import { timez } from "./util";
 const UI_UPDATE_RATE = 50;
 
 function useInputState(stage: SMXStage | undefined) {
+  const readTestData = useAtomValue(displayTestData$);
   const [panelStates, setPanelStates] = useState<Array<boolean> | null>();
   useEffect(() => {
     return stage?.inputState$.throttle(UI_UPDATE_RATE).onValue(setPanelStates);
@@ -25,7 +26,17 @@ function useTestData(stage: SMXStage | undefined) {
     if (!stage || !testDataMode) {
       return;
     }
-    const testMode = testDataMode === "raw" ? SensorTestMode.UncalibratedValues : SensorTestMode.CalibratedValues;
+    let testMode = SensorTestMode.UncalibratedValues;
+    switch (testDataMode) {
+      case "calibrated":
+        testMode = SensorTestMode.CalibratedValues;
+        break;
+      case "noise":
+        testMode = SensorTestMode.Noise;
+        break;
+      case "tare":
+        testMode = SensorTestMode.Tare;
+    }
     const handle = setInterval(() => stage.updateTestData(testMode), UI_UPDATE_RATE);
     return () => clearInterval(handle);
   }, [stage, testDataMode]);
diff --git a/ui/state.ts b/ui/state.ts
index 417c0fa..8178582 100644
--- a/ui/state.ts
+++ b/ui/state.ts
@@ -1,5 +1,5 @@
 import { atom, createStore } from "jotai";
-import type { SMXStage } from "../sdk";
+import { SensorTestMode, type SMXStage } from "../sdk";
 
 export const browserSupported = "hid" in navigator;
 
@@ -9,7 +9,7 @@ export const uiState = createStore();
 /** backing atom of all known stages */
 export const stages$ = atom<Record<number, SMXStage | undefined>>({});
 
-export const displayTestData$ = atom<"" | "raw" | "calibrated">("");
+export const displayTestData$ = atom<"" | "raw" | "calibrated" | "noise" | "tare">("");
 
 /** current p1 pad, derived from devices$ above */
 export const p1Stage$ = atom<SMXStage | undefined, [SMXStage | undefined], void>(
diff --git a/ui/ui.tsx b/ui/ui.tsx
index 6b92b80..9969a1f 100644
--- a/ui/ui.tsx
+++ b/ui/ui.tsx
@@ -6,6 +6,7 @@ import { open_smx_device, promptSelectDevice } from "./pad-coms.ts";
 import { browserSupported, displayTestData$, p1Stage$, p2Stage$, statusText$ } from "./state.ts";
 import { StageTest } from "./stage/stage-test.tsx";
 import { TypedSelect } from "./common/typed-select.tsx";
+import { SensorTestMode } from "../sdk/index.ts";
 
 export function UI() {
   useEffect(() => {
@@ -63,9 +64,11 @@ function TestDataDisplayToggle() {
       <TypedSelect
         value={testMode}
         options={[
-          ["", "None"],
+          ["", "Off"],
           ["calibrated", "Calibrated"],
           ["raw", "Raw"],
+          ["noise", "Noise"],
+          ["tare", "Tare"],
         ]}
         onOptSelected={(next) => setTestMode(next)}
       />