From f929955ca5b874f413bebf3fe8559b44ed37487c Mon Sep 17 00:00:00 2001
From: Bruno Albuquerque <bga@gmail.com>
Date: Sat, 18 May 2024 09:07:01 -0400
Subject: [PATCH] Add ChassisSpeedLevelType.

- Also add some comments.
---
 module/robot/chassis_speed_level_type.go | 20 ++++++++++++++++++++
 module/robot/robot.go                    | 12 +++++++-----
 2 files changed, 27 insertions(+), 5 deletions(-)
 create mode 100644 module/robot/chassis_speed_level_type.go

diff --git a/module/robot/chassis_speed_level_type.go b/module/robot/chassis_speed_level_type.go
new file mode 100644
index 0000000..31d7aaf
--- /dev/null
+++ b/module/robot/chassis_speed_level_type.go
@@ -0,0 +1,20 @@
+package robot
+
+// ChassisSpeedLevelType represents the speed level of the chassis.
+type ChassisSpeedLevelType int8
+
+const (
+	// ChassisSpeedLevelFast is the fast speed level.
+	ChassisSpeedLevelFast ChassisSpeedLevelType = iota
+	// ChassisSpeedLevelMedium is the medium speed level.
+	ChassisSpeedLevelMedium
+	// ChassisSpeedLevelSlow is the slow speed level.
+	ChassisSpeedLevelSlow
+	// ChassisSpeedLevelCustom is the custom speed level. This enables
+	// individually setting speeds for each direction/axis. This is not
+	// supported yet.
+	//
+	// TODO(bga): Add support for this.
+	ChassisSpeedLevelCustom
+	ChassisSpeedLevelTypeCount
+)
diff --git a/module/robot/robot.go b/module/robot/robot.go
index d36c67c..74ce5bb 100644
--- a/module/robot/robot.go
+++ b/module/robot/robot.go
@@ -183,22 +183,24 @@ func (r *Robot) BatteryPowerPercent() uint8 {
 	return *r.batteryPowerPercent.Load()
 }
 
-func (r *Robot) ChassisSpeedLevel() (uint8, error) {
+// ChassisSpeedLevel returns the current chassis speed level.
+func (r *Robot) ChassisSpeedLevel() (ChassisSpeedLevelType, error) {
 	res, err := r.UB().GetKeyValueSync(key.KeyRobomasterSystemChassisSpeedLevel, true)
 	if err != nil {
 		return 0, err
 	}
 
-	return uint8(res.Value().(*value.Uint64).Value), nil
+	return ChassisSpeedLevelType(res.Value().(*value.Uint64).Value - 1), nil
 }
 
-func (r *Robot) SetChassisSpeedLevel(speedLevel uint8) error {
-	if speedLevel < 1 || speedLevel > 4 {
+// SetChassisSpeedLevel sets the chassis speed level.
+func (r *Robot) SetChassisSpeedLevel(speedLevel ChassisSpeedLevelType) error {
+	if speedLevel >= ChassisSpeedLevelTypeCount {
 		return fmt.Errorf("invalid chassis speed level: %d", speedLevel)
 	}
 
 	return r.UB().SetKeyValueSync(key.KeyRobomasterSystemChassisSpeedLevel,
-		&value.Uint64{Value: uint64(speedLevel)})
+		&value.Uint64{Value: uint64(speedLevel + 1)})
 }
 
 // Stop stops the Robot module.