#!/usr/bin/env python3
"""
Unit tests for saltybot_can_bridge.balance_protocol.

No ROS2 or CAN hardware required — tests exercise encode/decode round-trips
and boundary conditions entirely in Python.

Run with:  pytest test/test_can_bridge.py -v
"""

import struct
import unittest

from saltybot_can_bridge.balance_protocol import (
    MAMBA_CMD_ESTOP,
    MAMBA_CMD_MODE,
    MAMBA_CMD_VELOCITY,
    MAMBA_TELEM_BATTERY,
    MAMBA_TELEM_IMU,
    VESC_TELEM_STATE,
    MODE_DRIVE,
    MODE_ESTOP,
    MODE_IDLE,
    BatteryTelemetry,
    ImuTelemetry,
    VescStateTelemetry,
    decode_battery_telem,
    decode_imu_telem,
    decode_vesc_state,
    encode_estop_cmd,
    encode_mode_cmd,
    encode_velocity_cmd,
)


class TestMessageIDs(unittest.TestCase):
    """Verify the CAN message ID constants are correct."""

    def test_command_ids(self):
        self.assertEqual(MAMBA_CMD_VELOCITY, 0x100)
        self.assertEqual(MAMBA_CMD_MODE, 0x101)
        self.assertEqual(MAMBA_CMD_ESTOP, 0x102)

    def test_telemetry_ids(self):
        self.assertEqual(MAMBA_TELEM_IMU, 0x200)
        self.assertEqual(MAMBA_TELEM_BATTERY, 0x201)
        self.assertEqual(VESC_TELEM_STATE, 0x300)


class TestVelocityEncode(unittest.TestCase):
    """Tests for encode_velocity_cmd."""

    def test_zero_velocity(self):
        payload = encode_velocity_cmd(0.0, 0.0)
        self.assertEqual(len(payload), 8)
        left, right = struct.unpack(">ff", payload)
        self.assertAlmostEqual(left, 0.0, places=5)
        self.assertAlmostEqual(right, 0.0, places=5)

    def test_forward_velocity(self):
        payload = encode_velocity_cmd(1.5, 1.5)
        left, right = struct.unpack(">ff", payload)
        self.assertAlmostEqual(left, 1.5, places=5)
        self.assertAlmostEqual(right, 1.5, places=5)

    def test_differential_velocity(self):
        payload = encode_velocity_cmd(-0.5, 0.5)
        left, right = struct.unpack(">ff", payload)
        self.assertAlmostEqual(left, -0.5, places=5)
        self.assertAlmostEqual(right, 0.5, places=5)

    def test_large_velocity(self):
        # No clamping at the protocol layer — values are sent as-is
        payload = encode_velocity_cmd(10.0, -10.0)
        left, right = struct.unpack(">ff", payload)
        self.assertAlmostEqual(left, 10.0, places=3)
        self.assertAlmostEqual(right, -10.0, places=3)

    def test_payload_is_big_endian(self):
        # Sanity check: first 4 bytes encode left speed
        payload = encode_velocity_cmd(1.0, 0.0)
        (left,) = struct.unpack(">f", payload[:4])
        self.assertAlmostEqual(left, 1.0, places=5)


class TestModeEncode(unittest.TestCase):
    """Tests for encode_mode_cmd."""

    def test_idle_mode(self):
        payload = encode_mode_cmd(MODE_IDLE)
        self.assertEqual(payload, b"\x00")

    def test_drive_mode(self):
        payload = encode_mode_cmd(MODE_DRIVE)
        self.assertEqual(payload, b"\x01")

    def test_estop_mode(self):
        payload = encode_mode_cmd(MODE_ESTOP)
        self.assertEqual(payload, b"\x02")

    def test_invalid_mode_raises(self):
        with self.assertRaises(ValueError):
            encode_mode_cmd(99)

    def test_invalid_mode_negative_raises(self):
        with self.assertRaises(ValueError):
            encode_mode_cmd(-1)


class TestEstopEncode(unittest.TestCase):
    """Tests for encode_estop_cmd."""

    def test_estop_assert(self):
        self.assertEqual(encode_estop_cmd(True), b"\x01")

    def test_estop_clear(self):
        self.assertEqual(encode_estop_cmd(False), b"\x00")

    def test_estop_default_is_stop(self):
        self.assertEqual(encode_estop_cmd(), b"\x01")


class TestImuDecodeRoundTrip(unittest.TestCase):
    """Round-trip tests for IMU telemetry."""

    def _encode_imu(self, ax, ay, az, gx, gy, gz) -> bytes:
        return struct.pack(">ffffff", ax, ay, az, gx, gy, gz)

    def test_zero_imu(self):
        data = self._encode_imu(0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
        telem = decode_imu_telem(data)
        self.assertIsInstance(telem, ImuTelemetry)
        self.assertAlmostEqual(telem.accel_x, 0.0, places=5)
        self.assertAlmostEqual(telem.gyro_z, 0.0, places=5)

    def test_nominal_imu(self):
        data = self._encode_imu(0.1, -0.2, 9.81, 0.01, -0.02, 0.03)
        telem = decode_imu_telem(data)
        self.assertAlmostEqual(telem.accel_x, 0.1, places=4)
        self.assertAlmostEqual(telem.accel_y, -0.2, places=4)
        self.assertAlmostEqual(telem.accel_z, 9.81, places=3)
        self.assertAlmostEqual(telem.gyro_x, 0.01, places=5)
        self.assertAlmostEqual(telem.gyro_y, -0.02, places=5)
        self.assertAlmostEqual(telem.gyro_z, 0.03, places=5)

    def test_imu_bad_length_raises(self):
        with self.assertRaises(struct.error):
            decode_imu_telem(b"\x00" * 10)  # too short


class TestBatteryDecodeRoundTrip(unittest.TestCase):
    """Round-trip tests for battery telemetry."""

    def _encode_bat(self, voltage, current) -> bytes:
        return struct.pack(">ff", voltage, current)

    def test_nominal_battery(self):
        data = self._encode_bat(24.6, 3.2)
        telem = decode_battery_telem(data)
        self.assertIsInstance(telem, BatteryTelemetry)
        self.assertAlmostEqual(telem.voltage, 24.6, places=3)
        self.assertAlmostEqual(telem.current, 3.2, places=4)

    def test_zero_battery(self):
        data = self._encode_bat(0.0, 0.0)
        telem = decode_battery_telem(data)
        self.assertAlmostEqual(telem.voltage, 0.0, places=5)
        self.assertAlmostEqual(telem.current, 0.0, places=5)

    def test_max_voltage(self):
        # 6S LiPo max ~25.2 V; test with a high value
        data = self._encode_bat(25.2, 0.0)
        telem = decode_battery_telem(data)
        self.assertAlmostEqual(telem.voltage, 25.2, places=3)

    def test_battery_bad_length_raises(self):
        with self.assertRaises(struct.error):
            decode_battery_telem(b"\x00" * 4)  # too short


class TestVescStateDecodeRoundTrip(unittest.TestCase):
    """Round-trip tests for VESC state telemetry."""

    def _encode_vesc(self, erpm, duty, voltage, current) -> bytes:
        return struct.pack(">ffff", erpm, duty, voltage, current)

    def test_nominal_vesc(self):
        data = self._encode_vesc(3000.0, 0.25, 24.0, 5.5)
        telem = decode_vesc_state(data)
        self.assertIsInstance(telem, VescStateTelemetry)
        self.assertAlmostEqual(telem.erpm, 3000.0, places=2)
        self.assertAlmostEqual(telem.duty, 0.25, places=5)
        self.assertAlmostEqual(telem.voltage, 24.0, places=4)
        self.assertAlmostEqual(telem.current, 5.5, places=4)

    def test_zero_vesc(self):
        data = self._encode_vesc(0.0, 0.0, 0.0, 0.0)
        telem = decode_vesc_state(data)
        self.assertAlmostEqual(telem.erpm, 0.0)
        self.assertAlmostEqual(telem.duty, 0.0)

    def test_reverse_erpm(self):
        data = self._encode_vesc(-1500.0, -0.15, 23.0, 2.0)
        telem = decode_vesc_state(data)
        self.assertAlmostEqual(telem.erpm, -1500.0, places=2)
        self.assertAlmostEqual(telem.duty, -0.15, places=5)

    def test_vesc_bad_length_raises(self):
        with self.assertRaises(struct.error):
            decode_vesc_state(b"\x00" * 8)  # too short (need 16)


class TestEncodeDecodeIdentity(unittest.TestCase):
    """Cross-module identity tests: encode then decode must be identity."""

    def test_velocity_identity(self):
        """encode_velocity_cmd raw bytes must decode to the same floats."""
        left, right = 0.75, -0.3
        payload = encode_velocity_cmd(left, right)
        decoded_l, decoded_r = struct.unpack(">ff", payload)
        self.assertAlmostEqual(decoded_l, left, places=5)
        self.assertAlmostEqual(decoded_r, right, places=5)

    def test_imu_identity(self):
        accel = (0.5, -0.5, 9.8)
        gyro = (0.1, -0.1, 0.2)
        raw = struct.pack(">ffffff", *accel, *gyro)
        telem = decode_imu_telem(raw)
        self.assertAlmostEqual(telem.accel_x, accel[0], places=4)
        self.assertAlmostEqual(telem.accel_y, accel[1], places=4)
        self.assertAlmostEqual(telem.accel_z, accel[2], places=3)
        self.assertAlmostEqual(telem.gyro_x, gyro[0], places=5)
        self.assertAlmostEqual(telem.gyro_y, gyro[1], places=5)
        self.assertAlmostEqual(telem.gyro_z, gyro[2], places=5)

    def test_battery_identity(self):
        voltage, current = 22.4, 8.1
        raw = struct.pack(">ff", voltage, current)
        telem = decode_battery_telem(raw)
        self.assertAlmostEqual(telem.voltage, voltage, places=3)
        self.assertAlmostEqual(telem.current, current, places=4)

    def test_vesc_identity(self):
        erpm, duty, voltage, current = 5000.0, 0.5, 24.5, 10.0
        raw = struct.pack(">ffff", erpm, duty, voltage, current)
        telem = decode_vesc_state(raw)
        self.assertAlmostEqual(telem.erpm, erpm, places=2)
        self.assertAlmostEqual(telem.duty, duty, places=5)
        self.assertAlmostEqual(telem.voltage, voltage, places=3)
        self.assertAlmostEqual(telem.current, current, places=4)


if __name__ == "__main__":
    unittest.main()
